fix(snap-account-service): faster :getLegacySnapKeyring

packages/snap-account-service/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Faster `:getLegacySnapKeyring` ([#8865](https://github.com/MetaMask/core/pull/8865))
+  - We now check if the keyring exists with `:withKeyringUnsafe` and returns it right away.
+  - If the keyring does not exist yet, we do create it with `:withController` (next calls will then be faster thanks to `:withKeyringUnsafe` pre-check).
+
 ### Fixed
 
 - Prevent double-lock in `:handleKeyringSnapMessage` for some events/methods ([#8860](https://github.com/MetaMask/core/pull/8860))

packages/snap-account-service/src/SnapAccountService.test.ts

@@ -630,40 +630,64 @@ describe('SnapAccountService', () => {
   });
 
   describe('getLegacySnapKeyring', () => {
-    it('returns the existing Snap keyring when one is already present', async () => {
+    it('returns the existing Snap keyring via the fast-path without acquiring the KeyringController mutex', async () => {
       const { service, mocks } = setup();
-      const existing = buildKeyringEntry(KeyringTypes.snap);
+      const existing = mockLegacySnapKeyring(mocks, {});
+
+      const result = await service.getLegacySnapKeyring();
+
+      expect(result).toBe(existing as unknown as SnapKeyring);
+      expect(mocks.KeyringController.withController).not.toHaveBeenCalled();
+    });
+
+    it('falls back to withController and creates a new Snap keyring when the fast-path reports it missing', async () => {
+      const { service, mocks } = setup();
+      mockLegacySnapKeyringMissing(mocks);
       const { addNewKeyring } = mockWithController(mocks, [
         buildKeyringEntry(KeyringTypes.hd),
-        existing,
       ]);
 
       const result = await service.getLegacySnapKeyring();
 
-      expect(result).toBe(existing.keyring as unknown as SnapKeyring);
-      expect(addNewKeyring).not.toHaveBeenCalled();
+      expect(mocks.KeyringController.withKeyringUnsafe).toHaveBeenCalled();
+      expect(addNewKeyring).toHaveBeenCalledWith(KeyringTypes.snap);
+      expect(result.type).toBe(KeyringTypes.snap);
     });
 
-    it('creates a new Snap keyring when none exists', async () => {
+    it('returns the existing Snap keyring found within withController when the fast-path reports it missing', async () => {
       const { service, mocks } = setup();
+      mockLegacySnapKeyringMissing(mocks);
+      const existing = buildKeyringEntry(KeyringTypes.snap);
       const { addNewKeyring } = mockWithController(mocks, [
         buildKeyringEntry(KeyringTypes.hd),
+        existing,
       ]);
 
       const result = await service.getLegacySnapKeyring();
 
-      expect(addNewKeyring).toHaveBeenCalledWith(KeyringTypes.snap);
-      expect(result.type).toBe(KeyringTypes.snap);
+      expect(result).toBe(existing.keyring as unknown as SnapKeyring);
+      expect(addNewKeyring).not.toHaveBeenCalled();
     });
 
     it('propagates errors thrown by withController', async () => {
       const { service, mocks } = setup();
+      mockLegacySnapKeyringMissing(mocks);
       mocks.KeyringController.withController.mockImplementation(async () => {
         throw new Error('boom');
       });
 
       await expect(service.getLegacySnapKeyring()).rejects.toThrow('boom');
     });
+
+    it('propagates non-KeyringNotFound errors thrown by the fast-path', async () => {
+      const { service, mocks } = setup();
+      mocks.KeyringController.withKeyringUnsafe.mockImplementation(async () => {
+        throw new Error('boom');
+      });
+
+      await expect(service.getLegacySnapKeyring()).rejects.toThrow('boom');
+      expect(mocks.KeyringController.withController).not.toHaveBeenCalled();
+    });
   });
 
   describe('handleKeyringSnapMessage', () => {

packages/snap-account-service/src/SnapAccountService.ts

@@ -319,6 +319,15 @@ export class SnapAccountService {
       snapKeyring: LegacySnapKeyring;
     };
 
+    // This is a fast-path for the common case where the keyring already exists, to avoid the
+    // overhead of acquiring the `KeyringController` mutex if we don't need to.
+    // NOTE: If it doesn't exist, we'll create it **safely** with `:withController` (which was
+    // not the case with the previous client's implementation).
+    const exists = await this.#getLegacySnapKeyringIfAvailable();
+    if (exists) {
+      return exists;
+    }
+
     // `KeyringController:withController` forbids returning a direct keyring
     // reference (it checks the result via `Object.is`), so we smuggle the
     // instance out wrapped in an object and unwrap it after the call.