Skip to content

Commit b5ae5ac

Browse files
salimtbsalimtb
committed
feat: parallelize balance fetch
1 parent 6cccc28 commit b5ae5ac

File tree

8 files changed

+370
-133
lines changed

8 files changed

+370
-133
lines changed

‎packages/assets-controllers/src/AccountTrackerController.test.ts‎

Lines changed: 35 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ jest.mock('@metamask/controller-utils', () => {
3131
return {
3232
...jest.requireActual('@metamask/controller-utils'),
3333
query: jest.fn(),
34+
safelyExecuteWithTimeout: jest.fn(),
3435
};
3536
});
3637

@@ -57,17 +58,34 @@ const mockedQuery = query as jest.Mock<
5758
Parameters<typeof query>
5859
>;
5960

61+
const { safelyExecuteWithTimeout } = jest.requireMock(
62+
'@metamask/controller-utils',
63+
);
64+
const mockedSafelyExecuteWithTimeout = safelyExecuteWithTimeout as jest.Mock;
65+
6066
describe('AccountTrackerController', () => {
6167
let clock: SinonFakeTimers;
6268

6369
beforeEach(() => {
6470
clock = useFakeTimers();
6571
mockedQuery.mockReturnValue(Promise.resolve('0x0'));
72+
73+
// Mock safelyExecuteWithTimeout to execute the operation normally by default
74+
mockedSafelyExecuteWithTimeout.mockImplementation(
75+
async (operation: () => Promise<unknown>) => {
76+
try {
77+
return await operation();
78+
} catch {
79+
return undefined;
80+
}
81+
},
82+
);
6683
});
6784

6885
afterEach(() => {
6986
clock.restore();
7087
mockedQuery.mockRestore();
88+
mockedSafelyExecuteWithTimeout.mockRestore();
7189
});
7290

7391
it('should set default state', async () => {
@@ -664,7 +682,6 @@ describe('AccountTrackerController', () => {
664682
});
665683

666684
it('should handle timeout error correctly', async () => {
667-
const originalSetTimeout = global.setTimeout;
668685
const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
669686

670687
await withController(
@@ -683,32 +700,28 @@ describe('AccountTrackerController', () => {
683700
selectedAccount: ACCOUNT_1,
684701
listAccounts: [ACCOUNT_1, ACCOUNT_2],
685702
},
686-
async ({ refresh }) => {
687-
// Mock setTimeout to immediately trigger the timeout callback
688-
global.setTimeout = ((callback: () => void, _delay: number) => {
689-
// This is the timeout callback from line 657 - trigger it immediately
690-
originalSetTimeout(callback, 0);
691-
return 123 as unknown as NodeJS.Timeout; // Return a fake timer id
692-
}) as typeof setTimeout;
693-
694-
// Mock the query to hang indefinitely
695-
const hangingPromise = new Promise(() => {
696-
// Intentionally empty - simulates hanging request
697-
});
698-
mockedQuery.mockReturnValue(hangingPromise);
703+
async ({ refresh, controller }) => {
704+
// Mock safelyExecuteWithTimeout to simulate timeout by returning undefined
705+
mockedSafelyExecuteWithTimeout.mockImplementation(
706+
async () => undefined, // Simulates timeout behavior
707+
);
699708

700-
// Start refresh and let the timeout trigger
709+
// Start refresh with the mocked timeout behavior
701710
await refresh(clock, ['mainnet']);
702711

703-
// Verify that the timeout error was logged (confirms line 657 was executed)
704-
expect(consoleWarnSpy).toHaveBeenCalledWith(
705-
expect.stringContaining(
706-
'Balance fetcher failed for chains 0x1: Error: Timeout after 15000ms',
707-
),
712+
// With safelyExecuteWithTimeout, timeouts are handled gracefully
713+
// The system should continue operating without throwing errors
714+
// No specific timeout error message should be logged
715+
expect(consoleWarnSpy).not.toHaveBeenCalledWith(
716+
expect.stringContaining('Timeout after'),
717+
);
718+
719+
// Verify that the controller state remains intact despite the timeout
720+
expect(controller.state.accountsByChainId).toHaveProperty('0x1');
721+
expect(controller.state.accountsByChainId['0x1']).toHaveProperty(
722+
CHECKSUM_ADDRESS_1,
708723
);
709724

710-
// Restore original setTimeout
711-
global.setTimeout = originalSetTimeout;
712725
consoleWarnSpy.mockRestore();
713726
},
714727
);

‎packages/assets-controllers/src/AccountTrackerController.ts‎

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ const controllerName = 'AccountTrackerController';
5353
export type ChainIdHex = Hex;
5454
export type ChecksumAddress = Hex;
5555

56-
const DEFAULT_TIMEOUT_MS = 15000;
5756
const ZERO_ADDRESS =
5857
'0x0000000000000000000000000000000000000000' as ChecksumAddress;
5958

@@ -92,9 +91,8 @@ class AccountTrackerRpcBalanceFetcher implements BalanceFetcher {
9291
selectedAccount,
9392
allAccounts,
9493
}: Parameters<BalanceFetcher['fetch']>[0]): Promise<ProcessedBalance[]> {
95-
const results: ProcessedBalance[] = [];
96-
97-
for (const chainId of chainIds) {
94+
// Process all chains in parallel for better performance
95+
const chainProcessingPromises = chainIds.map(async (chainId) => {
9896
const accountsToUpdate = queryAllAccounts
9997
? Object.values(allAccounts).map(
10098
(account) =>
@@ -104,6 +102,7 @@ class AccountTrackerRpcBalanceFetcher implements BalanceFetcher {
104102

105103
const { provider, blockTracker } = this.#getNetworkClient(chainId);
106104
const ethQuery = new EthQuery(provider);
105+
const chainResults: ProcessedBalance[] = [];
107106

108107
// Force fresh block data before multicall
109108
await safelyExecuteWithTimeout(() =>
@@ -133,7 +132,7 @@ class AccountTrackerRpcBalanceFetcher implements BalanceFetcher {
133132

134133
if (nativeBalances) {
135134
accountsToUpdate.forEach((address, index) => {
136-
results.push({
135+
chainResults.push({
137136
success: true,
138137
value: new BN(nativeBalances[index].toString()),
139138
account: address,
@@ -156,15 +155,15 @@ class AccountTrackerRpcBalanceFetcher implements BalanceFetcher {
156155
).catch(() => null);
157156

158157
if (balanceResult) {
159-
results.push({
158+
chainResults.push({
160159
success: true,
161160
value: new BN(balanceResult.replace('0x', ''), 16),
162161
account: address as ChecksumAddress,
163162
token: ZERO_ADDRESS,
164163
chainId,
165164
});
166165
} else {
167-
results.push({
166+
chainResults.push({
168167
success: false,
169168
account: address as ChecksumAddress,
170169
token: ZERO_ADDRESS,
@@ -201,7 +200,7 @@ class AccountTrackerRpcBalanceFetcher implements BalanceFetcher {
201200
if (stakingContractAddress) {
202201
Object.entries(stakedBalanceResult).forEach(
203202
([address, balance]) => {
204-
results.push({
203+
chainResults.push({
205204
success: true,
206205
value: balance
207206
? new BN(balance.replace('0x', ''), 16)
@@ -217,7 +216,22 @@ class AccountTrackerRpcBalanceFetcher implements BalanceFetcher {
217216
}
218217
}
219218
}
220-
}
219+
220+
return chainResults;
221+
});
222+
223+
// Wait for all chains to complete (or fail) and collect results
224+
const chainResultsArray = await Promise.allSettled(chainProcessingPromises);
225+
const results: ProcessedBalance[] = [];
226+
227+
chainResultsArray.forEach((chainResult) => {
228+
if (chainResult.status === 'fulfilled') {
229+
results.push(...chainResult.value);
230+
} else {
231+
// Log error but continue with other chains
232+
console.warn('Chain processing failed:', chainResult.reason);
233+
}
234+
});
221235

222236
return results;
223237
}
@@ -637,21 +651,14 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
637651
}
638652

639653
try {
640-
const balances = await Promise.race([
641-
fetcher.fetch({
642-
chainIds: supportedChains,
643-
queryAllAccounts: isMultiAccountBalancesEnabled,
644-
selectedAccount: toChecksumHexAddress(
645-
selectedAccount.address,
646-
) as ChecksumAddress,
647-
allAccounts,
648-
}),
649-
new Promise<never>((_resolve, reject) =>
650-
setTimeout(() => {
651-
reject(new Error(`Timeout after ${DEFAULT_TIMEOUT_MS}ms`));
652-
}, DEFAULT_TIMEOUT_MS),
653-
),
654-
]);
654+
const balances = await fetcher.fetch({
655+
chainIds: supportedChains,
656+
queryAllAccounts: isMultiAccountBalancesEnabled,
657+
selectedAccount: toChecksumHexAddress(
658+
selectedAccount.address,
659+
) as ChecksumAddress,
660+
allAccounts,
661+
});
655662

656663
if (balances && balances.length > 0) {
657664
aggregated.push(...balances);

‎packages/assets-controllers/src/TokenBalancesController.test.ts‎

Lines changed: 58 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,22 @@ import { advanceTime } from '../../../tests/helpers';
2121
import { createMockInternalAccount } from '../../accounts-controller/src/tests/mocks';
2222
import type { RpcEndpoint } from '../../network-controller/src/NetworkController';
2323

24+
// Mock safelyExecuteWithTimeout
25+
jest.mock('@metamask/controller-utils', () => ({
26+
...jest.requireActual('@metamask/controller-utils'),
27+
safelyExecuteWithTimeout: jest.fn(),
28+
}));
29+
2430
// Constants for native token and staking addresses used in tests
2531
const NATIVE_TOKEN_ADDRESS = '0x0000000000000000000000000000000000000000';
2632
const STAKING_CONTRACT_ADDRESS = '0x4FEF9D741011476750A243aC70b9789a63dd47Df';
2733

34+
// Mock function for safelyExecuteWithTimeout
35+
const { safelyExecuteWithTimeout } = jest.requireMock(
36+
'@metamask/controller-utils',
37+
);
38+
const mockedSafelyExecuteWithTimeout = safelyExecuteWithTimeout as jest.Mock;
39+
2840
const setupController = ({
2941
config,
3042
tokens = { allTokens: {}, allDetectedTokens: {} },
@@ -136,10 +148,22 @@ describe('TokenBalancesController', () => {
136148

137149
beforeEach(() => {
138150
clock = useFakeTimers();
151+
152+
// Mock safelyExecuteWithTimeout to execute the operation normally by default
153+
mockedSafelyExecuteWithTimeout.mockImplementation(
154+
async (operation: () => Promise<unknown>) => {
155+
try {
156+
return await operation();
157+
} catch {
158+
return undefined;
159+
}
160+
},
161+
);
139162
});
140163

141164
afterEach(() => {
142165
clock.restore();
166+
mockedSafelyExecuteWithTimeout.mockRestore();
143167
});
144168

145169
it('should set default state', () => {
@@ -1793,8 +1817,16 @@ describe('TokenBalancesController', () => {
17931817
const tokenAddress = '0x0000000000000000000000000000000000000001';
17941818
const mockError = new Error('Fetcher failed');
17951819

1796-
// Spy on console.warn
1797-
const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
1820+
// Spy on console.error since safelyExecuteWithTimeout logs errors there
1821+
const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation();
1822+
1823+
// Override the mock to use the real safelyExecuteWithTimeout for this test
1824+
const realSafelyExecuteWithTimeout = jest.requireActual(
1825+
'@metamask/controller-utils',
1826+
).safelyExecuteWithTimeout;
1827+
mockedSafelyExecuteWithTimeout.mockImplementation(
1828+
realSafelyExecuteWithTimeout,
1829+
);
17981830

17991831
// Set up tokens so there's something to fetch
18001832
const tokens = {
@@ -1821,14 +1853,13 @@ describe('TokenBalancesController', () => {
18211853

18221854
await controller.updateBalances({ chainIds: [chainId] });
18231855

1824-
// Verify the error was logged with the expected message
1825-
expect(consoleWarnSpy).toHaveBeenCalledWith(
1826-
`Balance fetcher failed for chains ${chainId}: Error: Fetcher failed`,
1827-
);
1856+
// With safelyExecuteWithTimeout, errors are logged as console.error
1857+
// and the operation continues gracefully
1858+
expect(consoleErrorSpy).toHaveBeenCalledWith(mockError);
18281859

18291860
// Restore mocks
18301861
multicallSpy.mockRestore();
1831-
consoleWarnSpy.mockRestore();
1862+
consoleErrorSpy.mockRestore();
18321863
});
18331864

18341865
it('should log error when updateBalances fails after token change', async () => {
@@ -1911,39 +1942,38 @@ describe('TokenBalancesController', () => {
19111942
// Use fake timers for precise control
19121943
jest.useFakeTimers();
19131944

1945+
// Mock safelyExecuteWithTimeout to simulate timeout by returning undefined
1946+
mockedSafelyExecuteWithTimeout.mockImplementation(
1947+
async () => undefined, // Simulates timeout behavior
1948+
);
1949+
1950+
// Mock the multicall function - this won't be reached due to timeout simulation
1951+
const multicallSpy = jest
1952+
.spyOn(multicall, 'getTokenBalancesForMultipleAddresses')
1953+
.mockResolvedValue({
1954+
tokenBalances: {},
1955+
stakedBalances: {},
1956+
});
1957+
19141958
try {
1915-
// Mock the multicall function to return a promise that never resolves
1916-
const multicallSpy = jest
1917-
.spyOn(multicall, 'getTokenBalancesForMultipleAddresses')
1918-
.mockImplementation(() => {
1919-
// Return a promise that never resolves (simulating a hanging request)
1920-
// eslint-disable-next-line no-empty-function
1921-
return new Promise(() => {});
1922-
});
1923-
1924-
// Start the balance update (don't await yet)
1925-
const updatePromise = controller.updateBalances({
1959+
// Start the balance update - should complete gracefully despite timeout
1960+
await controller.updateBalances({
19261961
chainIds: [chainId],
19271962
});
19281963

1929-
// Fast-forward time by 5000ms to trigger the timeout
1930-
jest.advanceTimersByTime(15000);
1964+
// With safelyExecuteWithTimeout, timeouts are handled gracefully
1965+
// The system should continue operating without throwing errors
1966+
// No specific timeout error message should be logged at controller level
19311967

1932-
// Now await the promise - it should have resolved due to timeout
1933-
await updatePromise;
1934-
1935-
// Verify the timeout error was logged with the correct format
1936-
expect(consoleWarnSpy).toHaveBeenCalledWith(
1937-
`Balance fetcher failed for chains ${chainId}: Error: Timeout after 15000ms`,
1938-
);
1968+
// Verify that the update completed without errors
1969+
expect(controller.state.tokenBalances).toBeDefined();
19391970

19401971
// Restore mocks
19411972
multicallSpy.mockRestore();
19421973
consoleWarnSpy.mockRestore();
19431974
} finally {
19441975
// Always restore timers
19451976
jest.useRealTimers();
1946-
consoleWarnSpy.mockRestore();
19471977
}
19481978
});
19491979
});

0 commit comments

Comments
 (0)