From 5e41d47f3a041086ca2a09290b3e39afef7a009d Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Thu, 27 Apr 2023 17:14:20 -0700
Subject: [PATCH] Add "close()" method to MP Web Tasks

PiperOrigin-RevId: 527726737
---
 mediapipe/tasks/web/core/task_runner.ts            |  5 +++++
 mediapipe/tasks/web/core/task_runner_test_utils.ts |  3 ++-
 mediapipe/web/graph_runner/graph_runner.ts         | 11 +++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/mediapipe/tasks/web/core/task_runner.ts b/mediapipe/tasks/web/core/task_runner.ts
index 9a36b648..91a38cd4 100644
--- a/mediapipe/tasks/web/core/task_runner.ts
+++ b/mediapipe/tasks/web/core/task_runner.ts
@@ -256,6 +256,11 @@ export abstract class TaskRunner {
 
     this.baseOptions.setAcceleration(acceleration);
   }
+
+  /** Closes and cleans up the resources held by this task. */
+  close(): void {
+    this.graphRunner.closeGraph();
+  }
 }
 
 
diff --git a/mediapipe/tasks/web/core/task_runner_test_utils.ts b/mediapipe/tasks/web/core/task_runner_test_utils.ts
index d911b2c8..edf1d0d3 100644
--- a/mediapipe/tasks/web/core/task_runner_test_utils.ts
+++ b/mediapipe/tasks/web/core/task_runner_test_utils.ts
@@ -36,7 +36,8 @@ export function createSpyWasmModule(): SpyWasmModule {
     '_setAutoRenderToScreen', 'stringToNewUTF8', '_attachProtoListener',
     '_attachProtoVectorListener', '_free', '_waitUntilIdle',
     '_addStringToInputStream', '_registerModelResourcesGraphService',
-    '_configureAudio', '_malloc', '_addProtoToInputStream', '_getGraphConfig'
+    '_configureAudio', '_malloc', '_addProtoToInputStream', '_getGraphConfig',
+    '_closeGraph'
   ]);
   spyWasmModule._getGraphConfig.and.callFake(() => {
     (spyWasmModule.simpleListeners![CALCULATOR_GRAPH_CONFIG_LISTENER_NAME] as
diff --git a/mediapipe/web/graph_runner/graph_runner.ts b/mediapipe/web/graph_runner/graph_runner.ts
index 0115312b..615971cb 100644
--- a/mediapipe/web/graph_runner/graph_runner.ts
+++ b/mediapipe/web/graph_runner/graph_runner.ts
@@ -63,6 +63,7 @@ export declare interface WasmModule {
   _bindTextureToCanvas: () => boolean;
   _changeBinaryGraph: (size: number, dataPtr: number) => void;
   _changeTextGraph: (size: number, dataPtr: number) => void;
+  _closeGraph: () => void;
   _free: (ptr: number) => void;
   _malloc: (size: number) => number;
   _processFrame: (width: number, height: number, timestamp: number) => void;
@@ -1148,6 +1149,16 @@ export class GraphRunner {
   finishProcessing(): void {
     this.wasmModule._waitUntilIdle();
   }
+
+  /**
+   * Closes the input streams and all calculators for this graph and frees up
+   * any C++ resources. The graph will not be usable once closed.
+   */
+  closeGraph(): void {
+    this.wasmModule._closeGraph();
+    this.wasmModule.simpleListeners = undefined;
+    this.wasmModule.emptyPacketListeners = undefined;
+  }
 }
 
 // Quick private helper to run the given script safely