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

packages/keyring-controller/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Prevent emitting `:stateChange` from `withKeyring` unnecessarily ([#5732](https://github.com/MetaMask/core/pull/5732))
 - Bump `@metamask/base-controller` from ^8.0.0 to ^8.0.1 ([#5722](https://github.com/MetaMask/core/pull/5722))
 
 ## [21.0.4]

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);
 
@@ -3080,6 +3097,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

@@ -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';
@@ -320,6 +321,13 @@ export type SerializedKeyring = {
   data: Json;
 };
 
+/**
+ * A serialized keyring and metadata object.
+ */
+export type SerializedKeyringAndMetadata = SerializedKeyring & {
+  metadata?: KeyringMetadata;
+};
+
 /**
  * A generic encryptor interface that supports encrypting and decrypting
  * serializable data with a password.
@@ -1028,7 +1036,8 @@ export class KeyringController extends BaseController<
    */
   async persistAllKeyrings(): Promise<boolean> {
     this.#assertIsUnlocked();
-    return this.#persistOrRollback(async () => true);
+
+    return this.#persistOrRollback(async () => true, { forceUpdate: true });
   }
 
   /**
@@ -1399,20 +1408,24 @@ export class KeyringController extends BaseController<
    */
   changePassword(password: string): Promise<void> {
     this.#assertIsUnlocked();
-    return this.#persistOrRollback(async () => {
-      assertIsValidPassword(password);
 
-      this.#password = password;
-      // We need to clear encryption key and salt from state
-      // to force the controller to re-encrypt the vault using
-      // the new password.
-      if (this.#cacheEncryptionKey) {
-        this.update((state) => {
-          delete state.encryptionKey;
-          delete state.encryptionSalt;
-        });
-      }
-    });
+    return this.#persistOrRollback(
+      async () => {
+        assertIsValidPassword(password);
+
+        this.#password = password;
+        // We need to clear encryption key and salt from state
+        // to force the controller to re-encrypt the vault using
+        // the new password.
+        if (this.#cacheEncryptionKey) {
+          this.update((state) => {
+            delete state.encryptionKey;
+            delete state.encryptionSalt;
+          });
+        }
+      },
+      { forceUpdate: true },
+    );
   }
 
   /**
@@ -2146,6 +2159,38 @@ export class KeyringController extends BaseController<
     return serializedKeyrings;
   }
 
+  /**
+   * Serialize the current array of keyring instances and their metadata,
+   * including unsupported keyrings by default.
+   *
+   * @param options - Method options.
+   * @param options.includeUnsupported - Whether to include unsupported keyrings.
+   * @returns The serialized keyrings.
+   */
+  async #getSerializedKeyringsAndMetadata(
+    { includeUnsupported }: { includeUnsupported: boolean } = {
+      includeUnsupported: true,
+    },
+  ): Promise<SerializedKeyringAndMetadata[]> {
+    const serializedKeyrings = await this.#getSerializedKeyrings({
+      includeUnsupported,
+    });
+
+    const serializedKeyringsAndMetadata: SerializedKeyringAndMetadata[] =
+      serializedKeyrings.map((serialized, index) => {
+        return {
+          ...serialized,
+          metadata: this.#keyringsMetadata[index],
+        };
+      });
+
+    if (includeUnsupported) {
+      serializedKeyringsAndMetadata.push(...this.#unsupportedKeyrings);
+    }
+
+    return serializedKeyringsAndMetadata;
+  }
+
   /**
    * Restore a serialized keyrings array.
    *
@@ -2617,19 +2662,31 @@ 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.
+   * @param options - Options.
+   * @param options.forceUpdate - Force the vault update.
    * @returns The result of the function.
    */
   async #persistOrRollback<Result>(
     callback: MutuallyExclusiveCallback<Result>,
+    { forceUpdate }: { forceUpdate: boolean } = { forceUpdate: false },
   ): Promise<Result> {
     return this.#withRollback(async ({ releaseLock }) => {
+      const oldState = !forceUpdate
+        ? await this.#getSerializedKeyringsAndMetadata()
+        : []; // No need to serialize anything when forcing the update.
       const callbackResult = await callback({ releaseLock });
-      // State is committed only if the operation is successful
-      await this.#updateVault();
+      const newState = !forceUpdate
+        ? await this.#getSerializedKeyringsAndMetadata()
+        : []; // Same.
+
+      // State is committed only if the operation is successful and need to trigger a vault update.
+      if (forceUpdate || !isEqual(oldState, newState)) {
+        await this.#updateVault();
+      }
 
       return callbackResult;
     });

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"