From 024f782cd9c26951dd72cacf4aaf57d59550c55e Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Tue, 16 May 2023 10:44:01 -0700
Subject: [PATCH] Warn users that do not invoke "close()"

PiperOrigin-RevId: 532507235
---
 mediapipe/tasks/web/vision/core/image.ts | 23 +++++++++++++++++++-
 mediapipe/tasks/web/vision/core/mask.ts  | 27 +++++++++++++++++++++---
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/mediapipe/tasks/web/vision/core/image.ts b/mediapipe/tasks/web/vision/core/image.ts
index 4a7b6dce7..3b067bd78 100644
--- a/mediapipe/tasks/web/vision/core/image.ts
+++ b/mediapipe/tasks/web/vision/core/image.ts
@@ -16,6 +16,9 @@
 
 import {assertNotNull, MPImageShaderContext} from '../../../../tasks/web/vision/core/image_shader_context';
 
+/** Number of instances a user can keep alive before we raise a warning. */
+const INSTANCE_COUNT_WARNING_THRESHOLD = 250;
+
 /** The underlying type of the image. */
 enum MPImageType {
   /** Represents the native `ImageData` type. */
@@ -52,6 +55,12 @@ export type MPImageContainer = ImageData|ImageBitmap|WebGLTexture;
 export class MPImage {
   private gl?: WebGL2RenderingContext;
 
+  /**
+   * A counter to track the number of instances of MPImage that own resources..
+   * This is used to raise a warning if the user does not close the instances.
+   */
+  private static instancesBeforeWarning = INSTANCE_COUNT_WARNING_THRESHOLD;
+
   /** @hideconstructor */
   constructor(
       private readonly containers: MPImageContainer[],
@@ -64,7 +73,16 @@ export class MPImage {
       readonly width: number,
       /** Returns the height of the image. */
       readonly height: number,
-  ) {}
+  ) {
+    if (this.ownsImageBitmap || this.ownsWebGLTexture) {
+      --MPImage.instancesBeforeWarning;
+      if (MPImage.instancesBeforeWarning === 0) {
+        console.error(
+            'You seem to be creating MPImage instances without invoking ' +
+            '.close(). This leaks resources.');
+      }
+    }
+  }
 
   /** Returns whether this `MPImage` contains a mask of type `ImageData`. */
   hasImageData(): boolean {
@@ -391,6 +409,9 @@ export class MPImage {
       const gl = this.getGL();
       gl.deleteTexture(this.getContainer(MPImageType.WEBGL_TEXTURE)!);
     }
+
+    // User called close(). We no longer issue warning.
+    MPImage.instancesBeforeWarning = -1;
   }
 }
 
diff --git a/mediapipe/tasks/web/vision/core/mask.ts b/mediapipe/tasks/web/vision/core/mask.ts
index 13deab3fc..d7cf59e5f 100644
--- a/mediapipe/tasks/web/vision/core/mask.ts
+++ b/mediapipe/tasks/web/vision/core/mask.ts
@@ -16,6 +16,9 @@
 
 import {assertNotNull, MPImageShaderContext} from '../../../../tasks/web/vision/core/image_shader_context';
 
+/** Number of instances a user can keep alive before we raise a warning. */
+const INSTANCE_COUNT_WARNING_THRESHOLD = 250;
+
 /** The underlying type of the image. */
 enum MPMaskType {
   /** Represents the native `UInt8Array` type. */
@@ -47,6 +50,12 @@ export type MPMaskContainer = Uint8Array|Float32Array|WebGLTexture;
 export class MPMask {
   private gl?: WebGL2RenderingContext;
 
+  /**
+   * A counter to track the number of instances of MPMask that own resources.
+   * This is used to raise a warning if the user does not close the instances.
+   */
+  private static instancesBeforeWarning = INSTANCE_COUNT_WARNING_THRESHOLD;
+
   /** @hideconstructor */
   constructor(
       private readonly containers: MPMaskContainer[],
@@ -58,7 +67,16 @@ export class MPMask {
       readonly width: number,
       /** Returns the height of the mask. */
       readonly height: number,
-  ) {}
+  ) {
+    if (this.ownsWebGLTexture) {
+      --MPMask.instancesBeforeWarning;
+      if (MPMask.instancesBeforeWarning === 0) {
+        console.error(
+            'You seem to be creating MPMask instances without invoking ' +
+            '.close(). This leaks resources.');
+      }
+    }
+  }
 
   /** Returns whether this `MPMask` contains a mask of type `Uint8Array`. */
   hasUint8Array(): boolean {
@@ -88,8 +106,8 @@ export class MPMask {
 
   /**
    * Returns the underlying mask as a single channel `Float32Array`. Note that
-   * this involves an expensive GPU to CPU transfer if the current mask is only
-   * available as a `WebGLTexture`.
+   * this involves an expensive GPU to CPU transfer if the current mask is
+   * only available as a `WebGLTexture`.
    *
    * @return The current mask as a Float32Array.
    */
@@ -311,6 +329,9 @@ export class MPMask {
       const gl = this.getGL();
       gl.deleteTexture(this.getContainer(MPMaskType.WEBGL_TEXTURE)!);
     }
+
+    // User called close(). We no longer issue warning.
+    MPMask.instancesBeforeWarning = -1;
   }
 }