diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index b574b460d69..87349864f97 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add a transient, non-persisted per-account assets loading state (`assetsLoadingStatus`, with `assetsLoadingTokens` tracking in-flight `getAssets` fetches) that marks accounts as loading while their assets are fetched, with selectors to read it ([#10230](https://github.com/MetaMask/core/pull/10230)) + ### Changed - Bump `@metamask/transaction-controller` from `^70.0.1` to `^70.1.0` ([#10262](https://github.com/MetaMask/core/pull/10262)) diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 78dda5b7762..ee5479a236d 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -1,5 +1,6 @@ -import { clientControllerSelectors } from '@metamask/client-controller'; /* eslint-disable jest/unbound-method */ +import { deriveStateFromMetadata } from '@metamask/base-controller'; +import { clientControllerSelectors } from '@metamask/client-controller'; import type { TraceCallback, TraceRequest } from '@metamask/controller-utils'; import type { ApiPlatformClient } from '@metamask/core-backend'; import type { InternalAccount } from '@metamask/keyring-internal-api'; @@ -12,6 +13,7 @@ import type { import type { NetworkState } from '@metamask/network-controller'; import { registerKeyringUnlockMock } from './__fixtures__/MockAssetControllerMessenger.js'; +import { waitFor } from './__fixtures__/test-utils.js'; import { AssetsController, getDefaultAssetsControllerState, @@ -89,6 +91,103 @@ function createMockQueryApiClient(): ApiPlatformClient { } as unknown as ApiPlatformClient; } +/** + * Fake accounts API client whose calls can be frozen and released one at a + * time, so tests can hold specific fetches in flight while the rest of the + * pipeline keeps running. `arm()` installs a fresh gate that the next gated + * call waits on; `release()` resolves the oldest pending gate and lets calls + * made after the release pass through. + * + * @returns The gated client plus `arm()`/`release()` controls and a count of + * calls made while a gate was armed. + */ +function createGatedQueryApiClient(): { + client: ApiPlatformClient; + arm: () => void; + release: () => void; + getGatedCallCount: () => number; +} { + let armed = false; + let gatedCallCount = 0; + const pendingGates: (() => void)[] = []; + let currentGate = new Promise(() => undefined); + + const gatedCall = + (value: () => Result): (() => Promise) => + () => { + if (!armed) { + return Promise.resolve(value()); + } + gatedCallCount += 1; + return currentGate.then(value); + }; + + const client = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: gatedCall(() => ({ + fullSupport: ['eip155:1'], + partialSupport: [], + })), + fetchV5MultiAccountBalances: gatedCall(() => ({ + balances: [], + unprocessedNetworks: [], + })), + fetchV6MultiAccountBalances: gatedCall(() => ({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + })), + }, + } as unknown as ApiPlatformClient; + + return { + client, + arm: (): void => { + armed = true; + currentGate = new Promise((resolve) => { + pendingGates.push(resolve); + }); + }, + release: (): void => { + const resolveGate = pendingGates.shift(); + if (resolveGate) { + resolveGate(); + // Calls made after this release pass through, while calls that + // already captured a pending gate stay frozen until it is released. + currentGate = Promise.resolve(); + } else { + armed = false; + } + }, + getGatedCallCount: (): number => gatedCallCount, + }; +} + +/** + * Fake accounts API client whose balance endpoints always fail, to test how + * the loading status settles when a fetch errors out. + * + * @returns The failing client. + */ +function createFailingQueryApiClient(): ApiPlatformClient { + return { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: ['eip155:1'], + partialSupport: [], + }), + fetchV5MultiAccountBalances: jest + .fn() + .mockRejectedValue(new Error('fetch failed')), + fetchV6MultiAccountBalances: jest + .fn() + .mockRejectedValue(new Error('fetch failed')), + }, + } as unknown as ApiPlatformClient; +} + type AllActions = MessengerActions; type AllEvents = MessengerEvents; @@ -334,6 +433,8 @@ describe('AssetsController', () => { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }); }); @@ -363,6 +464,8 @@ describe('AssetsController', () => { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }); }); }); @@ -520,6 +623,8 @@ describe('AssetsController', () => { assetsPrice: {}, customAssets: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }); // Action handlers should be registered @@ -3918,6 +4023,285 @@ describe('AssetsController', () => { }); }); + describe('assets loading status', () => { + it('marks the selected accounts as loading on account switch, then loaded once the fetch settles', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + arm(); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ), + ); + + release(); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ), + ); + }, + ); + }); + + it('runs a queued account switch after the previous refresh finishes, marking its accounts loading then loaded', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger, getSelectedAccountsMock }) => { + await activateTracking(messenger); + + arm(); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); + + // First switch is frozen mid-flight and holds the refresh mutex. + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ), + ); + + const accountB = createMockInternalAccount({ + id: 'mock-account-id-2', + }); + getSelectedAccountsMock.mockReturnValue([accountB]); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/2', + 'entropy:mock-keyring-id-1/1', + ); + + // The queued switch has not fetched yet, so its accounts are not + // marked while the previous refresh still holds the mutex. + expect( + controller.state.assetsLoadingStatus[accountB.id], + ).toBeUndefined(); + + release(); + + // Once the previous refresh finishes, the queued switch runs, marks + // its own accounts, and settles them. + await waitFor(() => + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [MOCK_ACCOUNT_ID]: 'loaded', + [accountB.id]: 'loaded', + }), + ); + + getSelectedAccountsMock.mockClear(); + }, + ); + }); + + it('does not let an older fetch mark an account loaded while a newer overlapping fetch owns the marker', async () => { + const { client, arm, release, getGatedCallCount } = + createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + const account = createMockInternalAccount(); + + // Two overlapping getAssets calls for the same account: the older + // one is frozen mid-flight, then the newer one takes the marker. + arm(); + const olderFetch = controller.getAssets([account], { + forceUpdate: true, + }); + await waitFor(() => expect(getGatedCallCount()).toBeGreaterThan(0)); + + arm(); + const newerFetch = controller.getAssets([account], { + forceUpdate: true, + }); + await waitFor(() => expect(getGatedCallCount()).toBeGreaterThan(1)); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loading', + ), + ); + + release(); + await olderFetch; + + // The older fetch has settled, but the newer one still owns the + // marker, so the account stays loading. + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loading', + ); + + release(); + await newerFetch; + + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loaded', + ); + }, + ); + }); + + it('marks accounts as loading on unlock, then loaded once the fetch settles', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { + queryApiClient: client, + clientControllerState: { isUiOpen: true }, + }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + messenger.publish('KeyringController:lock'); + arm(); + messenger.publish('KeyringController:unlock'); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ), + ); + + release(); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ), + ); + }, + ); + }); + + it('marks accounts as loaded even when the startup fetch fails', async () => { + await withController( + { queryApiClient: createFailingQueryApiClient() }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); + }, + ); + }); + + it('marks the loading status for direct getAssets calls as well', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + const account = createMockInternalAccount(); + + arm(); + const fetchPromise = controller.getAssets([account], { + forceUpdate: true, + }); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loading', + ), + ); + + release(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loaded', + ); + }, + ); + }); + + it('does not mark anything when getAssets is called with no accounts', async () => { + await withController(async ({ controller }) => { + await controller.getAssets([], { forceUpdate: true }); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }); + }); + + it('emits state change events when the loading status is set and settled', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ messenger }) => { + const stateChanges: AssetsControllerState[] = []; + messenger.subscribe('AssetsController:stateChanged', (state) => { + stateChanges.push(state); + }); + + await activateTracking(messenger); + + arm(); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); + + await waitFor(() => + expect( + stateChanges.find( + (state) => + state.assetsLoadingStatus[MOCK_ACCOUNT_ID] === 'loading', + ), + ).toBeDefined(), + ); + + release(); + + await waitFor(() => + expect( + stateChanges.at(-1)?.assetsLoadingStatus[MOCK_ACCOUNT_ID], + ).toBe('loaded'), + ); + }, + ); + }); + + it('does not persist the loading status in the persisted state snapshot', async () => { + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); + + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); + + const persisted = deriveStateFromMetadata( + controller.state, + controller.metadata, + 'persist', + ); + expect(persisted).not.toHaveProperty('assetsLoadingStatus'); + }); + }); + }); + describe('account tree initialized', () => { it('triggers start when the tree initializes after unlock with empty accounts', async () => { const getAccountsMock = jest.fn().mockReturnValue([]); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 5912b609d89..836eec4afab 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -128,6 +128,7 @@ import { buildFastFetchSources, executeAssetsPipeline, } from './pipeline/index.js'; +import { trackAssetsLoading } from './trackAssetsLoading.js'; import type { AccountId, AssetPreferences, @@ -143,6 +144,7 @@ import type { FungibleAssetBalance, AccountWithSupportedChains, AssetType, + AssetsLoadingStatus, DataType, DataRequest, DataResponse, @@ -260,6 +262,8 @@ export type AssetsControllerState = { assetPreferences: { [assetId: string]: AssetPreferences }; /** Currently-active ISO 4217 currency code */ selectedCurrency: SupportedCurrency; + assetsLoadingStatus: Record; + assetsLoadingTokens: Record; }; /** @@ -282,6 +286,8 @@ export function getDefaultAssetsControllerState(): AssetsControllerState { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }; } @@ -503,6 +509,18 @@ const stateMetadata: StateMetadata = { includeInDebugSnapshot: false, usedInUi: true, }, + assetsLoadingStatus: { + persist: false, + includeInStateLogs: true, + includeInDebugSnapshot: true, + usedInUi: true, + }, + assetsLoadingTokens: { + persist: false, + includeInStateLogs: false, + includeInDebugSnapshot: false, + usedInUi: false, + }, }; // ============================================================================ @@ -922,13 +940,14 @@ export class AssetsController extends BaseController< // TEMPORARY: heal assetsInfo metadata wiped by a prior defect // (see extension migration #215 / ASSETS-3346). Remove in a future release. if (tempMigrateAssetsInfoMetadataAssets3346) { - this.update(() => - tempHealAssetsInfoMetadata({ + this.update(() => ({ + ...this.state, + ...tempHealAssetsInfoMetadata({ state: this.state, getMigrationState: tempMigrateAssetsInfoMetadataAssets3346, captureException, }), - ); + })); } this.#initializeNativeAssetsMap(queryApiClient); @@ -1485,6 +1504,7 @@ export class AssetsController extends BaseController< // PUBLIC API: QUERY METHODS // ============================================================================ + @trackAssetsLoading async getAssets( accounts: InternalAccount[], options?: { diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index fca7ca96476..178336facf4 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -93,6 +93,7 @@ export type { DataRequest, DataResponse, AssetsUpdateMode, + AssetsLoadingStatus, // Middleware types Context, NextFunction, @@ -199,6 +200,14 @@ export { getInternalAccountsForGroup, } from './selectors/balance.js'; +export { + getAccountGroupLoadingStatus, + getAccountLoadingStatus, + getIsAssetsLoadingForSelectedAccountGroup, + isAccountGroupLoading, + isAccountLoading, +} from './selectors/loading.js'; + export type { AccountGroupBalance, AccountsById, diff --git a/packages/assets-controller/src/selectors/loading.test.ts b/packages/assets-controller/src/selectors/loading.test.ts new file mode 100644 index 00000000000..2d78dfbe0f4 --- /dev/null +++ b/packages/assets-controller/src/selectors/loading.test.ts @@ -0,0 +1,178 @@ +import type { AccountTreeControllerState } from '@metamask/account-tree-controller'; + +import type { AssetsControllerState } from '../AssetsController.js'; +import type { AccountId } from '../types.js'; +import { + getAccountGroupLoadingStatus, + getAccountLoadingStatus, + getIsAssetsLoadingForSelectedAccountGroup, + isAccountGroupLoading, + isAccountLoading, +} from './loading.js'; + +const ACCOUNT_ID_A = 'mock-account-id-a' as AccountId; +const ACCOUNT_ID_B = 'mock-account-id-b' as AccountId; +const ACCOUNT_ID_C = 'mock-account-id-c' as AccountId; +const GROUP_ID_1 = 'mock-group-id-1'; +const GROUP_ID_2 = 'mock-group-id-2'; + +const createState = ( + assetsLoadingStatus: Record, +): Pick => ({ + assetsLoadingStatus, +}); + +const createAccountTreeState = ( + selectedAccountGroup: string, +): AccountTreeControllerState => + ({ + accountTree: { + wallets: { + 'mock-wallet-id': { + groups: { + [GROUP_ID_1]: { accounts: [ACCOUNT_ID_A, ACCOUNT_ID_B] }, + [GROUP_ID_2]: { accounts: [ACCOUNT_ID_C] }, + }, + }, + }, + }, + selectedAccountGroup, + }) as unknown as AccountTreeControllerState; + +describe('getAccountLoadingStatus', () => { + it('returns the loading status for an account', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', + }); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBe('loading'); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBe('loaded'); + }); + + it('returns undefined for an account with no loading status', () => { + expect( + getAccountLoadingStatus(createState({}), ACCOUNT_ID_C), + ).toBeUndefined(); + }); +}); + +describe('isAccountLoading', () => { + it('returns true while an account is loading', () => { + expect( + isAccountLoading( + createState({ [ACCOUNT_ID_A]: 'loading' }), + ACCOUNT_ID_A, + ), + ).toBe(true); + }); + + it('returns false once the account is loaded or has no status', () => { + expect( + isAccountLoading(createState({ [ACCOUNT_ID_A]: 'loaded' }), ACCOUNT_ID_A), + ).toBe(false); + expect(isAccountLoading(createState({}), ACCOUNT_ID_A)).toBe(false); + }); +}); + +describe('getAccountGroupLoadingStatus', () => { + it('returns statuses for the accounts in the group', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', + }); + expect( + getAccountGroupLoadingStatus( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toStrictEqual({ + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', + }); + }); + + it('omits accounts in the group that have no loading status', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + expect( + getAccountGroupLoadingStatus( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toStrictEqual({ [ACCOUNT_ID_A]: 'loading' }); + }); + + it('returns an empty object for an unknown group', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + expect( + getAccountGroupLoadingStatus(state, createAccountTreeState(''), 'nope'), + ).toStrictEqual({}); + }); +}); + +describe('isAccountGroupLoading', () => { + it('returns true while any account in the group is loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loaded', + [ACCOUNT_ID_B]: 'loading', + }); + expect( + isAccountGroupLoading( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toBe(true); + }); + + it('returns false when no account in the group is loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loaded', + [ACCOUNT_ID_B]: 'loaded', + [ACCOUNT_ID_C]: 'loading', + }); + expect( + isAccountGroupLoading( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toBe(false); + }); +}); + +describe('getIsAssetsLoadingForSelectedAccountGroup', () => { + it('returns true while any account in the selected group is loading', () => { + const state = createState({ [ACCOUNT_ID_C]: 'loading' }); + expect( + getIsAssetsLoadingForSelectedAccountGroup( + state, + createAccountTreeState(GROUP_ID_2), + ), + ).toBe(true); + }); + + it('returns false when the selected group is not loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loaded', + [ACCOUNT_ID_B]: 'loading', + }); + expect( + getIsAssetsLoadingForSelectedAccountGroup( + state, + createAccountTreeState(GROUP_ID_2), + ), + ).toBe(false); + }); + + it('returns false when no group is selected', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + expect( + getIsAssetsLoadingForSelectedAccountGroup( + state, + createAccountTreeState(''), + ), + ).toBe(false); + }); +}); diff --git a/packages/assets-controller/src/selectors/loading.ts b/packages/assets-controller/src/selectors/loading.ts new file mode 100644 index 00000000000..3c1cfd19872 --- /dev/null +++ b/packages/assets-controller/src/selectors/loading.ts @@ -0,0 +1,96 @@ +import type { AccountTreeControllerState } from '@metamask/account-tree-controller'; + +import type { AssetsControllerState } from '../AssetsController.js'; +import type { AccountId, AssetsLoadingStatus } from '../types.js'; +import { getAccountIdsForGroup } from './balance.js'; + +/** + * Get the loading status for a single account. + * + * @param state - AssetsController state slice. + * @param accountId - The account id (`InternalAccount.id`). + * @returns `'loading'` while the account's assets are loading, `'loaded'` + * after its fetch has settled, or `undefined` if no fetch was triggered. + */ +export function getAccountLoadingStatus( + state: Pick, + accountId: AccountId, +): AssetsLoadingStatus | undefined { + return state.assetsLoadingStatus[accountId]; +} + +/** + * Check whether an account's assets are currently loading. + * + * @param state - AssetsController state slice. + * @param accountId - The account id (`InternalAccount.id`). + * @returns True while the account's assets are loading. + */ +export function isAccountLoading( + state: Pick, + accountId: AccountId, +): boolean { + return getAccountLoadingStatus(state, accountId) === 'loading'; +} + +/** + * Get the loading statuses for every account in an account group. + * + * @param state - AssetsController state slice. + * @param accountTreeState - AccountTreeController state slice. + * @param groupId - The account group id. + * @returns A record containing an entry for each account in the group that + * has a loading status, keyed by account id. + */ +export function getAccountGroupLoadingStatus( + state: Pick, + accountTreeState: AccountTreeControllerState, + groupId: string, +): Record { + const loadingStatus = state.assetsLoadingStatus; + const result: Record = {}; + for (const accountId of getAccountIdsForGroup(accountTreeState, groupId)) { + const status = loadingStatus[accountId]; + if (status !== undefined) { + result[accountId] = status; + } + } + return result; +} + +/** + * Check whether any account in an account group is currently loading. + * + * @param state - AssetsController state slice. + * @param accountTreeState - AccountTreeController state slice. + * @param groupId - The account group id. + * @returns True while at least one account in the group is loading. + */ +export function isAccountGroupLoading( + state: Pick, + accountTreeState: AccountTreeControllerState, + groupId: string, +): boolean { + const loadingStatus = state.assetsLoadingStatus; + return getAccountIdsForGroup(accountTreeState, groupId).some( + (accountId) => loadingStatus[accountId] === 'loading', + ); +} + +/** + * Check whether the selected account group is currently loading its assets. + * + * @param state - AssetsController state slice. + * @param accountTreeState - AccountTreeController state slice. + * @returns True while any account in the selected account group is loading. + */ +export function getIsAssetsLoadingForSelectedAccountGroup( + state: Pick, + accountTreeState: AccountTreeControllerState, +): boolean { + const groupId = accountTreeState.selectedAccountGroup; + if (!groupId) { + return false; + } + return isAccountGroupLoading(state, accountTreeState, groupId); +} diff --git a/packages/assets-controller/src/trackAssetsLoading.test.ts b/packages/assets-controller/src/trackAssetsLoading.test.ts new file mode 100644 index 00000000000..9fd5d945f97 --- /dev/null +++ b/packages/assets-controller/src/trackAssetsLoading.test.ts @@ -0,0 +1,153 @@ +import type { InternalAccount } from '@metamask/keyring-internal-api'; + +import type { AssetsControllerState } from './AssetsController.js'; +import { trackAssetsLoading } from './trackAssetsLoading.js'; + +type FakeControllerState = Pick< + AssetsControllerState, + 'assetsLoadingStatus' | 'assetsLoadingTokens' +>; + +const ACCOUNT_1: InternalAccount = { + id: 'account-1', + address: '0x1234567890123456789012345678901234567890', + options: {}, + methods: [], + type: 'eip155:eoa', + scopes: ['eip155:1'], + metadata: { + name: 'Test Account 1', + keyring: { type: 'HD Key Tree' }, + importTime: 0, + lastSelected: 0, + }, +}; + +const ACCOUNT_2: InternalAccount = { + ...ACCOUNT_1, + id: 'account-2', + metadata: { ...ACCOUNT_1.metadata, name: 'Test Account 2' }, +}; + +/** + * Minimal stand-in for the controller: just the state and `update` the + * decorator needs, with a fetch method that stays in flight until released. + */ +class FakeAssetsController { + readonly holds: (() => void)[] = []; + + state: FakeControllerState = { + assetsLoadingStatus: {}, + assetsLoadingTokens: {}, + }; + + update(callback: (state: FakeControllerState) => void): void { + callback(this.state); + } + + /** + * Let the oldest fetch that is still in flight finish. + */ + releaseNextFetch(): void { + const resolve = this.holds.shift(); + resolve?.(); + } + + @trackAssetsLoading + async getAssets( + _accounts: InternalAccount[], + options?: { forceUpdate?: boolean; fail?: boolean }, + ): Promise { + if (options?.fail) { + throw new Error('fetch failed'); + } + await new Promise((resolve) => { + this.holds.push(resolve); + }); + } +} + +describe('trackAssetsLoading', () => { + it('marks accounts as loading while the fetch is in flight and loaded once it settles', async () => { + const controller = new FakeAssetsController(); + const fetchPromise = controller.getAssets([ACCOUNT_1, ACCOUNT_2], { + forceUpdate: true, + }); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [ACCOUNT_1.id]: 'loading', + [ACCOUNT_2.id]: 'loading', + }); + expect(controller.state.assetsLoadingTokens[ACCOUNT_1.id]).toBeDefined(); + + controller.releaseNextFetch(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [ACCOUNT_1.id]: 'loaded', + [ACCOUNT_2.id]: 'loaded', + }); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + }); + + it('settles the loading status when the fetch fails', async () => { + const controller = new FakeAssetsController(); + + await expect( + controller.getAssets([ACCOUNT_1], { forceUpdate: true, fail: true }), + ).rejects.toThrow('fetch failed'); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [ACCOUNT_1.id]: 'loaded', + }); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + }); + + it('does not let an older fetch settle an account a newer overlapping fetch owns', async () => { + const controller = new FakeAssetsController(); + const olderFetch = controller.getAssets([ACCOUNT_1], { + forceUpdate: true, + }); + const newerFetch = controller.getAssets([ACCOUNT_1], { + forceUpdate: true, + }); + + expect(controller.state.assetsLoadingStatus[ACCOUNT_1.id]).toBe('loading'); + + controller.releaseNextFetch(); + await olderFetch; + + expect(controller.state.assetsLoadingStatus[ACCOUNT_1.id]).toBe('loading'); + + controller.releaseNextFetch(); + await newerFetch; + + expect(controller.state.assetsLoadingStatus[ACCOUNT_1.id]).toBe('loaded'); + }); + + it('does not mark anything when there are no accounts', async () => { + const controller = new FakeAssetsController(); + const fetchPromise = controller.getAssets([], { forceUpdate: true }); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + + controller.releaseNextFetch(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }); + + it('does not track calls that do not force an update', async () => { + const controller = new FakeAssetsController(); + const fetchPromise = controller.getAssets([ACCOUNT_1]); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + + controller.releaseNextFetch(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }); +}); diff --git a/packages/assets-controller/src/trackAssetsLoading.ts b/packages/assets-controller/src/trackAssetsLoading.ts new file mode 100644 index 00000000000..33d721c3722 --- /dev/null +++ b/packages/assets-controller/src/trackAssetsLoading.ts @@ -0,0 +1,113 @@ +import type { InternalAccount } from '@metamask/keyring-internal-api'; + +import type { AssetsController } from './AssetsController.js'; +import type { AccountId, AssetsLoadingStatus } from './types.js'; + +/** + * The controller surface the loading tracker needs. `update` is spelled out + * here because it is protected on the controller and so cannot be picked. + */ +type AssetsLoadingStateUpdater = Pick & { + update: ( + callback: (state: { + assetsLoadingStatus: Record; + assetsLoadingTokens: Record; + }) => void, + ) => void; +}; + +let loadingTokenCounter = 0; + +/** + * Mark the given accounts as `'loading'` and record an ownership token that + * identifies this invocation as the owner of those markers. + * + * @param controller - The controller whose state should be updated. + * @param accounts - The accounts being fetched. + * @returns The ownership token, or `0` when there is nothing to mark. + */ +function markAssetsLoading( + controller: AssetsLoadingStateUpdater, + accounts: InternalAccount[], +): number { + if (accounts.length === 0) { + return 0; + } + loadingTokenCounter += 1; + const token = loadingTokenCounter; + controller.update((state) => { + for (const account of accounts) { + state.assetsLoadingTokens[account.id] = token; + state.assetsLoadingStatus[account.id] = 'loading'; + } + }); + return token; +} + +/** + * Settle the markers owned by the given token to `'loaded'`. Accounts whose + * marker was re-claimed by a newer invocation are left untouched, so an older + * overlapping fetch cannot clobber a newer one's loading status. + * + * @param controller - The controller whose state should be updated. + * @param accounts - The accounts the invocation fetched. + * @param token - The ownership token recorded by {@link markAssetsLoading}. + */ +function markAssetsSettled( + controller: AssetsLoadingStateUpdater, + accounts: InternalAccount[], + token: number, +): void { + if (token === 0) { + return; + } + const ownedAccounts = accounts.filter( + (account) => controller.state.assetsLoadingTokens[account.id] === token, + ); + if (ownedAccounts.length === 0) { + return; + } + controller.update((state) => { + for (const account of ownedAccounts) { + state.assetsLoadingStatus[account.id] = 'loaded'; + delete state.assetsLoadingTokens[account.id]; + } + }); +} + +/** + * Method decorator that tracks a per-account assets loading status around the + * decorated fetch method: the requested accounts are marked `'loading'` when + * the method is invoked and settle to `'loaded'` once it completes, whether it + * succeeded or failed. Calls that do not force an update are cache reads and + * are not tracked. + * + * @param target - The decorated fetch method. + * @param _context - The decorator context. + * @returns The wrapped fetch method. + */ +export function trackAssetsLoading< + This, + Args extends [InternalAccount[], { forceUpdate?: boolean }?, ...unknown[]], + Return, +>( + target: (this: This, ...args: Args) => Promise, + _context: ClassMethodDecoratorContext< + This, + (this: This, ...args: Args) => Promise + >, +): (this: This, ...args: Args) => Promise { + return async function (this: This, ...args: Args): Promise { + const [accounts, options] = args; + if (options?.forceUpdate !== true) { + return target.call(this, ...args); + } + const controller = this as unknown as AssetsLoadingStateUpdater; + const token = markAssetsLoading(controller, accounts); + try { + return await target.call(this, ...args); + } finally { + markAssetsSettled(controller, accounts, token); + } + }; +} diff --git a/packages/assets-controller/src/types.ts b/packages/assets-controller/src/types.ts index a1c7bd80d28..912bbf68dbc 100644 --- a/packages/assets-controller/src/types.ts +++ b/packages/assets-controller/src/types.ts @@ -414,6 +414,8 @@ export type DataResponse = { */ export type AssetsUpdateMode = 'full' | 'merge' | 'update'; +export type AssetsLoadingStatus = 'loading' | 'loaded'; + // ============================================================================ // DATA SOURCE <-> CONTROLLER (DIRECT CALLS, NO MESSENGER PER SOURCE) // ============================================================================