perf(keyring-controller): do not fire unnecessary :stageChange in withKeyring

packages/keyring-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Prevent emitting `:stateChange` from `withKeyring` unnecessarily ([#5732](https://github.com/MetaMask/core/pull/5732))
+
 ## [21.0.5]
 
 ### Changed

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

@@ -262,14 +262,26 @@ describe('KeyringController', () => {
     });
 
     it('should throw error if the account is duplicated', async () => {
-      jest
-        .spyOn(HdKeyring.prototype, 'addAccounts')
-        .mockResolvedValue(['0x123']);
-      jest.spyOn(HdKeyring.prototype, 'getAccounts').mockReturnValue(['0x123']);
+      const mockAddress = '0x123';
+      const addAccountsSpy = jest.spyOn(HdKeyring.prototype, 'addAccounts');
+      const getAccountsSpy = jest.spyOn(HdKeyring.prototype, 'getAccounts');
+      const serializeSpy = jest.spyOn(HdKeyring.prototype, 'serialize');
+
+      addAccountsSpy.mockResolvedValue([mockAddress]);
+      getAccountsSpy.mockReturnValue([mockAddress]);
       await withController(async ({ controller }) => {
-        jest
-          .spyOn(HdKeyring.prototype, 'getAccounts')
-          .mockReturnValue(['0x123', '0x123']);
+        getAccountsSpy.mockReturnValue([mockAddress, mockAddress]);
+        serializeSpy
+          .mockResolvedValueOnce({
+            mnemonic: '',
+            numberOfAccounts: 1,
+            hdPath: "m/44'/60'/0'/0",
+          })
+          .mockResolvedValueOnce({
+            mnemonic: '',
+            numberOfAccounts: 2,
+            hdPath: "m/44'/60'/0'/0",
+          });
         await expect(controller.addNewAccount()).rejects.toThrow(
           KeyringControllerError.DuplicatedAccount,
         );
@@ -315,6 +327,11 @@ describe('KeyringController', () => {
               MockShallowGetAccountsKeyring.type,
             )[0] as EthKeyring;
 
+            jest
+              .spyOn(mockKeyring, 'serialize')
+              .mockResolvedValueOnce({ numberOfAccounts: 1 })
+              .mockResolvedValueOnce({ numberOfAccounts: 2 });
+
             const addedAccountAddress =
               await controller.addNewAccountForKeyring(mockKeyring);
 
@@ -3115,6 +3132,74 @@ describe('KeyringController', () => {
             },
           );
         });
+
+        it('should update the vault if the keyring is being updated', async () => {
+          const mockAddress = '0x4584d2B4905087A100420AFfCe1b2d73fC69B8E4';
+          stubKeyringClassWithAccount(MockKeyring, mockAddress);
+          await withController(
+            { keyringBuilders: [keyringBuilderFactory(MockKeyring)] },
+            async ({ controller, messenger }) => {
+              const selector = { type: MockKeyring.type };
+
+              await controller.addNewKeyring(MockKeyring.type);
+              const serializeSpy = jest.spyOn(
+                MockKeyring.prototype,
+                'serialize',
+              );
+              serializeSpy.mockResolvedValueOnce({
+                foo: 'bar', // Initial keyring state.
+              });
+
+              const mockStateChange = jest.fn();
+              messenger.subscribe(
+                'KeyringController:stateChange',
+                mockStateChange,
+              );
+
+              await controller.withKeyring(selector, async () => {
+                serializeSpy.mockResolvedValueOnce({
+                  foo: 'zzz', // Mock keyring state change.
+                });
+              });
+
+              expect(mockStateChange).toHaveBeenCalled();
+            },
+          );
+        });
+
+        it('should not update the vault if the keyring has not been updated', async () => {
+          const mockAddress = '0x4584d2B4905087A100420AFfCe1b2d73fC69B8E4';
+          stubKeyringClassWithAccount(MockKeyring, mockAddress);
+          await withController(
+            {
+              keyringBuilders: [keyringBuilderFactory(MockKeyring)],
+            },
+            async ({ controller, messenger }) => {
+              const selector = { type: MockKeyring.type };
+
+              await controller.addNewKeyring(MockKeyring.type);
+              const serializeSpy = jest.spyOn(
+                MockKeyring.prototype,
+                'serialize',
+              );
+              serializeSpy.mockResolvedValue({
+                foo: 'bar', // Initial keyring state.
+              });
+
+              const mockStateChange = jest.fn();
+              messenger.subscribe(
+                'KeyringController:stateChange',
+                mockStateChange,
+              );
+
+              await controller.withKeyring(selector, async () => {
+                // No-op, keyring state won't be updated.
+              });
+
+              expect(mockStateChange).not.toHaveBeenCalled();
+            },
+          );
+        });
       });
     });
 

packages/keyring-controller/src/KeyringController.ts

@@ -37,6 +37,7 @@ import { Mutex } from 'async-mutex';
 import type { MutexInterface } from 'async-mutex';
 import Wallet, { thirdparty as importers } from 'ethereumjs-wallet';
 import type { Patch } from 'immer';
+import { isEqual } from 'lodash';
 // When generating a ULID within the same millisecond, monotonicFactory provides some guarantees regarding sort order.
 import { ulid } from 'ulid';
 
@@ -320,6 +321,15 @@ export type SerializedKeyring = {
   data: Json;
 };
 
+/**
+ * State/data that can be updated during a `withKeyring` operation.
+ */
+type SessionState = {
+  keyrings: SerializedKeyring[];
+  keyringsMetadata: KeyringMetadata[];
+  password?: string;
+};
+
 /**
  * A generic encryptor interface that supports encrypting and decrypting
  * serializable data with a password.
@@ -1027,8 +1037,12 @@ export class KeyringController extends BaseController<
    * operation completes.
    */
   async persistAllKeyrings(): Promise<boolean> {
-    this.#assertIsUnlocked();
-    return this.#persistOrRollback(async () => true);
+    return this.#withRollback(async () => {
+      this.#assertIsUnlocked();
+
+      await this.#updateVault();
+      return true;
+    });
   }
 
   /**
@@ -1399,6 +1413,7 @@ export class KeyringController extends BaseController<
    */
   changePassword(password: string): Promise<void> {
     this.#assertIsUnlocked();
+
     return this.#persistOrRollback(async () => {
       assertIsValidPassword(password);
 
@@ -2158,6 +2173,20 @@ export class KeyringController extends BaseController<
     return serializedKeyrings;
   }
 
+  /**
+   * Get a snapshot of session data held by class variables.
+   *
+   * @returns An object with serialized keyrings, keyrings metadata,
+   * and the user password.
+   */
+  async #getSessionState(): Promise<SessionState> {
+    return {
+      keyrings: await this.#getSerializedKeyrings(),
+      keyringsMetadata: this.#keyringsMetadata.slice(), // Force copy.
+      password: this.#password,
+    };
+  }
+
   /**
    * Restore a serialized keyrings array.
    *
@@ -2635,7 +2664,7 @@ export class KeyringController extends BaseController<
 
   /**
    * Execute the given function after acquiring the controller lock
-   * and save the keyrings to state after it, or rollback to their
+   * and save the vault to state after it (only if needed), or rollback to their
    * previous state in case of error.
    *
    * @param callback - The function to execute.
@@ -2645,9 +2674,14 @@ export class KeyringController extends BaseController<
     callback: MutuallyExclusiveCallback<Result>,
   ): Promise<Result> {
     return this.#withRollback(async ({ releaseLock }) => {
+      const oldState = await this.#getSessionState();
       const callbackResult = await callback({ releaseLock });
-      // State is committed only if the operation is successful
-      await this.#updateVault();
+      const newState = await this.#getSessionState();
+
+      // State is committed only if the operation is successful and need to trigger a vault update.
+      if (!isEqual(oldState, newState)) {
+        await this.#updateVault();
+      }
 
       return callbackResult;
     });

yarn.lock

@@ -3555,6 +3555,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"