Add controller messaging system (#377)

Adds a controller messaging system, which will be used to facilitate inter-controller communication. The controller messenger acts as a message broker, passing actions and events back and forth between controllers. It is _fully type safe_. This idea was described in the Controller Messaging System proposal[1]. This controller messenger doesn't yet support attenuation, so it's not yet possible to restrict which actions and events are interacted with. That will be coming in a later PR. [1]: https://www.notion.so/Controller-Messaging-System-617efb02b9e54bd0a0b0e44c6f776d85 * Remove `subscribe` and `unsubscribe` BaseController methods The base controller no longer exposes methods to subscribe and unsubscribe to state changes. Instead all subscriptions and unsubscriptions are done with the controller messaging system.
MetaMask · Oct 11, 2023 · 39bacb1 · 39bacb1
1 parent 3fd6386
commit 39bacb1
Show file tree

Hide file tree

Showing 4 changed files with 574 additions and 57 deletions.
diff --git a/src/BaseControllerV2.test.ts b/src/BaseControllerV2.test.ts
@@ -1,21 +1,27 @@
-import type { Draft } from 'immer';
+import type { Draft, Patch } from 'immer';
 import * as sinon from 'sinon';
 
 import { BaseController, getAnonymizedState, getPersistentState } from './BaseControllerV2';
+import { ControllerMessenger } from './ControllerMessenger';
 
-type MockControllerState = {
+type CountControllerState = {
   count: number;
 };
 
-const mockControllerStateMetadata = {
+type CountControllerEvent = {
+  type: `CountController:stateChange`;
+  payload: [CountControllerState, Patch[]];
+};
+
+const CountControllerStateMetadata = {
   count: {
     persist: true,
     anonymous: true,
   },
 };
 
-class MockController extends BaseController<MockControllerState> {
-  update(callback: (state: Draft<MockControllerState>) => void | MockControllerState) {
+class MockController extends BaseController<'CountController', CountControllerState> {
+  update(callback: (state: Draft<CountControllerState>) => void | CountControllerState) {
     super.update(callback);
   }
 
@@ -26,27 +32,51 @@ class MockController extends BaseController<MockControllerState> {
 
 describe('BaseController', () => {
   it('should set initial state', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
 
     expect(controller.state).toEqual({ count: 0 });
   });
 
   it('should set initial schema', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
 
-    expect(controller.metadata).toEqual(mockControllerStateMetadata);
+    expect(controller.metadata).toEqual(CountControllerStateMetadata);
   });
 
   it('should not allow mutating state directly', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
 
     expect(() => {
       controller.state = { count: 1 };
     }).toThrow();
   });
 
   it('should allow updating state by modifying draft', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
 
     controller.update((draft) => {
       draft.count += 1;
@@ -56,7 +86,13 @@ describe('BaseController', () => {
   });
 
   it('should allow updating state by return a value', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
 
     controller.update(() => {
       return { count: 1 };
@@ -66,7 +102,13 @@ describe('BaseController', () => {
   });
 
   it('should throw an error if update callback modifies draft and returns value', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
 
     expect(() => {
       controller.update((draft) => {
@@ -77,12 +119,18 @@ describe('BaseController', () => {
   });
 
   it('should inform subscribers of state changes', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
     const listener1 = sinon.stub();
     const listener2 = sinon.stub();
 
-    controller.subscribe(listener1);
-    controller.subscribe(listener2);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener2);
     controller.update(() => {
       return { count: 1 };
     });
@@ -94,11 +142,18 @@ describe('BaseController', () => {
   });
 
   it('should inform a subscriber of each state change once even after multiple subscriptions', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
     const listener1 = sinon.stub();
 
-    controller.subscribe(listener1);
-    controller.subscribe(listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+
     controller.update(() => {
       return { count: 1 };
     });
@@ -108,11 +163,17 @@ describe('BaseController', () => {
   });
 
   it('should no longer inform a subscriber about state changes after unsubscribing', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
     const listener1 = sinon.stub();
 
-    controller.subscribe(listener1);
-    controller.unsubscribe(listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    controllerMessenger.unsubscribe('CountController:stateChange', listener1);
     controller.update(() => {
       return { count: 1 };
     });
@@ -121,35 +182,48 @@ describe('BaseController', () => {
   });
 
   it('should no longer inform a subscriber about state changes after unsubscribing once, even if they subscribed many times', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
     const listener1 = sinon.stub();
 
-    controller.subscribe(listener1);
-    controller.subscribe(listener1);
-    controller.unsubscribe(listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    controllerMessenger.unsubscribe('CountController:stateChange', listener1);
     controller.update(() => {
       return { count: 1 };
     });
 
     expect(listener1.callCount).toEqual(0);
   });
 
-  it('should allow unsubscribing listeners who were never subscribed', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+  it('should throw when unsubscribing listener who was never subscribed', () => {
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    new MockController(controllerMessenger, 'CountController', { count: 0 }, CountControllerStateMetadata);
     const listener1 = sinon.stub();
 
     expect(() => {
-      controller.unsubscribe(listener1);
-    }).not.toThrow();
+      controllerMessenger.unsubscribe('CountController:stateChange', listener1);
+    }).toThrow();
   });
 
   it('should no longer update subscribers after being destroyed', () => {
-    const controller = new MockController({ count: 0 }, mockControllerStateMetadata);
+    const controllerMessenger = new ControllerMessenger<never, CountControllerEvent>();
+    const controller = new MockController(
+      controllerMessenger,
+      'CountController',
+      { count: 0 },
+      CountControllerStateMetadata,
+    );
     const listener1 = sinon.stub();
     const listener2 = sinon.stub();
 
-    controller.subscribe(listener1);
-    controller.subscribe(listener2);
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    controllerMessenger.subscribe('CountController:stateChange', listener2);
     controller.destroy();
     controller.update(() => {
       return { count: 1 };

diff --git a/src/BaseControllerV2.ts b/src/BaseControllerV2.ts
@@ -4,6 +4,8 @@ import { enablePatches, produceWithPatches } from 'immer';
 // eslint-disable-next-line no-duplicate-imports
 import type { Draft, Patch } from 'immer';
 
+import type { ControllerMessenger } from './ControllerMessenger';
+
 enablePatches();
 
 /**
@@ -96,21 +98,31 @@ type Json = null | boolean | number | string | Json[] | { [prop: string]: Json }
 /**
  * Controller class that provides state management, subscriptions, and state metadata
  */
-export class BaseController<S extends Record<string, unknown>> {
+export class BaseController<N extends string, S extends Record<string, unknown>> {
   private internalState: IsJsonable<S>;
 
-  private internalListeners: Set<Listener<S>> = new Set();
+  private messagingSystem: ControllerMessenger<never, { type: `${N}:stateChange`; payload: [S, Patch[]] }>;
+
+  private name: N;
 
   public readonly metadata: StateMetadata<S>;
 
   /**
    * Creates a BaseController instance.
    *
+   * @param messagingSystem - Controller messaging system
    * @param state - Initial controller state
    * @param metadata - State metadata, describing how to "anonymize" the state,
    *   and which parts should be persisted.
    */
-  constructor(state: IsJsonable<S>, metadata: StateMetadata<S>) {
+  constructor(
+    messagingSystem: ControllerMessenger<never, { type: `${N}:stateChange`; payload: [S, Patch[]] }>,
+    name: N,
+    state: IsJsonable<S>,
+    metadata: StateMetadata<S>,
+  ) {
+    this.messagingSystem = messagingSystem;
+    this.name = name;
     this.internalState = state;
     this.metadata = metadata;
   }
@@ -128,24 +140,6 @@ export class BaseController<S extends Record<string, unknown>> {
     throw new Error(`Controller state cannot be directly mutated; use 'update' method instead.`);
   }
 
-  /**
-   * Adds new listener to be notified of state changes
-   *
-   * @param listener - Callback triggered when state changes
-   */
-  subscribe(listener: Listener<S>) {
-    this.internalListeners.add(listener);
-  }
-
-  /**
-   * Removes existing listener from receiving state changes
-   *
-   * @param listener - Callback to remove
-   */
-  unsubscribe(listener: Listener<S>) {
-    this.internalListeners.delete(listener);
-  }
-
   /**
    * Updates controller state. Accepts a callback that is passed a draft copy
    * of the controller state. If a value is returned, it is set as the new
@@ -158,9 +152,7 @@ export class BaseController<S extends Record<string, unknown>> {
   protected update(callback: (state: Draft<IsJsonable<S>>) => void | IsJsonable<S>) {
     const [nextState, patches] = produceWithPatches(this.internalState, callback);
     this.internalState = nextState as IsJsonable<S>;
-    for (const listener of this.internalListeners) {
-      listener(nextState as S, patches);
-    }
+    this.messagingSystem.publish(`${this.name}:stateChange` as `${N}:stateChange`, nextState as S, patches);
   }
 
   /**
@@ -173,7 +165,7 @@ export class BaseController<S extends Record<string, unknown>> {
    * listeners from being garbage collected.
    */
   protected destroy() {
-    this.internalListeners.clear();
+    this.messagingSystem.clearEventSubscriptions(`${this.name}:stateChange` as `${N}:stateChange`);
   }
 }