fix: message proxy support multiple subscribers (#2392)

* fix: message proxy support multiple subscribers

Signed-off-by: axel7083 <[email protected]>

* fix: rpc readable tests

Signed-off-by: axel7083 <[email protected]>

---------

Signed-off-by: axel7083 <[email protected]>

Author: axel7083

packages/frontend/src/stores/rpcReadable.spec.ts

@@ -37,7 +37,7 @@ vi.mock('../utils/client', async () => {
     postMessage: (message: unknown) => {
       if (message && typeof message === 'object' && 'channel' in message) {
         const f = rpcBrowser.subscribers.get(message.channel as string);
-        f?.('');
+        f?.forEach(listener => listener(''));
       }
     },
   } as unknown as PodmanDesktopApi;

packages/shared/src/messages/MessageProxy.spec.ts

@@ -187,6 +187,97 @@ test('getChannel should use CHANNEL property of classType provided', () => {
   expect(channel).toBe('dummy-ping');
 });
 
+describe('subscribe', () => {
+  beforeEach(() => {
+    window.addEventListener = vi.fn();
+
+    (defaultNoTimeoutChannels.noTimeoutChannels as string[]) = [];
+  });
+
+  function getMessageListener(): (event: MessageEvent) => void {
+    expect(window.addEventListener).toHaveBeenCalledOnce();
+    expect(window.addEventListener).toHaveBeenCalledWith('message', expect.any(Function));
+    return vi.mocked(window.addEventListener).mock.calls[0][1] as (event: MessageEvent) => void;
+  }
+
+  test('subscriber should be called on event received', async () => {
+    const rpcBrowser = new RpcBrowser(window, api);
+    const messageListener = getMessageListener();
+
+    const listener = vi.fn();
+    rpcBrowser.subscribe('example', listener);
+
+    messageListener({
+      data: {
+        id: 'example',
+        body: 'hello',
+      },
+    } as unknown as MessageEvent);
+
+    expect(listener).toHaveBeenCalledOnce();
+  });
+
+  test('all subscribers should be called if multiple exists', async () => {
+    const rpcBrowser = new RpcBrowser(window, api);
+    const messageListener = getMessageListener();
+
+    const listeners = Array.from({ length: 10 }, _ => vi.fn());
+
+    listeners.forEach(listener => rpcBrowser.subscribe('example', listener));
+
+    messageListener({
+      data: {
+        id: 'example',
+        body: 'hello',
+      },
+    } as unknown as MessageEvent);
+
+    for (const listener of listeners) {
+      expect(listener).toHaveBeenCalledWith('hello');
+    }
+  });
+
+  test('subscribers which unsubscribe should not be called', async () => {
+    const rpcBrowser = new RpcBrowser(window, api);
+    const messageListener = getMessageListener();
+
+    const [listenerA, listenerB] = [vi.fn(), vi.fn()];
+
+    const unsubscriberA = rpcBrowser.subscribe('example', listenerA);
+    const unsubscriberB = rpcBrowser.subscribe('example', listenerB);
+
+    messageListener({
+      data: {
+        id: 'example',
+        body: 'hello',
+      },
+    } as unknown as MessageEvent);
+
+    // unsubscriber the listener B
+    unsubscriberB.unsubscribe();
+
+    messageListener({
+      data: {
+        id: 'example',
+        body: 'hello',
+      },
+    } as unknown as MessageEvent);
+
+    // unsubscriber the listener A
+    unsubscriberA.unsubscribe();
+
+    messageListener({
+      data: {
+        id: 'example',
+        body: 'hello',
+      },
+    } as unknown as MessageEvent);
+
+    expect(listenerA).toHaveBeenCalledTimes(2);
+    expect(listenerB).toHaveBeenCalledOnce();
+  });
+});
+
 describe('no timeout channel', () => {
   beforeEach(() => {
     vi.resetAllMocks();

packages/shared/src/messages/MessageProxy.ts

@@ -129,10 +129,12 @@ export interface Subscriber {
   unsubscribe(): void;
 }
 
+export type Listener = (value: any) => void;
+
 export class RpcBrowser {
   counter: number = 0;
   promises: Map<number, { resolve: (value: unknown) => unknown; reject: (value: unknown) => void }> = new Map();
-  subscribers: Map<string, (msg: any) => void> = new Map();
+  subscribers: Map<string, Set<Listener>> = new Map();
 
   getUniqueId(): number {
     return ++this.counter;
@@ -164,8 +166,7 @@ export class RpcBrowser {
         }
         this.promises.delete(message.id);
       } else if (this.isSubscribedMessage(message)) {
-        const handler = this.subscribers.get(message.id);
-        handler?.(message.body);
+        this.subscribers.get(message.id)?.forEach(handler => handler(message.body));
       } else {
         console.error('Received incompatible message.', message);
         return;
@@ -220,12 +221,12 @@ export class RpcBrowser {
     return promise;
   }
 
-  // TODO(feloy) need to subscribe several times?
-  subscribe(msgId: string, f: (msg: any) => void): Subscriber {
-    this.subscribers.set(msgId, f);
+  subscribe(msgId: string, f: Listener): Subscriber {
+    this.subscribers.set(msgId, (this.subscribers.get(msgId) ?? new Set()).add(f));
+
     return {
       unsubscribe: (): void => {
-        this.subscribers.delete(msgId);
+        this.subscribers.get(msgId)?.delete(f);
       },
     };
   }