Skip to content
Closed
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
1 change: 1 addition & 0 deletions packages/assets-controller/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Changed

- `MulticallClient` memoizes `balanceOf` and `getEthBalance` call encodings per account address when building multicall batches, reducing redundant ABI encoding for wallets with many tokens ([#9423](https://github.com/MetaMask/core/pull/9423))
- Bump `@metamask/transaction-controller` from `^68.2.2` to `^68.3.0` ([#9421](https://github.com/MetaMask/core/pull/9421))

## [10.1.0]
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { defaultAbiCoder, Interface } from '@ethersproject/abi';
import * as controllerUtils from '@metamask/controller-utils';
import type { Hex } from '@metamask/utils';

import type { Address, BalanceOfRequest, ChainId, Provider } from '../types';
Expand Down Expand Up @@ -212,6 +213,42 @@ describe('MulticallClient', () => {
expect(result[1].balance).toBe('2000000000');
});

it('encodes balance call data once per account address in a batch', async () => {
const encodeSpy = jest.spyOn(controllerUtils, 'encodeFunctionData');
const otherAccount: Address =
'0xabcdefabcdefabcdefabcdefabcdefabcdefabcd' as Address;

const requests: BalanceOfRequest[] = [
{ tokenAddress: TEST_TOKEN_1, accountAddress: TEST_ACCOUNT },
{ tokenAddress: TEST_TOKEN_2, accountAddress: TEST_ACCOUNT },
{ tokenAddress: TEST_TOKEN_1, accountAddress: otherAccount },
{ tokenAddress: ZERO_ADDRESS, accountAddress: TEST_ACCOUNT },
];

const mockResponse = buildMockAggregate3Response([
{ success: true, balance: '1000000000' },
{ success: true, balance: '2000000000' },
{ success: true, balance: '3000000000' },
{ success: true, balance: '1000000000000000000' },
]);

mockProvider.call.mockResolvedValue(mockResponse);

await client.batchBalanceOf(MAINNET_CHAIN_ID, requests);

const balanceOfEncodings = encodeSpy.mock.calls.filter(
([, method]) => method === 'balanceOf',
);
const getEthBalanceEncodings = encodeSpy.mock.calls.filter(
([, method]) => method === 'getEthBalance',
);

expect(balanceOfEncodings).toHaveLength(2);
expect(getEthBalanceEncodings).toHaveLength(1);

encodeSpy.mockRestore();
});

it('should fetch native token balance using getEthBalance', async () => {
const requests: BalanceOfRequest[] = [
{ tokenAddress: ZERO_ADDRESS, accountAddress: TEST_ACCOUNT },
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,72 @@ function encodeGetEthBalance(accountAddress: Address): Hex {
]);
}

type BalanceCallDataCache = {
getErc20BalanceCallData: (accountAddress: Address) => Hex;
getNativeBalanceCallData: (accountAddress: Address) => Hex;
};

/**
* Cache balance call encodings per account address for a single batch request.
* ERC-20 `balanceOf` and native `getEthBalance` call data depend only on the
* account being queried, not the token contract (which is the multicall target).
*
* @returns A cache scoped to one `batchBalanceOf` invocation.
*/
function createBalanceCallDataCache(): BalanceCallDataCache {
const erc20ByAccount = new Map<string, Hex>();
const nativeByAccount = new Map<string, Hex>();

return {
getErc20BalanceCallData(accountAddress: Address): Hex {
const key = accountAddress.toLowerCase();
const existing = erc20ByAccount.get(key);
if (existing !== undefined) {
return existing;
}
const callData = encodeBalanceOf(accountAddress);
erc20ByAccount.set(key, callData);
return callData;
},
getNativeBalanceCallData(accountAddress: Address): Hex {
const key = accountAddress.toLowerCase();
const existing = nativeByAccount.get(key);
if (existing !== undefined) {
return existing;
}
const callData = encodeGetEthBalance(accountAddress);
nativeByAccount.set(key, callData);
return callData;
},
};
}

/**
* Build Multicall3 aggregate3 calls for a batch of balance requests.
*
* @param batch - Balance requests in the current batch.
* @param multicallAddress - Multicall3 contract address for native balance calls.
* @param callDataCache - Per-request cache for encoded balance call data.
* @returns Aggregate3 call descriptors.
*/
function buildAggregate3BalanceCalls(
batch: BalanceOfRequest[],
multicallAddress: Hex,
callDataCache: BalanceCallDataCache,
): { target: Address; allowFailure: boolean; callData: Hex }[] {
return batch.map((req) => {
const isNative = req.tokenAddress === ZERO_ADDRESS;
const target = isNative ? multicallAddress : req.tokenAddress;
return {
target,
allowFailure: true,
callData: isNative
? callDataCache.getNativeBalanceCallData(req.accountAddress)
: callDataCache.getErc20BalanceCallData(req.accountAddress),
};
});
}

/**
* Encode a Multicall3 aggregate3 call.
*
Expand Down Expand Up @@ -527,10 +593,11 @@ export class MulticallClient {

const multicallAddress = MULTICALL3_ADDRESS_BY_CHAIN[chainId];
const provider = this.#getProvider(chainId);
const callDataCache = createBalanceCallDataCache();

if (!multicallAddress) {
return options.fallbackToSingleCalls
? this.#fallbackBatchBalanceOf(provider, requests)
? this.#fallbackBatchBalanceOf(provider, requests, callDataCache)
: this.#createFailedResponses(requests);
}

Expand All @@ -540,6 +607,7 @@ export class MulticallClient {
multicallAddress,
requests,
options.fallbackToSingleCalls,
callDataCache,
);
}

Expand All @@ -550,13 +618,15 @@ export class MulticallClient {
* @param multicallAddress - The Multicall3 contract address.
* @param requests - Array of balance requests.
* @param fallbackToSingleCalls - Whether to fall back to individual RPC calls on batch failure.
* @param callDataCache - The cache for encoded balance call data.
* @returns Array of balance responses.
*/
async #multicallBatchBalanceOf(
provider: Provider,
multicallAddress: Hex,
requests: BalanceOfRequest[],
fallbackToSingleCalls: boolean,
callDataCache: BalanceCallDataCache,
): Promise<BalanceOfResponse[]> {
const batchSize = this.#config.maxCallsPerBatch;

Expand All @@ -568,23 +638,16 @@ export class MulticallClient {
batchSize,
initialResult: [],
eachBatch: async (workingResult, batch) => {
const calls = buildAggregate3BalanceCalls(
batch,
multicallAddress,
callDataCache,
);

try {
await createServicePolicy({
maxRetries: MULTICALL_MAX_RETRIES,
}).execute(async () => {
// Build aggregate3 calls
const calls = batch.map((req) => {
const isNative = req.tokenAddress === ZERO_ADDRESS;
const target = isNative ? multicallAddress : req.tokenAddress;
return {
target,
allowFailure: true,
callData: isNative
? encodeGetEthBalance(req.accountAddress)
: encodeBalanceOf(req.accountAddress),
};
});

// Encode and send aggregate3 call
const callData = encodeAggregate3(calls);
const result = await provider.call({
Expand Down Expand Up @@ -625,7 +688,9 @@ export class MulticallClient {
// #fetchSingleBalance never rejects - it catches all errors internally
// and returns a failed response, so we use Promise.all here.
const fallbackResults = await Promise.all(
batch.map((req) => this.#fetchSingleBalance(provider, req)),
batch.map((req) =>
this.#fetchSingleBalance(provider, req, callDataCache),
),
);

for (const result of fallbackResults) {
Expand All @@ -649,10 +714,12 @@ export class MulticallClient {
* @param provider - The RPC provider.
* @param requests - Array of balance requests.
* @returns Array of balance responses.
* @param callDataCache - The cache for encoded balance call data.
*/
async #fallbackBatchBalanceOf(
provider: Provider,
requests: BalanceOfRequest[],
callDataCache: BalanceCallDataCache,
): Promise<BalanceOfResponse[]> {
// Use smaller batch size for parallel individual calls to avoid overwhelming RPC
const batchSize = Math.min(this.#config.maxCallsPerBatch, 50);
Expand All @@ -668,7 +735,9 @@ export class MulticallClient {
// #fetchSingleBalance never rejects - it catches all errors internally
// and returns a failed response, so we use Promise.all here.
const batchResults = await Promise.all(
batch.map((req) => this.#fetchSingleBalance(provider, req)),
batch.map((req) =>
this.#fetchSingleBalance(provider, req, callDataCache),
),
);

for (const result of batchResults) {
Expand All @@ -688,10 +757,12 @@ export class MulticallClient {
* @param provider - The RPC provider.
* @param request - The balance request.
* @returns The balance response.
* @param callDataCache - The cache for encoded balance call data.
*/
async #fetchSingleBalance(
provider: Provider,
request: BalanceOfRequest,
callDataCache: BalanceCallDataCache,
): Promise<BalanceOfResponse> {
// Destructure inside try block to ensure any errors are caught
// and don't cause promise rejections that bypass error handling
Expand All @@ -710,7 +781,7 @@ export class MulticallClient {
}

// ERC-20 token
const callData = encodeBalanceOf(accountAddress);
const callData = callDataCache.getErc20BalanceCallData(accountAddress);
const result = await provider.call({
to: tokenAddress,
data: callData,
Expand Down
Loading