Use disposable to remove registered listeners

TypeFox · Nov 22, 2024 · 2d3babb · 2d3babb
1 parent 9dc66a2
commit 2d3babb
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 19 deletions.
diff --git a/packages/vscode-messenger-common/src/messages.ts b/packages/vscode-messenger-common/src/messages.ts
@@ -202,9 +202,18 @@ export class Deferred<R = any> {
  */
 export interface CancellationToken {
     readonly isCancellationRequested: boolean;
-    onCancellationRequested(callBack: (reason: string) => void): void;
+    onCancellationRequested(callBack: (reason: string) => void): Disposable;
 }
 
+/**
+ * Interface for objects that can be disposed.
+ */
+export interface Disposable {
+    /**
+     * Dispose this object.
+     */
+    dispose(): void;
+}
 /**
 * Implementation of the CancellationToken interface.
 * Allows to trigger cancelation.
@@ -219,19 +228,21 @@ export class CancellationTokenImpl implements CancellationToken {
         }
         this.canceled = true;
         this.listeners.forEach(callBack => callBack(reason));
-        this.clearCancelListeners();
+        this.listeners = [];
     }
 
     get isCancellationRequested(): boolean {
         return this.canceled;
     }
 
-    public onCancellationRequested(callBack: (reason: string) => void): void {
+    public onCancellationRequested(callBack: (reason: string) => void): Disposable {
         this.listeners.push(callBack);
-    }
-
-    public clearCancelListeners(): void {
-        this.listeners = [];
+        const listeners = this.listeners;
+        return {
+            dispose() {
+                listeners.splice(listeners.indexOf(callBack), 1);
+            }
+        };
     }
 }
 

diff --git a/packages/vscode-messenger-webview/src/messenger.ts b/packages/vscode-messenger-webview/src/messenger.ts
@@ -171,15 +171,15 @@ export class Messenger implements MessengerAPI {
         const pending = new Deferred<R>();
         this.requests.set(msgId, pending);
         if (cancelable) {
-            cancelable.onCancellationRequested((reason) => {
+            const listener = cancelable.onCancellationRequested((reason) => {
                 // Send cancel message for pending request
                 this.vscode.postMessage(createCancelRequestMessage(receiver, {  msgId }));
                 pending.reject(new Error(reason));
                 this.requests.delete(msgId);
             });
             pending.result.finally(() => {
-                // Request finished, nothing to do on cancel.
-                cancelable.clearCancelListeners();
+                // Request finished, remove the listener
+                listener.dispose();
             }).catch((err: unknown) =>
                 this.log(`Pending request rejected: ${String(err)}`)
             );

diff --git a/packages/vscode-messenger-webview/tests/messenger-webview.test.ts b/packages/vscode-messenger-webview/tests/messenger-webview.test.ts
@@ -8,7 +8,7 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 
 import crypto from 'crypto';
-import { CancellationTokenImpl, createCancelRequestMessage, HOST_EXTENSION, isRequestMessage, Message, MessageParticipant, NotificationType, RequestType } from 'vscode-messenger-common';
+import { CancellationTokenImpl, createCancelRequestMessage, Disposable, HOST_EXTENSION, isRequestMessage, Message, MessageParticipant, NotificationType, RequestType } from 'vscode-messenger-common';
 import { Messenger, VsCodeApi } from '../src';
 
 Object.defineProperty(globalThis, 'crypto', {
@@ -285,12 +285,13 @@ describe('Webview Messenger', () => {
         let started = false;
         let canceled = false;
         let handled = false;
+        const toDispose: Disposable[] = [];
         new Messenger(vsCodeApi).onRequest(stringRequest, async (param: string, sender, cancelation) => {
             let timeOut: any;
-            cancelation.onCancellationRequested(() => {
+            toDispose.push(cancelation.onCancellationRequested(() => {
                 clearTimeout(timeOut);
                 canceled = true;
-            });
+            }));
             started = true;
             // simulate work in progress
             await new Promise<void>(resolve => {
@@ -311,6 +312,7 @@ describe('Webview Messenger', () => {
         // send cancel request
         postWindowMsg(createCancelRequestMessage({ type: 'webview', webviewId: 'test-view' }, { msgId: 'request_id' }));
 
+        toDispose.forEach(d => d.dispose());
         expect(started).toBe(true);
         expect(canceled).toBe(true);
         expect(handled).toBe(false);

diff --git a/packages/vscode-messenger/src/messenger.ts b/packages/vscode-messenger/src/messenger.ts
@@ -355,15 +355,15 @@ export class Messenger implements MessengerAPI {
         const pendingRequest = new Deferred<R>();
         this.requests.set(msgId, pendingRequest);
         if (cancelable) {
-            cancelable.onCancellationRequested((reason) => {
+            const listener = cancelable.onCancellationRequested((reason) => {
                 // Send cancel message for pending request
                 view.webview.postMessage(createCancelRequestMessage(receiver, { msgId }));
                 pendingRequest.reject(new Error(reason));
                 this.requests.delete(msgId);
             });
             pendingRequest.result.finally(() => {
-                // Request finished, nothing to do on cancel.
-                cancelable.clearCancelListeners();
+                // Request finished, remove listener
+                listener.dispose();
             }).catch((err) =>
                 this.log(`Pending request rejected: ${String(err)}`)
             );

diff --git a/packages/vscode-messenger/tests/messenger.test.ts b/packages/vscode-messenger/tests/messenger.test.ts
@@ -7,7 +7,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable @typescript-eslint/no-unused-vars */
 
-import { BROADCAST, CancellationTokenImpl, createCancelRequestMessage, HOST_EXTENSION, isCancelRequestNotification, isRequestMessage, MessageParticipant, NotificationType, RequestType, WebviewTypeMessageParticipant } from 'vscode-messenger-common';
+import { BROADCAST, CancellationTokenImpl, createCancelRequestMessage, Disposable, HOST_EXTENSION, isCancelRequestNotification, isRequestMessage, MessageParticipant, NotificationType, RequestType, WebviewTypeMessageParticipant } from 'vscode-messenger-common';
 import { MessengerEvent } from '../src/diagnostic-api';
 import { Messenger } from '../src/messenger';
 
@@ -528,11 +528,12 @@ describe('Extension Messenger', () => {
         messenger.registerWebviewView(view1);
         let started = false;
         let handled = false;
+        const toDispose: Disposable[] = [];
         messenger.onRequest(simpleRequest, async (params: string, sender, cancelation) => {
             let timeOut: any;
-            cancelation.onCancellationRequested(() => {
+            toDispose.push(cancelation.onCancellationRequested(() => {
                 clearTimeout(timeOut);
-            });
+            }));
             started = true;
             // simulate work in progress
             await new Promise<void>(resolve => {
@@ -548,6 +549,7 @@ describe('Extension Messenger', () => {
         const cancelMsg = createCancelRequestMessage(HOST_EXTENSION, { msgId: 'fake_req_id' });
         await view1.messageCallback(cancelMsg);
 
+        toDispose.forEach(disposable => disposable.dispose());
         expect(started).toBe(true);
         expect(handled).toBe(false);
         expect(view1.messages[0]).toBeUndefined(); // don't expect cancelation succeed