feat(multichain-account-service): add Solana account batching

Author: ccharly

packages/multichain-account-service/CHANGELOG.md

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Optimize `SolAccountProvider.createAccounts` for range operations ([#8131](https://github.com/MetaMask/core/pull/8131))
+  - Batch account creation with the new `SnapKeyring.createAccounts` method.
+  - Significantly reduces lock acquisitions and API calls for batch operations.
 - Optimize `EvmAccountProvider.createAccounts` for range operations ([#7801](https://github.com/MetaMask/core/pull/7801))
   - Batch account creation with single a `withKeyring` call for entire range instead of one call per account.
   - Batch account creation with single `keyring.addAccounts` call.

packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts

@@ -98,6 +98,8 @@ class MockBtcKeyring {
 
       return account;
     });
+
+  createAccounts: SnapKeyring['createAccounts'] = jest.fn();
 }
 class MockBtcAccountProvider extends BtcAccountProvider {
   override async ensureCanUseSnapPlatform(): Promise<void> {

packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts

@@ -213,6 +213,7 @@ const setup = ({
 
   const keyring = {
     createAccount: jest.fn(),
+    createAccounts: jest.fn(),
     removeAccount: jest.fn(),
   };
 

packages/multichain-account-service/src/providers/SnapAccountProvider.ts

@@ -22,6 +22,7 @@ import { createSentryError } from '../utils';
 
 export type RestrictedSnapKeyring = {
   createAccount: (options: Record<string, Json>) => Promise<KeyringAccount>;
+  createAccounts: (options: CreateAccountOptions) => Promise<KeyringAccount[]>;
   removeAccount: (address: string) => Promise<void>;
 };
 
@@ -126,6 +127,10 @@ export abstract class SnapAccountProvider extends BaseBip44AccountProvider {
       SnapKeyring['createAccount']
     >(async ({ keyring }) => keyring.createAccount.bind(keyring));
 
+    const createAccounts = await this.#withSnapKeyring<
+      SnapKeyring['createAccounts']
+    >(async ({ keyring }) => keyring.createAccounts.bind(keyring));
+
     return {
       createAccount: async (options) =>
         // We use the "unguarded" account creation here (see explanation above).
@@ -134,6 +139,8 @@ export abstract class SnapAccountProvider extends BaseBip44AccountProvider {
           displayConfirmation: false,
           setSelectedAccount: false,
         }),
+      createAccounts: async (options) =>
+        await createAccounts(this.snapId, options),
       removeAccount: async (address: string) =>
         // Though, when removing account, we can use the normal flow.
         await this.#withSnapKeyring(async ({ keyring }) => {

packages/multichain-account-service/src/providers/SolAccountProvider.test.ts

@@ -19,6 +19,7 @@ import { TraceName } from '../constants/traces';
 import {
   getMultichainAccountServiceMessenger,
   getRootMessenger,
+  toGroupIndexRangeArray,
   MOCK_HD_ACCOUNT_1,
   MOCK_HD_KEYRING_1,
   MOCK_SOL_ACCOUNT_1,
@@ -82,6 +83,35 @@ class MockSolanaKeyring {
 
       return account;
     });
+
+  createAccounts: SnapKeyring['createAccounts'] = jest
+    .fn()
+    .mockImplementation((_, options) => {
+      const groupIndices =
+        options.type === 'bip44:derive-index'
+          ? [options.groupIndex]
+          : toGroupIndexRangeArray(options.range);
+
+      return groupIndices.map((groupIndex) => {
+        const found = this.accounts.find(
+          (account) =>
+            isBip44Account(account) &&
+            account.options.entropy.groupIndex === groupIndex,
+        );
+
+        if (found) {
+          return found; // Idempotent.
+        }
+
+        const account = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
+          .withUuid()
+          .withAddressSuffix(`${groupIndex}`)
+          .withGroupIndex(groupIndex)
+          .get();
+        this.accounts.push(account);
+        return account;
+      });
+    });
 }
 
 class MockSolAccountProvider extends SolAccountProvider {
@@ -115,6 +145,7 @@ function setup({
     handleRequest: jest.Mock;
     keyring: {
       createAccount: jest.Mock;
+      createAccounts: jest.Mock;
     };
     trace: jest.Mock;
   };
@@ -192,6 +223,7 @@ function setup({
       handleRequest: mockHandleRequest,
       keyring: {
         createAccount: keyring.createAccount as jest.Mock,
+        createAccounts: keyring.createAccounts as jest.Mock,
       },
       trace: mockTrace,
     },
@@ -252,7 +284,7 @@ describe('SolAccountProvider', () => {
 
   it('creates accounts', async () => {
     const accounts = [MOCK_SOL_ACCOUNT_1];
-    const { provider, keyring } = setup({
+    const { provider, mocks } = setup({
       accounts,
     });
 
@@ -263,7 +295,16 @@ describe('SolAccountProvider', () => {
       groupIndex: newGroupIndex,
     });
     expect(newAccounts).toHaveLength(1);
-    expect(keyring.createAccount).toHaveBeenCalled();
+    // Batch endpoint must be called, NOT the singular one.
+    expect(mocks.keyring.createAccounts).toHaveBeenCalledWith(
+      SolAccountProvider.SOLANA_SNAP_ID,
+      {
+        type: AccountCreationType.Bip44DeriveIndex,
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: newGroupIndex,
+      },
+    );
+    expect(mocks.keyring.createAccount).not.toHaveBeenCalled();
   });
 
   it('does not re-create accounts (idempotent)', async () => {
@@ -283,7 +324,7 @@ describe('SolAccountProvider', () => {
 
   it('creates multiple accounts using Bip44DeriveIndexRange', async () => {
     const accounts = [MOCK_SOL_ACCOUNT_1];
-    const { provider, keyring } = setup({
+    const { provider, mocks } = setup({
       accounts,
     });
 
@@ -297,7 +338,9 @@ describe('SolAccountProvider', () => {
     });
 
     expect(newAccounts).toHaveLength(3);
-    expect(keyring.createAccount).toHaveBeenCalledTimes(3);
+    // Single batch call, NOT three individual calls.
+    expect(mocks.keyring.createAccounts).toHaveBeenCalledTimes(1);
+    expect(mocks.keyring.createAccount).not.toHaveBeenCalled();
 
     // Verify each account has the correct group index
     expect(
@@ -315,7 +358,7 @@ describe('SolAccountProvider', () => {
   });
 
   it('creates accounts with range starting from 0', async () => {
-    const { provider, keyring } = setup({
+    const { provider, mocks } = setup({
       accounts: [],
     });
 
@@ -329,11 +372,12 @@ describe('SolAccountProvider', () => {
     });
 
     expect(newAccounts).toHaveLength(3);
-    expect(keyring.createAccount).toHaveBeenCalledTimes(3);
+    expect(mocks.keyring.createAccounts).toHaveBeenCalledTimes(1);
+    expect(mocks.keyring.createAccount).not.toHaveBeenCalled();
   });
 
   it('creates a single account when range from equals to', async () => {
-    const { provider, keyring } = setup({
+    const { provider, mocks } = setup({
       accounts: [],
     });
 
@@ -347,7 +391,8 @@ describe('SolAccountProvider', () => {
     });
 
     expect(newAccounts).toHaveLength(1);
-    expect(keyring.createAccount).toHaveBeenCalledTimes(1);
+    expect(mocks.keyring.createAccounts).toHaveBeenCalledTimes(1);
+    expect(mocks.keyring.createAccount).not.toHaveBeenCalled();
     expect(
       isBip44Account(newAccounts[0]) &&
         newAccounts[0].options.entropy.groupIndex,
@@ -359,10 +404,10 @@ describe('SolAccountProvider', () => {
       accounts: [],
     });
 
-    mocks.keyring.createAccount.mockImplementation(() => {
+    mocks.keyring.createAccounts.mockImplementation(() => {
       return new Promise((resolve) => {
         setTimeout(() => {
-          resolve(MOCK_SOL_ACCOUNT_1);
+          resolve([MOCK_SOL_ACCOUNT_1]);
         }, 4000);
       });
     });

packages/multichain-account-service/src/providers/SolAccountProvider.ts

@@ -113,52 +113,50 @@ export class SolAccountProvider extends SnapAccountProvider {
       `${AccountCreationType.Bip44DeriveIndexRange}`,
     ]);
 
-    const { entropySource } = options;
-
-    if (options.type === AccountCreationType.Bip44DeriveIndexRange) {
-      const { range } = options;
-      return this.withSnap(async ({ keyring }) => {
-        const accounts: Bip44Account<KeyringAccount>[] = [];
-
-        // TODO: Use `SnapKeyring.createAccounts` when that functionality is implemented on the Snap
-        // itself, instead of creating accounts one by one.
-        for (
-          let groupIndex = range.from;
-          groupIndex <= range.to;
-          groupIndex++
-        ) {
-          const account = await this.withMaxConcurrency(async () => {
-            const derivationPath = `m/44'/501'/${groupIndex}'/0'`;
-            return this.#createAccount({
-              keyring,
-              entropySource,
-              groupIndex,
-              derivationPath,
-            });
-          });
-
-          this.accounts.add(account.id);
-          accounts.push(account);
+    return this.withSnap(async ({ keyring }) => {
+      return this.withMaxConcurrency(async () => {
+        let groupIndexOffset = 0;
+        let snapAccounts: KeyringAccount[];
+
+        const { entropySource } = options;
+
+        if (options.type === `${AccountCreationType.Bip44DeriveIndexRange}`) {
+          snapAccounts = await withTimeout(
+            keyring.createAccounts(options),
+            this.config.createAccounts.timeoutMs,
+          );
+
+          // Group indices are sequential, so we just need the starting index.
+          groupIndexOffset = options.range.from;
+        } else {
+          const [snapAccount] = await withTimeout(
+            keyring.createAccounts(options),
+            this.config.createAccounts.timeoutMs,
+          );
+          snapAccounts = [snapAccount];
+
+          // For single account, there will only be 1 account, so we can use the
+          // provided group index directly.
+          groupIndexOffset = options.groupIndex;
         }
 
-        return accounts;
-      });
-    }
+        // NOTE: We still need to convert accounts to proper BIP-44 accounts for now.
+        return snapAccounts.map((snapAccount, index) => {
+          const groupIndex = groupIndexOffset + index;
+          const derivationPath = `m/44'/501'/${groupIndex}'/0'`;
 
-    const { groupIndex } = options;
+          // Ensure entropy is present before type assertion validation
+          snapAccount.options.entropy = {
+            type: KeyringAccountEntropyTypeOption.Mnemonic,
+            id: entropySource,
+            groupIndex,
+            derivationPath,
+          };
 
-    return this.withSnap(async ({ keyring }) => {
-      return this.withMaxConcurrency(async () => {
-        const derivationPath = `m/44'/501'/${groupIndex}'/0'`;
-        const account = await this.#createAccount({
-          keyring,
-          entropySource,
-          groupIndex,
-          derivationPath,
-        });
+          assertIsBip44Account(snapAccount);
 
-        this.accounts.add(account.id);
-        return [account];
+          return snapAccount;
+        });
       });
     });
   }

packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts

@@ -69,6 +69,8 @@ class MockTronKeyring {
       return account;
     });
 
+  createAccounts: SnapKeyring['createAccounts'] = jest.fn();
+
   // Add discoverAccounts method to match the provider's usage
   discoverAccounts = jest.fn().mockResolvedValue([]);
 }

packages/multichain-account-service/src/tests/providers.ts

@@ -12,6 +12,7 @@ import type {
   KeyringAccount,
   KeyringCapabilities,
 } from '@metamask/keyring-api';
+import { GroupIndexRange } from 'src/utils';
 
 import {
   AccountProviderWrapper,
@@ -201,3 +202,19 @@ export function mockCreateAccountsOnce(
     return created;
   });
 }
+
+/**
+ * Helper to convert a group index range to an array of group indices, inclusive of the
+ * start and end indices.
+ *
+ * @param range - The range.
+ * @param range.from - The starting index of the range (inclusive).
+ * @param range.to - The ending index of the range (inclusive).
+ * @returns An array of group indices from `from` to `to`, inclusive.
+ */
+export function toGroupIndexRangeArray({
+  from = 0,
+  to,
+}: GroupIndexRange): number[] {
+  return Array.from({ length: to - from + 1 }, (_, i) => from + i);
+}