Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions packages/multichain-account-service/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@ class MockBtcKeyring {

return account;
});

createAccounts: SnapKeyring['createAccounts'] = jest.fn();
}
class MockBtcAccountProvider extends BtcAccountProvider {
override async ensureCanUseSnapPlatform(): Promise<void> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,7 @@ const setup = ({

const keyring = {
createAccount: jest.fn(),
createAccounts: jest.fn(),
removeAccount: jest.fn(),
};

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
};

Expand Down Expand Up @@ -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).
Expand All @@ -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 }) => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -115,6 +145,7 @@ function setup({
handleRequest: jest.Mock;
keyring: {
createAccount: jest.Mock;
createAccounts: jest.Mock;
};
trace: jest.Mock;
};
Expand Down Expand Up @@ -192,6 +223,7 @@ function setup({
handleRequest: mockHandleRequest,
keyring: {
createAccount: keyring.createAccount as jest.Mock,
createAccounts: keyring.createAccounts as jest.Mock,
},
trace: mockTrace,
},
Expand Down Expand Up @@ -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,
});

Expand All @@ -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 () => {
Expand All @@ -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,
});

Expand All @@ -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(
Expand All @@ -315,7 +358,7 @@ describe('SolAccountProvider', () => {
});

it('creates accounts with range starting from 0', async () => {
const { provider, keyring } = setup({
const { provider, mocks } = setup({
accounts: [],
});

Expand All @@ -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: [],
});

Expand All @@ -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,
Expand All @@ -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);
});
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
});
});
});
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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([]);
}
Expand Down
17 changes: 17 additions & 0 deletions packages/multichain-account-service/src/tests/providers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import type {
KeyringAccount,
KeyringCapabilities,
} from '@metamask/keyring-api';
import { GroupIndexRange } from 'src/utils';

import {
AccountProviderWrapper,
Expand Down Expand Up @@ -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);
}
Loading