diff --git a/ui/store/background-connection.test.ts b/ui/store/background-connection.test.ts index cd6c3736334b..f9c96b5c60a8 100644 --- a/ui/store/background-connection.test.ts +++ b/ui/store/background-connection.test.ts @@ -20,24 +20,17 @@ function setup() { notificationListeners.add(listener); }); - const removeOnNotification = jest - .fn() - .mockImplementation((listener: NotificationListener) => { - notificationListeners.delete(listener); - }); - const submitNotification = (notification: JsonRpcNotification) => { notificationListeners.forEach((listener) => listener(notification)); }; - const messengerSubscribe = jest.fn(); + const messengerSubscribe = jest.fn().mockResolvedValue(undefined); const messengerUnsubscribe = jest.fn(); // @ts-expect-error Partial mock. setBackgroundConnection({ onNotification, - removeOnNotification, messengerSubscribe, messengerUnsubscribe, }); @@ -45,7 +38,6 @@ function setup() { return { submitNotification, onNotification, - removeOnNotification, messengerSubscribe, messengerUnsubscribe, }; @@ -80,7 +72,6 @@ describe('subscribeToMessengerEvent', () => { messengerSubscribe, messengerUnsubscribe, onNotification, - removeOnNotification, } = setup(); const listener = jest.fn(); @@ -101,7 +92,90 @@ describe('subscribeToMessengerEvent', () => { expect(listener).not.toHaveBeenCalled(); expect(messengerUnsubscribe).toHaveBeenCalledWith(event); - expect(removeOnNotification).toHaveBeenCalledWith(expect.any(Function)); + }); + + it('only calls messengerSubscribe once if there are multiple subscriptions for the same event requested', async () => { + const { submitNotification, messengerSubscribe, onNotification } = setup(); + + const listenerA = jest.fn(); + const listenerB = jest.fn(); + + await subscribeToMessengerEvent(event, listenerA); + await subscribeToMessengerEvent(event, listenerB); + + expect(messengerSubscribe).toHaveBeenCalledTimes(1); + expect(messengerSubscribe).toHaveBeenCalledWith(event); + expect(onNotification).toHaveBeenCalledTimes(1); + + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + expect(listenerA).toHaveBeenCalledWith([{ foo: 'bar' }, []]); + expect(listenerB).toHaveBeenCalledWith([{ foo: 'bar' }, []]); + }); + + it('does not resolve subscribe calls until the upstream messengerSubscribe call resolves', async () => { + const { messengerSubscribe } = setup(); + + const { promise: subscribeRpcPromise, resolve: resolveSubscribe } = + createDeferredPromise(); + messengerSubscribe.mockReturnValueOnce(subscribeRpcPromise); + + let resolvedA = false; + let resolvedB = false; + + const subscribeA = subscribeToMessengerEvent(event, jest.fn()).then(() => { + resolvedA = true; + }); + const subscribeB = subscribeToMessengerEvent(event, jest.fn()).then(() => { + resolvedB = true; + }); + + // Let any already-scheduled microtasks run; neither subscribe call + // should have resolved yet because both await the same in-flight call. + await new Promise((r) => setImmediate(r)); + + expect(resolvedA).toBe(false); + expect(resolvedB).toBe(false); + + resolveSubscribe(); + + await subscribeA; + await subscribeB; + + expect(resolvedA).toBe(true); + expect(resolvedB).toBe(true); + }); + + it('refcounts subscribers and only unsubscribes upstream on the last unsubscribe', async () => { + const { submitNotification, messengerUnsubscribe } = setup(); + + const listenerA = jest.fn(); + const listenerB = jest.fn(); + + const unsubscribeA = await subscribeToMessengerEvent(event, listenerA); + const unsubscribeB = await subscribeToMessengerEvent(event, listenerB); + + await unsubscribeA(); + + expect(messengerUnsubscribe).not.toHaveBeenCalled(); + + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + expect(listenerA).not.toHaveBeenCalled(); + expect(listenerB).toHaveBeenCalledWith([{ foo: 'bar' }, []]); + + await unsubscribeB(); + + expect(messengerUnsubscribe).toHaveBeenCalledTimes(1); + expect(messengerUnsubscribe).toHaveBeenCalledWith(event); }); it('ignores other JSON-RPC notifications', async () => { @@ -119,4 +193,227 @@ describe('subscribeToMessengerEvent', () => { expect(listener).not.toHaveBeenCalled(); }); + + it('clears the entry when the upstream messengerSubscribe call rejects, allowing retry', async () => { + const { messengerSubscribe } = setup(); + + messengerSubscribe.mockRejectedValueOnce(new Error('subscribe failed')); + messengerSubscribe.mockResolvedValueOnce(undefined); + + const listener = jest.fn(); + + await expect(subscribeToMessengerEvent(event, listener)).rejects.toThrow( + 'subscribe failed', + ); + + // A fresh subscribe attempt should send a new messengerSubscribe call, + // not silently reuse a rejected entry. + await subscribeToMessengerEvent(event, listener); + + expect(messengerSubscribe).toHaveBeenCalledTimes(2); + }); + + it('rejects all concurrent subscribers when the upstream messengerSubscribe RPC rejects', async () => { + const { messengerSubscribe } = setup(); + + messengerSubscribe.mockRejectedValueOnce(new Error('subscribe failed')); + + const listenerA = jest.fn(); + const listenerB = jest.fn(); + + const subscribeA = subscribeToMessengerEvent(event, listenerA); + const subscribeB = subscribeToMessengerEvent(event, listenerB); + + expect(messengerSubscribe).toHaveBeenCalledTimes(1); + + await expect(subscribeA).rejects.toThrow('subscribe failed'); + await expect(subscribeB).rejects.toThrow('subscribe failed'); + + // Entry is cleared, so a retry sends a fresh messengerSubscribe call. + messengerSubscribe.mockResolvedValueOnce(undefined); + await subscribeToMessengerEvent(event, listenerA); + expect(messengerSubscribe).toHaveBeenCalledTimes(2); + }); + + it('continues invoking remaining callbacks when one callback throws', async () => { + const { submitNotification } = setup(); + + const consoleErrorSpy = jest + .spyOn(console, 'error') + .mockImplementation(() => undefined); + + const throwingListener = jest.fn(() => { + throw new Error('callback boom'); + }); + const otherListener = jest.fn(); + + await subscribeToMessengerEvent(event, throwingListener); + await subscribeToMessengerEvent(event, otherListener); + + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + expect(throwingListener).toHaveBeenCalledTimes(1); + expect(otherListener).toHaveBeenCalledTimes(1); + expect(consoleErrorSpy).toHaveBeenCalledWith(expect.any(Error)); + + consoleErrorSpy.mockRestore(); + }); + + it('is idempotent when the same unsubscribe function is called twice', async () => { + const { messengerUnsubscribe } = setup(); + + const listener = jest.fn(); + const unsubscribe = await subscribeToMessengerEvent(event, listener); + + await unsubscribe(); + await unsubscribe(); + + expect(messengerUnsubscribe).toHaveBeenCalledTimes(1); + }); + + it('keeps subscriptions to different events independent', async () => { + const { submitNotification, messengerSubscribe } = setup(); + + const otherEvent = 'OtherController:stateChange'; + + const listenerA = jest.fn(); + const listenerB = jest.fn(); + + await subscribeToMessengerEvent(event, listenerA); + await subscribeToMessengerEvent(otherEvent, listenerB); + + expect(messengerSubscribe).toHaveBeenCalledTimes(2); + + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + expect(listenerA).toHaveBeenCalledTimes(1); + expect(listenerB).not.toHaveBeenCalled(); + }); + + it('treats two subscriptions sharing the same callback reference as independent', async () => { + const { submitNotification, messengerUnsubscribe } = setup(); + + const listener = jest.fn(); + + const unsubscribeFirst = await subscribeToMessengerEvent(event, listener); + const unsubscribeSecond = await subscribeToMessengerEvent(event, listener); + + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + // Each subscribe call is an independent registration, so a notification + // fires the shared callback once per registration. + expect(listener).toHaveBeenCalledTimes(2); + + // Unsubscribing one registration leaves the other intact, so no upstream + // unsubscribe yet and the remaining registration still receives events. + await unsubscribeFirst(); + expect(messengerUnsubscribe).not.toHaveBeenCalled(); + + listener.mockClear(); + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + expect(listener).toHaveBeenCalledTimes(1); + + // Removing the last registration sends the upstream unsubscribe. + await unsubscribeSecond(); + expect(messengerUnsubscribe).toHaveBeenCalledTimes(1); + }); + + it('ignores subscription notifications that have no params', async () => { + const { submitNotification } = setup(); + + const listener = jest.fn(); + await subscribeToMessengerEvent(event, listener); + + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + // No params field. + }); + + expect(listener).not.toHaveBeenCalled(); + }); + + it('propagates messengerUnsubscribe rejection to the awaiter', async () => { + const { messengerUnsubscribe } = setup(); + + messengerUnsubscribe.mockRejectedValueOnce(new Error('unsubscribe failed')); + + const listener = jest.fn(); + const unsubscribe = await subscribeToMessengerEvent(event, listener); + + await expect(unsubscribe()).rejects.toThrow('unsubscribe failed'); + }); + + it('dispatches notifications that arrive before the upstream messengerSubscribe RPC resolves', async () => { + const { messengerSubscribe, submitNotification } = setup(); + + const { promise: subscribeRpcPromise, resolve: resolveSubscribe } = + createDeferredPromise(); + messengerSubscribe.mockReturnValueOnce(subscribeRpcPromise); + + const listener = jest.fn(); + const subscribePromise = subscribeToMessengerEvent(event, listener); + + // The upstream RPC is still pending, but the notification router is + // attached and the callback set is populated. A notification arriving + // now must still reach the listener. + submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + expect(listener).toHaveBeenCalledWith([{ foo: 'bar' }, []]); + + resolveSubscribe(); + await subscribePromise; + }); + + it('clears subscription state when setBackgroundConnection is called again', async () => { + const firstConnection = setup(); + + const firstListener = jest.fn(); + await subscribeToMessengerEvent(event, firstListener); + + expect(firstConnection.messengerSubscribe).toHaveBeenCalledTimes(1); + expect(firstConnection.onNotification).toHaveBeenCalledTimes(1); + + // Replace the background connection. The new connection should start + // with no in-memory subscription state — a subscribe for the same + // event must send a fresh upstream messengerSubscribe call and attach a + // fresh notification router. + const secondConnection = setup(); + + const secondListener = jest.fn(); + await subscribeToMessengerEvent(event, secondListener); + + expect(secondConnection.messengerSubscribe).toHaveBeenCalledTimes(1); + expect(secondConnection.messengerSubscribe).toHaveBeenCalledWith(event); + expect(secondConnection.onNotification).toHaveBeenCalledTimes(1); + + secondConnection.submitNotification({ + jsonrpc: '2.0', + method: MESSENGER_SUBSCRIPTION_NOTIFICATION, + params: [event, [{ foo: 'bar' }, []]], + }); + + expect(secondListener).toHaveBeenCalledWith([{ foo: 'bar' }, []]); + expect(firstListener).not.toHaveBeenCalled(); + }); }); diff --git a/ui/store/background-connection.ts b/ui/store/background-connection.ts index 3896c9e57147..f4e8ef318d2e 100644 --- a/ui/store/background-connection.ts +++ b/ui/store/background-connection.ts @@ -51,8 +51,55 @@ export function submitRequestToBackground( return background[method](...rpcArgs) as unknown as Promise; } +type MessengerEventSubscription = { + // Each subscribe call is an independent registration keyed by a unique + // symbol, so unsubscribing one never affects another even when they share + // the same callback reference. + callbacks: Map void>; + subscribePromise: Promise; +}; + +const messengerEventSubscriptions = new Map< + NamespacedName, + MessengerEventSubscription +>(); +// Tracks whether the router is attached to the current background reference. +// setBackgroundConnection resets it so the router reattaches to the new +// connection on the next subscribe. +let notificationRouterAttached = false; + +function routeMessengerEventNotification( + notification: JsonRpcNotification<[string, Json]>, +) { + if ( + notification.method !== MESSENGER_SUBSCRIPTION_NOTIFICATION || + !notification.params + ) { + return; + } + const [eventName, payload] = notification.params; + // `eventName` is `string` from the notification params; mismatches are caught by the entry-not-found check below. + const subscription = messengerEventSubscriptions.get( + eventName as NamespacedName, + ); + if (!subscription) { + return; + } + for (const callback of subscription.callbacks.values()) { + try { + callback(payload); + } catch (error) { + console.error(error); + } + } +} + /** - * Sets/replaces the background connection reference + * Sets/replaces the background connection reference. + * + * Clears any in-memory subscription state because subscriptions registered + * against the previous background connection are stale once the connection + * has been replaced. * * @param backgroundConnection */ @@ -60,34 +107,88 @@ export async function setBackgroundConnection( backgroundConnection: BackgroundRpcClient, ) { background = backgroundConnection; + messengerEventSubscriptions.clear(); + notificationRouterAttached = false; } /** - * Subscribe to a given messenger event emitted by the background. + * Subscribe to a given event emitted by the background via the root messenger. + * + * Because callbacks cannot be sent to the background, we create the + * subscription in two steps: + * + * 1. First, we send a `messengerSubscribe` request to the background, which + * will attach an event listener to the root messenger. When the event occurs, + * it will send a notification. + * 2. Second, we use `onNotification` on the background client to wait for the + * notification. When it arrives, it will call the given callback. + * + * To prevent unnecessary calls to the background, if this function is called + * more than once for the same pending event, callbacks will be consolidated; + * when the event occurs, they will be called in the order they were defined. * * @param event - The event name. * @param callback - The callback to invoke when the event is emitted. - * @returns A cleanup function that can be invoked to unsubscribe. + * @returns A cleanup function that can be invoked to remove the subscription on + * the messenger event. If multiple subscriptions exist for the same messenger + * event, the unsubscribe function will only take effect once there is only one + * subscriber left. */ export async function subscribeToMessengerEvent( event: NamespacedName, callback: (data: Data) => void, ): Promise<() => Promise> { - await submitRequestToBackground('messengerSubscribe', [event]); - - const listener = (notification: JsonRpcNotification<[string, Data]>) => { - if ( - notification.method === MESSENGER_SUBSCRIPTION_NOTIFICATION && - notification.params?.[0] === event - ) { - callback(notification.params[1]); + // `Data extends Json` but `(data: Data) => void` is not assignable to `(data: Json) => void` due to contravariant function parameters; the cast is safe because all callbacks receive `Json`-shaped data at runtime. + const looselyTypedCallback = callback as (data: Json) => void; + + // Unique key for this registration so it can be removed independently of any + // other registration that happens to share the same callback reference. + const registrationId = Symbol('messengerEventSubscription'); + + let subscription = messengerEventSubscriptions.get(event); + + if (subscription) { + subscription.callbacks.set(registrationId, looselyTypedCallback); + } else { + if (!notificationRouterAttached) { + background.onNotification(routeMessengerEventNotification); + notificationRouterAttached = true; } - }; - background.onNotification(listener); + const subscribePromise = submitRequestToBackground( + 'messengerSubscribe', + [event], + ); + + subscription = { + callbacks: new Map([[registrationId, looselyTypedCallback]]), + subscribePromise, + }; + messengerEventSubscriptions.set(event, subscription); + + // Side-effect handler: clear the entry on rejection so future subscribe + // attempts retry cleanly. Do NOT reassign `subscription.subscribePromise` + // here — the original promise must retain its rejection so awaiters below + // still see it. + subscribePromise.catch(() => { + messengerEventSubscriptions.delete(event); + }); + } + + await subscription.subscribePromise; + + return async () => { + const currentSubscription = messengerEventSubscriptions.get(event); + if (!currentSubscription) { + return; + } + + const removed = currentSubscription.callbacks.delete(registrationId); + if (!removed || currentSubscription.callbacks.size > 0) { + return; + } - return () => { - background.removeOnNotification(listener); - return submitRequestToBackground('messengerUnsubscribe', [event]); + messengerEventSubscriptions.delete(event); + await submitRequestToBackground('messengerUnsubscribe', [event]); }; }