perf(keyring-controller): do not fire unecessary :stageChange with withKeyring

Author: ccharly

packages/keyring-controller/package.json

@@ -60,6 +60,7 @@
     "async-mutex": "^0.5.0",
     "ethereumjs-wallet": "^1.0.1",
     "immer": "^9.0.6",
+    "lodash": "^4.17.21",
     "ulid": "^2.3.0"
   },
   "devDependencies": {

packages/keyring-controller/src/KeyringController.test.ts

@@ -41,6 +41,7 @@ import MockEncryptor, {
 } from '../tests/mocks/mockEncryptor';
 import { MockErc4337Keyring } from '../tests/mocks/mockErc4337Keyring';
 import { MockKeyring } from '../tests/mocks/mockKeyring';
+import { MockMutableKeyring } from '../tests/mocks/mockMutableKeyring';
 import MockShallowGetAccountsKeyring from '../tests/mocks/mockShallowGetAccountsKeyring';
 import { buildMockTransaction } from '../tests/mocks/mockTransaction';
 
@@ -3080,6 +3081,71 @@ describe('KeyringController', () => {
             },
           );
         });
+
+        it('should update the vault if the keyring is being updated', async () => {
+          await withController(
+            { keyringBuilders: [keyringBuilderFactory(MockMutableKeyring)] },
+            async ({ controller, messenger }) => {
+              const selector = { type: MockMutableKeyring.type };
+
+              const mockStateChange = jest.fn();
+              messenger.subscribe(
+                'KeyringController:stateChange',
+                mockStateChange,
+              );
+
+              await controller.withKeyring(
+                selector,
+                async ({ keyring }) => {
+                  (keyring as MockMutableKeyring).update(); // Update the keyring state.
+                },
+                {
+                  createIfMissing: true,
+                },
+              );
+
+              expect(mockStateChange).toHaveBeenCalled();
+              expect(controller.state.keyrings).toHaveLength(2);
+            },
+          );
+        });
+
+        it('should not update the vault if the keyring has not been updated', async () => {
+          await withController(
+            { keyringBuilders: [keyringBuilderFactory(MockMutableKeyring)] },
+            async ({ controller, messenger }) => {
+              const selector = { type: MockMutableKeyring.type };
+
+              const mockStateChange = jest.fn();
+              messenger.subscribe(
+                'KeyringController:stateChange',
+                mockStateChange,
+              );
+
+              // First create the reading before reading from it.
+              await controller.withKeyring(
+                selector,
+                async () => {
+                  // No-op, but the keyring will still be persisted since we
+                  // used `createIfMissing`.
+                },
+                {
+                  createIfMissing: true,
+                },
+              );
+
+              expect(mockStateChange).toHaveBeenCalled();
+              expect(controller.state.keyrings).toHaveLength(2);
+
+              mockStateChange.mockReset();
+              await controller.withKeyring(selector, async () => {
+                // No-op, keyring state won't be updated.
+              });
+
+              expect(mockStateChange).not.toHaveBeenCalled();
+            },
+          );
+        });
       });
     });
 

packages/keyring-controller/src/KeyringController.ts

@@ -38,6 +38,7 @@ import type { MutexInterface } from 'async-mutex';
 import Wallet, { thirdparty as importers } from 'ethereumjs-wallet';
 import type { Patch } from 'immer';
 // When generating a ULID within the same millisecond, monotonicFactory provides some guarantees regarding sort order.
+import { isEqual } from 'lodash';
 import { ulid } from 'ulid';
 
 import { KeyringControllerError } from './constants';
@@ -1551,8 +1552,9 @@ export class KeyringController extends BaseController<
   ): Promise<CallbackResult> {
     this.#assertIsUnlocked();
 
-    return this.#persistOrRollback(async () => {
+    return this.#withRollback(async () => {
       let keyring: SelectedKeyring | undefined;
+      let forceUpdate = false;
 
       if ('address' in selector) {
         keyring = (await this.getKeyringForAccount(selector.address)) as
@@ -1568,6 +1570,9 @@ export class KeyringController extends BaseController<
             selector.type,
             options.createWithData,
           )) as SelectedKeyring;
+
+          // This is a new keyring, so we force the vault update in that case.
+          forceUpdate = true;
         }
       } else if ('id' in selector) {
         keyring = this.#getKeyringById(selector.id) as SelectedKeyring;
@@ -1577,6 +1582,8 @@ export class KeyringController extends BaseController<
         throw new Error(KeyringControllerError.KeyringNotFound);
       }
 
+      const oldKeyringState = await keyring.serialize();
+
       const result = await operation({
         keyring,
         metadata: this.#getKeyringMetadata(keyring),
@@ -1590,6 +1597,11 @@ export class KeyringController extends BaseController<
         throw new Error(KeyringControllerError.UnsafeDirectKeyringAccess);
       }
 
+      if (forceUpdate || !isEqual(oldKeyringState, await keyring.serialize())) {
+        // Keyring has been updated, we need to update the vault.
+        await this.#updateVault();
+      }
+
       return result;
     });
   }

packages/keyring-controller/tests/mocks/mockMutableKeyring.ts

@@ -0,0 +1,19 @@
+import { MockKeyring } from './mockKeyring';
+
+export class MockMutableKeyring extends MockKeyring {
+  static type = 'Mock Mutable Keyring';
+
+  public type = 'Mock Mutable Keyring';
+
+  #updated: boolean = false;
+
+  update() {
+    this.#updated = true;
+  }
+
+  async serialize() {
+    return Promise.resolve({
+      updated: this.#updated,
+    });
+  }
+}

yarn.lock

@@ -3553,6 +3553,7 @@ __metadata:
     immer: "npm:^9.0.6"
     jest: "npm:^27.5.1"
     jest-environment-node: "npm:^27.5.1"
+    lodash: "npm:^4.17.21"
     sinon: "npm:^9.2.4"
     ts-jest: "npm:^27.1.4"
     typedoc: "npm:^0.24.8"