diff --git a/packages/multichain-account-service/CHANGELOG.md b/packages/multichain-account-service/CHANGELOG.md index 75a85c259ab..61ebfa0a28a 100644 --- a/packages/multichain-account-service/CHANGELOG.md +++ b/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. diff --git a/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts b/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts index a118295b2ce..2d3b5fe113d 100644 --- a/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts +++ b/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 { diff --git a/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts b/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts index 4c987169b29..ccd44702190 100644 --- a/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts +++ b/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(), }; diff --git a/packages/multichain-account-service/src/providers/SnapAccountProvider.ts b/packages/multichain-account-service/src/providers/SnapAccountProvider.ts index dba319101aa..bee1f2e1aaf 100644 --- a/packages/multichain-account-service/src/providers/SnapAccountProvider.ts +++ b/packages/multichain-account-service/src/providers/SnapAccountProvider.ts @@ -22,6 +22,7 @@ import { createSentryError } from '../utils'; export type RestrictedSnapKeyring = { createAccount: (options: Record) => Promise; + createAccounts: (options: CreateAccountOptions) => Promise; removeAccount: (address: string) => Promise; }; @@ -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 }) => { diff --git a/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts b/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts index 4963f9d609e..1c1963d0a06 100644 --- a/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts +++ b/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); }); }); diff --git a/packages/multichain-account-service/src/providers/SolAccountProvider.ts b/packages/multichain-account-service/src/providers/SolAccountProvider.ts index 455ee26a102..63d0111f26b 100644 --- a/packages/multichain-account-service/src/providers/SolAccountProvider.ts +++ b/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[] = []; - - // 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; + }); }); }); } diff --git a/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts b/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts index 10bf6ce44e1..5a0b51929c6 100644 --- a/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts +++ b/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([]); } diff --git a/packages/multichain-account-service/src/tests/providers.ts b/packages/multichain-account-service/src/tests/providers.ts index d56f117ea7d..6a01ed717cb 100644 --- a/packages/multichain-account-service/src/tests/providers.ts +++ b/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); +}