fix: unsubscribe from RpcBrowser

containers · Jan 17, 2024 · 48870fa · 48870fa
1 parent f71ddd3
commit 48870fa
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git a/packages/frontend/src/pages/ModelPlayground.svelte b/packages/frontend/src/pages/ModelPlayground.svelte
@@ -19,7 +19,7 @@
     }
     studioClient.startPlayground(model.id);
 
-    playgroundQueries.subscribe((queries: QueryState[]) => {
+    const unsubscribe = playgroundQueries.subscribe((queries: QueryState[]) => {
       if (queryId === -1) {
         return;
       }
@@ -32,6 +32,10 @@
       }
       displayQuery(myQuery);
     });
+
+    return () => {
+      unsubscribe();
+    };
   });
 
   function displayQuery(query: QueryState) {

diff --git a/packages/frontend/src/stores/playground-queries.ts b/packages/frontend/src/stores/playground-queries.ts
@@ -6,8 +6,11 @@ import { MSG_NEW_PLAYGROUND_QUERIES_STATE } from '@shared/Messages';
 import { studioClient } from '/@/utils/client';
 
 export const playgroundQueries: Readable<QueryState[]> = readable<QueryState[]>([], (set) => {  
-  rpcBrowser.subscribe(MSG_NEW_PLAYGROUND_QUERIES_STATE, (msg) => {
+  const sub = rpcBrowser.subscribe(MSG_NEW_PLAYGROUND_QUERIES_STATE, (msg) => {
     set(msg);
   });
   studioClient.askPlaygroundQueries();
+  return () => {
+    sub.unsubscribe();
+  }
 });
diff --git a/packages/shared/MessageProxy.ts b/packages/shared/MessageProxy.ts
@@ -83,6 +83,10 @@ export class RpcExtension {
   }
 }
 
+export interface Subscriber {
+  unsubscribe(): void;
+}
+
 export class RpcBrowser {
   counter: number = 0;
   promises: Map<number, {resolve: (value: unknown) => any, reject: (value: unknown) => void}> = new Map();
@@ -166,10 +170,16 @@ export class RpcBrowser {
     })
   }
 
-  // TODO(feloy) need to unsubscribe
   // TODO(feloy) need to subscribe several times?
-  subscribe(msgId: string, f: (msg: any) => void) {
+  subscribe(msgId: string, f: (msg: any) => void): Subscriber {
+    console.log('add subscriber');
     this.subscribers.set(msgId, f);
+    return {
+      unsubscribe: () => {
+        console.log('delete subscriber');
+        this.subscribers.delete(msgId);
+      } 
+    }
   }
 
   isSubscribedMessage(content: any): content is ISubscribedMessage {