From f3a2f18b1ffe815e559a537f6bf868c334973b57 Mon Sep 17 00:00:00 2001 From: NPC Date: Tue, 18 Aug 2026 01:16:31 +0700 Subject: [PATCH] feat: dedupe and epoch-fence address/balance fetches against wallet changes - Add RequestCoordinator + a single global worklet epoch: dedupes concurrent fetches by cache key, and discards a result if the worklet's loaded wallet changed (switch/lock/delete) while the fetch was in flight - Wire the epoch bump into WorkletLifecycleService.reset() - the one choke point every wallet-changing flow already routes through, so no other lifecycle code needs to change - Fix AddressService: the loading flag could get stuck at `true` forever when a stale result was discarded, permanently blocking useAddressLoader from ever retrying that network/account - Fix BalanceService: updateBalance/updateLastBalanceUpdate now receive the wallet captured at fetch kickoff instead of re-resolving activeWalletId at write time, which could attribute a balance to the wrong wallet after a switch - Unify fetchBalance onto the same per-network fetch/commit path as fetchBalances so they dedupe against each other instead of risking a cache-key collision --- src/services/addressService.ts | 75 ++++--- src/services/balanceService.ts | 160 +++++++++------ src/services/workletLifecycleService.ts | 12 ++ src/utils/requestCoordinator.ts | 64 ++++++ src/utils/workletEpoch.ts | 31 +++ tests/hooks/useBalance.test.tsx | 4 +- tests/services/addressService.test.ts | 133 ++++++++++++ tests/services/balanceService.test.ts | 191 +++++++++++++++++- .../services/workletLifecycleService.test.ts | 67 ++++++ tests/utils/requestCoordinator.test.ts | 190 +++++++++++++++++ tests/utils/workletEpoch.test.ts | 61 ++++++ 11 files changed, 897 insertions(+), 91 deletions(-) create mode 100644 src/utils/requestCoordinator.ts create mode 100644 src/utils/workletEpoch.ts create mode 100644 tests/utils/requestCoordinator.test.ts create mode 100644 tests/utils/workletEpoch.test.ts diff --git a/src/services/addressService.ts b/src/services/addressService.ts index b45503d..38923e3 100644 --- a/src/services/addressService.ts +++ b/src/services/addressService.ts @@ -23,6 +23,7 @@ import { produce } from 'immer' import { getWalletStore } from '../store/walletStore' import { getWorkletStore } from '../store/workletStore' import { handleServiceError } from '../utils/errorHandling' +import { RequestCoordinator } from '../utils/requestCoordinator' import { requireInitialized, resolveWalletId, @@ -38,6 +39,12 @@ import { AddressInfoResult } from '../types' * Provides methods for retrieving and caching wallet addresses. */ export class AddressService { + /** + * Dedupes concurrent getAddress calls and fences stale writes if the + * worklet's loaded wallet changes (switch/lock/delete) mid-flight + */ + private static coordinator = new RequestCoordinator() + /** * Get address for a specific network and account index * Caches the address in walletStore for future use @@ -64,17 +71,36 @@ export class AddressService { return cachedAddress } - const hrpc = await requireInitialized() + const cacheKey = `${targetWalletId}:${network}:${accountIndex}` + return this.coordinator.run( + cacheKey, + () => this.fetchAddressFromWorklet(network, accountIndex, targetWalletId), + (address) => this.cacheAddress(targetWalletId, network, accountIndex, address), + ) + } + + private static async fetchAddressFromWorklet( + network: string, + accountIndex: number, + targetWalletId: string, + ): Promise { + const walletStore = getWalletStore() const loadingKey = `${network}-${accountIndex}` - try { + const setLoading = (loading: boolean): void => { walletStore.setState((prev) => produce(prev, (state) => { state.walletLoading[targetWalletId] ??= {} - state.walletLoading[targetWalletId][loadingKey] = true + state.walletLoading[targetWalletId][loadingKey] = loading }), ) + } + + try { + const hrpc = await requireInitialized() + + setLoading(true) const response = await hrpc.callMethod({ methodName: 'getAddress', @@ -86,13 +112,12 @@ export class AddressService { throw new Error('Failed to get address from worklet') } - let address: string try { const parsed = JSON.parse(response.result) if (typeof parsed !== 'string') { throw new Error('Address must be a string') } - address = parsed + return parsed } catch (error) { throw new Error( `Failed to parse address from worklet response: ${ @@ -100,41 +125,29 @@ export class AddressService { }`, ) } - - // Cache the address using helper (per-wallet) - walletStore.setState((prev) => - produce( - updateAddressInState( - prev, - targetWalletId, - network, - accountIndex, - address, - ), - (state) => { - state.walletLoading[targetWalletId] ??= {} - state.walletLoading[targetWalletId][loadingKey] = false - }, - ), - ) - - return address } catch (error) { - walletStore.setState((prev) => - produce(prev, (state) => { - state.walletLoading[targetWalletId] ??= {} - state.walletLoading[targetWalletId][loadingKey] = false - }), - ) - handleServiceError(error, 'AddressService', 'getAddress', { network, accountIndex, walletId: targetWalletId, }) + } finally { + setLoading(false) } } + private static cacheAddress( + targetWalletId: string, + network: string, + accountIndex: number, + address: string, + ): void { + const walletStore = getWalletStore() + walletStore.setState((prev) => + updateAddressInState(prev, targetWalletId, network, accountIndex, address), + ) + } + /** * Get addresses for multiple accounts and networks. */ diff --git a/src/services/balanceService.ts b/src/services/balanceService.ts index 6b07d65..5f3e6a1 100644 --- a/src/services/balanceService.ts +++ b/src/services/balanceService.ts @@ -48,12 +48,20 @@ import { produce } from 'immer' import { getWalletStore } from '../store/walletStore' import { resolveWalletId, updateBalanceInState, getNestedState } from '../utils/storeHelpers' +import { RequestCoordinator } from '../utils/requestCoordinator' import { validateBalance, validateWalletParams } from '../utils/validation' import { BalanceFetchResult, IAsset } from '../types' import { AccountService } from './accountService' import { convertBalanceToString } from '../utils/balanceUtils' import { log, logError, logWarn } from '../utils/logger' +/** + * Dedupes concurrent fetchBalance/fetchBalances calls and fences stale + * writes if the worklet's loaded wallet changes (switch/lock/delete) + * mid-flight. + */ +const balanceCoordinator = new RequestCoordinator() + /** * Balance Service * @@ -293,6 +301,7 @@ type FetchBalancesResult = | { success: false; asset: IAsset; error: unknown }; const BATCH_NOT_SUPPORTED = 'Batch balance fetching is not supported'; +const DEFAULT_NETWORK_TIMEOUT_MS = 15_000; async function fetchNonNativeBalancesInBatch(network: string, accountIndex: number, nonNativeAssets: IAsset[]): Promise { try { @@ -414,8 +423,10 @@ async function fetchBalancesForNetwork(accountIndex: number, networkAssets: IAss export async function fetchBalances( accountIndex: number, assets: IAsset[], - networkTimeoutMs: number = 15_000, + networkTimeoutMs: number = DEFAULT_NETWORK_TIMEOUT_MS, ): Promise { + const targetWalletId = resolveWalletId() + const assetsByNetwork = new Map(); assets.forEach(asset => { const network = asset.getNetwork(); @@ -426,47 +437,8 @@ export async function fetchBalances( }); const allNetworkPromises = Array.from(assetsByNetwork.entries()).map( - async ([network, networkAssets]) => { - let timedOut = false; - let timeoutId: ReturnType - const timeout = new Promise(resolve => { - timeoutId = setTimeout(() => { - timedOut = true; - const error = new Error(`Network ${network} timed out after ${networkTimeoutMs}ms`); - logWarn(`[fetchBalances] ${error.message}`); - resolve(networkAssets.map(asset => ({ success: false, asset, error }))); - }, networkTimeoutMs); - }); - - const networkResults = await Promise.race([ - fetchBalancesForNetwork(accountIndex, networkAssets), - timeout - ]); - clearTimeout(timeoutId!); - - if (!timedOut) { - let hasSuccessfulUpdate = false; - for (const result of networkResults) { - if (!result.success) { - continue; - } - hasSuccessfulUpdate = true; - BalanceService.updateBalance( - accountIndex, - network, - result.asset.getId(), - result.balance, - ); - } - - if (hasSuccessfulUpdate) { - BalanceService.updateLastBalanceUpdate(network, accountIndex); - log(`[fetchBalances] Fetched balances for network ${network}:${accountIndex}`); - } - } - - return networkResults; - }, + ([network, networkAssets]) => + fetchAndCacheNetworkBalances(accountIndex, network, networkAssets, networkTimeoutMs, targetWalletId), ); const allResults = (await Promise.all(allNetworkPromises)).flat(); @@ -501,39 +473,113 @@ export async function fetchBalances( }); } +/** Fetches and caches balances for one network, deduped by cache key. */ +function fetchAndCacheNetworkBalances( + accountIndex: number, + network: string, + networkAssets: IAsset[], + networkTimeoutMs: number, + targetWalletId: string, +): Promise { + const assetIds = networkAssets.map((asset) => asset.getId()).sort().join(',') + const cacheKey = `${targetWalletId}:${network}:${accountIndex}:${assetIds}` + + return balanceCoordinator.run( + cacheKey, + () => fetchNetworkBalancesWithTimeout(accountIndex, network, networkAssets, networkTimeoutMs), + (networkResults) => commitNetworkBalances(accountIndex, network, targetWalletId, networkResults), + ) +} + +async function fetchNetworkBalancesWithTimeout( + accountIndex: number, + network: string, + networkAssets: IAsset[], + networkTimeoutMs: number, +): Promise { + let timeoutId: ReturnType + const timeout = new Promise(resolve => { + timeoutId = setTimeout(() => { + const error = new Error(`Network ${network} timed out after ${networkTimeoutMs}ms`); + logWarn(`[fetchBalances] ${error.message}`); + resolve(networkAssets.map(asset => ({ success: false, asset, error }))); + }, networkTimeoutMs); + }); + + const results = await Promise.race([ + fetchBalancesForNetwork(accountIndex, networkAssets), + timeout, + ]); + clearTimeout(timeoutId!); + return results; +} + +function commitNetworkBalances( + accountIndex: number, + network: string, + targetWalletId: string, + networkResults: FetchBalancesResult[], +): void { + let hasSuccessfulUpdate = false; + for (const result of networkResults) { + if (!result.success) { + continue; + } + hasSuccessfulUpdate = true; + BalanceService.updateBalance(accountIndex, network, result.asset.getId(), result.balance, targetWalletId); + } + + if (hasSuccessfulUpdate) { + BalanceService.updateLastBalanceUpdate(network, accountIndex, targetWalletId); + log(`[fetchBalances] Fetched balances for network ${network}:${accountIndex}`); + } +} + /** - * Fetch balance for a specific asset + * Fetch balance for a specific asset. + * + * Routes through the same per-network dedupe/commit path as fetchBalances, + * treating this as a batch of one - so a concurrent fetchBalances call for + * the same wallet/network/account/asset shares this fetch instead of firing + * a duplicate RPC. * * @param accountIndex - Account index * @param asset - Asset entity (contains ID and contract details) - * @param walletId - Optional wallet identifier (defaults to activeWalletId) * @returns Promise with balance fetch result */ -export async function fetchBalance(accountIndex: number, asset: IAsset): Promise { - const results = await fetchBalancesForNetwork(accountIndex, [asset]); - const result = results[0]!; +export async function fetchBalance( + accountIndex: number, + asset: IAsset, +): Promise { + const targetWalletId = resolveWalletId() + const network = asset.getNetwork() - if (result.success) { - BalanceService.updateBalance(accountIndex, result.asset.getNetwork(), result.asset.getId(), result.balance); - BalanceService.updateLastBalanceUpdate(result.asset.getNetwork(), accountIndex); + const [result] = await fetchAndCacheNetworkBalances( + accountIndex, + network, + [asset], + DEFAULT_NETWORK_TIMEOUT_MS, + targetWalletId, + ) + if (result!.success) { return { success: true, - network: result.asset.getNetwork(), + network, accountIndex, - assetId: result.asset.getId(), - balance: result.balance, + assetId: result!.asset.getId(), + balance: result!.balance, }; } - const errorMessage = result.error instanceof Error ? result.error.message : String(result.error); - logError(`Failed to fetch balance for ${result.asset.getNetwork()}:${accountIndex}:${result.asset.getId()}:`, result.error); + const errorMessage = result!.error instanceof Error ? result!.error.message : String(result!.error); + logError(`Failed to fetch balance for ${network}:${accountIndex}:${result!.asset.getId()}:`, result!.error); return { success: false, - network: result.asset.getNetwork(), + network, accountIndex, - assetId: result.asset.getId(), + assetId: result!.asset.getId(), balance: null, error: errorMessage, }; diff --git a/src/services/workletLifecycleService.ts b/src/services/workletLifecycleService.ts index 2a2b1a1..5459fe8 100644 --- a/src/services/workletLifecycleService.ts +++ b/src/services/workletLifecycleService.ts @@ -31,6 +31,7 @@ import type { WdkConfigs, BundleConfig } from '../types' import type { WorkletState } from '../store/workletStore' import HRPC from '@tetherto/pear-wrk-wdk/hrpc' import { createResolvablePromise } from '../utils/promise' +import { bumpEpoch } from '../utils/workletEpoch' /** * Worklet Lifecycle Service @@ -183,6 +184,11 @@ export class WorkletLifecycleService { throw new Error('Worklet must be started before initializing WDK') } + // Also bump here, not just in reset(): this is the moment the worklet + // actually gets a new seed, so a fetch started after reset() but before + // this point needs its own fencing too. + bumpEpoch() + try { store.setState({ error: null, isLoading: true }) @@ -532,6 +538,12 @@ export class WorkletLifecycleService { workletStore.setState({ isWorkletInitializedPromise: createResolvablePromise() }) + // Bump the worklet epoch so any address/balance fetch still in flight for + // the previously-loaded wallet skips writing once it resolves. reset() is + // the sole choke point every flow that changes the worklet's loaded + // wallet (delete, lock, switch) routes through - see workletEpoch.ts. + bumpEpoch() + // Clear addresses from wallet store walletStore.setState({ addresses: {}, diff --git a/src/utils/requestCoordinator.ts b/src/utils/requestCoordinator.ts new file mode 100644 index 0000000..7bb734e --- /dev/null +++ b/src/utils/requestCoordinator.ts @@ -0,0 +1,64 @@ +// Copyright 2026 Tether Operations Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * Coordinates concurrent fetches that get cached into the wallet store: + * dedupes concurrent calls for the same key, and discards a result if the + * worklet's loaded wallet changed while it was in flight, instead of + * committing stale data. Each domain (addresses, balances, ...) should use + * its own instance. + */ +import { getEpoch } from './workletEpoch' +import { log } from './logger' + +export class RequestCoordinator { + private inflight = new Map>() + + /** + * Runs `fetch` for `key`, sharing the in-flight promise with any + * concurrent caller using the same key. If the worklet's loaded wallet is + * still the same one as when `fetch` started, `commit` is called with the + * result; otherwise the result is discarded. + */ + async run( + key: string, + fetch: () => Promise, + commit: (result: T) => void, + ): Promise { + const epoch = getEpoch() + const fullKey = `${epoch}:${key}` + + const existing = this.inflight.get(fullKey) + if (existing) { + return existing as Promise + } + + const promise = fetch() + this.inflight.set(fullKey, promise) + + try { + const result = await promise + if (epoch === getEpoch()) { + commit(result) + } else { + log('[RequestCoordinator] Discarding stale result', { key }) + } + return result + } finally { + if (this.inflight.get(fullKey) === promise) { + this.inflight.delete(fullKey) + } + } + } +} diff --git a/src/utils/workletEpoch.ts b/src/utils/workletEpoch.ts new file mode 100644 index 0000000..88178d7 --- /dev/null +++ b/src/utils/workletEpoch.ts @@ -0,0 +1,31 @@ +// Copyright 2026 Tether Operations Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * Global counter tracking the worklet's loaded wallet. Bumped by + * `WorkletLifecycleService.reset()` whenever it changes (delete, lock, + * switch). Consumers snapshot `getEpoch()` at fetch kickoff and compare it + * again after the fetch resolves, to detect whether the result is stale. + */ +let epoch = 0 + +/** Current worklet epoch. */ +export function getEpoch(): number { + return epoch +} + +/** Bump the worklet epoch. Call only from WorkletLifecycleService.reset(). */ +export function bumpEpoch(): void { + epoch += 1 +} diff --git a/tests/hooks/useBalance.test.tsx b/tests/hooks/useBalance.test.tsx index cf68002..1a1e984 100644 --- a/tests/hooks/useBalance.test.tsx +++ b/tests/hooks/useBalance.test.tsx @@ -292,8 +292,8 @@ describe('useBalance', () => { }); expect(mockAccountService.callAccountMethod).toHaveBeenCalledWith(mockAssetNative.getNetwork(), mockAccountIndex, 'getBalance'); expect(mockConvertBalanceToString).toHaveBeenCalledWith(mockBalanceValue); - expect(mockBalanceService.updateBalance).toHaveBeenCalledWith(mockAccountIndex, mockAssetNative.getNetwork(), mockAssetNative.getId(), mockBalanceValue); - expect(mockBalanceService.updateLastBalanceUpdate).toHaveBeenCalledWith('ethereum', mockAccountIndex); + expect(mockBalanceService.updateBalance).toHaveBeenCalledWith(mockAccountIndex, mockAssetNative.getNetwork(), mockAssetNative.getId(), mockBalanceValue, 'mock-active-wallet'); + expect(mockBalanceService.updateLastBalanceUpdate).toHaveBeenCalledWith('ethereum', mockAccountIndex, 'mock-active-wallet'); }); it('should handle fetch errors correctly', async () => { diff --git a/tests/services/addressService.test.ts b/tests/services/addressService.test.ts index 5b792e0..10ea43f 100644 --- a/tests/services/addressService.test.ts +++ b/tests/services/addressService.test.ts @@ -21,6 +21,7 @@ import { AddressService } from '../../src/services/addressService' import { getWorkletStore } from '../../src/store/workletStore' import { getWalletStore } from '../../src/store/walletStore' +import { bumpEpoch } from '../../src/utils/workletEpoch' // Mock stores jest.mock('../../src/store/workletStore', () => ({ @@ -323,6 +324,138 @@ describe('AddressService', () => { expect(loadingWasSetToTrue).toBe(true) expect(loadingWasSetToFalse).toBe(true) }) + + it('should dedupe concurrent calls for the same key', async () => { + const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0' + mockHRPC.callMethod.mockResolvedValue({ + result: JSON.stringify(address), + }) + + const p1 = AddressService.getAddress('ethereum', 0) + const p2 = AddressService.getAddress('ethereum', 0) + const results = await Promise.all([p1, p2]) + + expect(mockHRPC.callMethod).toHaveBeenCalledTimes(1) + expect(results).toEqual([address, address]) + }) + + it('should not dedupe different keys', async () => { + mockHRPC.callMethod.mockResolvedValue({ + result: JSON.stringify('0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'), + }) + + await Promise.all([ + AddressService.getAddress('ethereum', 0), + AddressService.getAddress('ethereum', 1), + AddressService.getAddress('polygon', 0), + ]) + + expect(mockHRPC.callMethod).toHaveBeenCalledTimes(3) + }) + + it('should not dedupe calls for different wallets', async () => { + mockHRPC.callMethod.mockResolvedValue({ + result: JSON.stringify('0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0'), + }) + + await Promise.all([ + AddressService.getAddress('ethereum', 0, 'wallet-a'), + AddressService.getAddress('ethereum', 0, 'wallet-b'), + ]) + + expect(mockHRPC.callMethod).toHaveBeenCalledTimes(2) + }) + + // Replays the sequence of setState updater functions against an initial + // state, mirroring how the real zustand store would apply them. + function replaySetStateCalls(setStateMock: jest.Mock) { + let state: any = { + walletLoading: {}, + addresses: {}, + balanceLoading: {}, + lastBalanceUpdate: {}, + balances: {}, + activeWalletId: 'test-wallet-1', + } + for (const [updater] of setStateMock.mock.calls) { + state = typeof updater === 'function' ? { ...state, ...updater(state) } : { ...state, ...updater } + } + return state + } + + // requireInitialized() awaits a couple of already-resolved promises + // before hrpc.callMethod is actually invoked - flush enough microtask + // ticks that it's guaranteed to have been called by the time we return. + async function flushMicrotasks(times = 10) { + for (let i = 0; i < times; i++) { + await Promise.resolve() + } + } + + it('should skip the address-cache write after a worklet epoch bump, but still clear the loading flag', async () => { + const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0' + let resolveCallMethod: (value: { result: string }) => void + mockHRPC.callMethod.mockImplementation( + () => new Promise((resolve) => { resolveCallMethod = resolve }), + ) + + const pending = AddressService.getAddress('ethereum', 0, 'test-wallet-1') + await flushMicrotasks() + expect(mockHRPC.callMethod).toHaveBeenCalled() + + // Simulate a wallet switch/lock happening while the fetch is in flight. + bumpEpoch() + resolveCallMethod!({ result: JSON.stringify(address) }) + + const result = await pending + expect(result).toBe(address) + + // The stale fetch must not resurrect address data after the epoch changed... + const finalState = replaySetStateCalls(mockWalletStore.setState) + expect(finalState.addresses['test-wallet-1']).toBeUndefined() + + // ...but it must still clear the loading flag it set, so a future + // getAddress call for this network/account isn't blocked forever. + expect(finalState.walletLoading['test-wallet-1']?.['ethereum-0']).toBe(false) + }) + + it('should clear the loading flag even when the fetch fails after an epoch bump', async () => { + let rejectCallMethod: (reason: Error) => void + mockHRPC.callMethod.mockImplementation( + () => new Promise((_resolve, reject) => { rejectCallMethod = reject }), + ) + + const pending = AddressService.getAddress('ethereum', 0, 'test-wallet-1') + await flushMicrotasks() + expect(mockHRPC.callMethod).toHaveBeenCalled() + + bumpEpoch() + rejectCallMethod!(new Error('Worklet error')) + + await expect(pending).rejects.toThrow('Worklet error') + + const finalState = replaySetStateCalls(mockWalletStore.setState) + expect(finalState.walletLoading['test-wallet-1']?.['ethereum-0']).toBe(false) + }) + + it('should start a fresh fetch after an epoch bump, instead of joining the stale in-flight promise', async () => { + const staleAddress = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0' + const freshAddress = '0x842d35Cc6634C0532925a3b844Bc9e7595f0bEb0' + + mockHRPC.callMethod + .mockResolvedValueOnce({ result: JSON.stringify(staleAddress) }) + .mockResolvedValueOnce({ result: JSON.stringify(freshAddress) }) + + const stale = AddressService.getAddress('ethereum', 0, 'test-wallet-1') + bumpEpoch() + + const fresh = AddressService.getAddress('ethereum', 0, 'test-wallet-1') + const [staleResult, freshResult] = await Promise.all([stale, fresh]) + + expect(mockHRPC.callMethod).toHaveBeenCalledTimes(2) + expect(staleResult).toBe(staleAddress) + expect(freshResult).toBe(freshAddress) + }) }) }) diff --git a/tests/services/balanceService.test.ts b/tests/services/balanceService.test.ts index 5a29c37..372b9c0 100644 --- a/tests/services/balanceService.test.ts +++ b/tests/services/balanceService.test.ts @@ -12,14 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -import { BalanceService } from '../../src/services/balanceService' +import { BalanceService, fetchBalance, fetchBalances } from '../../src/services/balanceService' import { getWalletStore } from '../../src/store/walletStore' +import { AccountService } from '../../src/services/accountService' +import { bumpEpoch } from '../../src/utils/workletEpoch' +import type { IAsset } from '../../src/types' // Mock stores jest.mock('../../src/store/walletStore', () => ({ getWalletStore: jest.fn(), })) +jest.mock('../../src/services/accountService', () => ({ + AccountService: { callAccountMethod: jest.fn() }, +})) + const MOCK_NATIVE_TOKEN_ID = 'eth-native' describe('BalanceService', () => { @@ -417,4 +424,186 @@ describe('BalanceService', () => { }) }) }) +}) + +describe('fetchBalance / fetchBalances', () => { + let mockWalletStore: any + const mockAccountService = AccountService as jest.Mocked + + const nativeAsset: IAsset = { + getId: () => 'eth-native', + getNetwork: () => 'ethereum', + isNative: () => true, + getContractAddress: () => null, + getSymbol: () => 'ETH', + getName: () => 'Ethereum', + getDecimals: () => 18, + } as unknown as IAsset + + const usdtAsset: IAsset = { + getId: () => 'usdt', + getNetwork: () => 'ethereum', + isNative: () => false, + getContractAddress: () => '0xabc', + getSymbol: () => 'USDT', + getName: () => 'Tether', + getDecimals: () => 6, + } as unknown as IAsset + + function replaySetStateCalls(setStateMock: jest.Mock) { + let state: any = { balances: {}, balanceLoading: {}, lastBalanceUpdate: {} } + for (const [updater] of setStateMock.mock.calls) { + state = typeof updater === 'function' ? { ...state, ...updater(state) } : { ...state, ...updater } + } + return state + } + + // Overwrites getState with a fresh mock returning the given active wallet - + // used to simulate a wallet switch by calling this again mid-test. + function mockActiveWallet(walletId: string) { + mockWalletStore.getState = jest.fn(() => ({ + balances: {}, + balanceLoading: {}, + lastBalanceUpdate: {}, + activeWalletId: walletId, + })) + } + + beforeEach(() => { + jest.clearAllMocks() + + mockWalletStore = { getState: jest.fn(), setState: jest.fn() } + mockActiveWallet('active-wallet') + ;(getWalletStore as jest.Mock).mockReturnValue(mockWalletStore) + }) + + describe('wallet capture at kickoff', () => { + it('attributes the write to the wallet active when the fetch started, not whatever is active when it resolves', async () => { + let resolveCallMethod: (value: string) => void + mockAccountService.callAccountMethod.mockImplementation( + () => new Promise((resolve) => { resolveCallMethod = resolve }), + ) + + const pending = fetchBalance(0, nativeAsset) // captures 'active-wallet' + + // Simulate a wallet switch happening while the fetch is in flight. + mockActiveWallet('wallet-b') + resolveCallMethod!('1000000000000000000') + + await pending + + const finalState = replaySetStateCalls(mockWalletStore.setState) + expect(finalState.balances['active-wallet']?.ethereum?.[0]?.[nativeAsset.getId()]).toBe( + '1000000000000000000', + ) + expect(finalState.balances['wallet-b']).toBeUndefined() + }) + }) + + describe('dedupe', () => { + it('shares a single in-flight call for concurrent fetchBalance calls to the same wallet/network/account/asset', async () => { + mockAccountService.callAccountMethod.mockResolvedValue('1000000000000000000') + + const [a, b] = await Promise.all([ + fetchBalance(0, nativeAsset), + fetchBalance(0, nativeAsset), + ]) + + expect(mockAccountService.callAccountMethod).toHaveBeenCalledTimes(1) + expect(a).toEqual(b) + }) + + it('shares a single in-flight fetch between fetchBalance and fetchBalances for the same asset', async () => { + mockAccountService.callAccountMethod.mockResolvedValue('1000000000000000000') + + const [single, batch] = await Promise.all([ + fetchBalance(0, nativeAsset), + fetchBalances(0, [nativeAsset]), + ]) + + expect(mockAccountService.callAccountMethod).toHaveBeenCalledTimes(1) + expect(single.balance).toBe('1000000000000000000') + expect(batch[0]?.balance).toBe('1000000000000000000') + }) + + it('does not dedupe calls for different wallets', async () => { + mockAccountService.callAccountMethod.mockResolvedValue('1000000000000000000') + + // Both calls are issued synchronously (no await in between), so each + // captures its wallet and stakes its cache-key claim before the other runs. + const p1 = fetchBalance(0, nativeAsset) // captures 'active-wallet' + mockActiveWallet('other-wallet') + const p2 = fetchBalance(0, nativeAsset) // captures 'other-wallet' + + await Promise.all([p1, p2]) + + expect(mockAccountService.callAccountMethod).toHaveBeenCalledTimes(2) + }) + + it('does not dedupe calls for different assets', async () => { + mockAccountService.callAccountMethod.mockResolvedValue('1000000000000000000') + + await Promise.all([ + fetchBalance(0, nativeAsset), + fetchBalance(0, usdtAsset), + ]) + + expect(mockAccountService.callAccountMethod).toHaveBeenCalledTimes(2) + }) + }) + + describe('epoch fencing', () => { + it('fetchBalance skips the store write after a worklet epoch bump', async () => { + let resolveCallMethod: (value: string) => void + mockAccountService.callAccountMethod.mockImplementation( + () => new Promise((resolve) => { resolveCallMethod = resolve }), + ) + + const pending = fetchBalance(0, nativeAsset) + + // Simulate a wallet switch/lock/delete happening mid-flight. + bumpEpoch() + resolveCallMethod!('1000000000000000000') + + const result = await pending + expect(result.success).toBe(true) + + const finalState = replaySetStateCalls(mockWalletStore.setState) + expect(finalState.balances['active-wallet']).toBeUndefined() + }) + + it('fetchBalances skips the store write after a worklet epoch bump', async () => { + let resolveCallMethod: (value: string) => void + mockAccountService.callAccountMethod.mockImplementation( + () => new Promise((resolve) => { resolveCallMethod = resolve }), + ) + + const pending = fetchBalances(0, [nativeAsset]) + + bumpEpoch() + resolveCallMethod!('1000000000000000000') + + const results = await pending + expect(results[0]?.success).toBe(true) + + const finalState = replaySetStateCalls(mockWalletStore.setState) + expect(finalState.balances['active-wallet']).toBeUndefined() + }) + + it('starts a fresh fetchBalance call after an epoch bump, instead of joining the stale in-flight promise', async () => { + mockAccountService.callAccountMethod + .mockResolvedValueOnce('1000000000000000000') + .mockResolvedValueOnce('2000000000000000000') + + const stale = fetchBalance(0, nativeAsset) + bumpEpoch() + const fresh = fetchBalance(0, nativeAsset) + + const [staleResult, freshResult] = await Promise.all([stale, fresh]) + + expect(mockAccountService.callAccountMethod).toHaveBeenCalledTimes(2) + expect(staleResult.balance).toBe('1000000000000000000') + expect(freshResult.balance).toBe('2000000000000000000') + }) + }) }) \ No newline at end of file diff --git a/tests/services/workletLifecycleService.test.ts b/tests/services/workletLifecycleService.test.ts index 79fc4a0..59ca570 100644 --- a/tests/services/workletLifecycleService.test.ts +++ b/tests/services/workletLifecycleService.test.ts @@ -13,9 +13,11 @@ // limitations under the License. import { WorkletLifecycleService } from '../../src/services/workletLifecycleService' +import { getWalletStore } from '../../src/store/walletStore' import { getWorkletStore } from '../../src/store/workletStore' import type { WdkConfigs, BundleConfig } from '../../src/types' import { createResolvablePromise } from '../../src/utils/promise' +import { getEpoch } from '../../src/utils/workletEpoch' import HRPC from '@tetherto/pear-wrk-wdk/hrpc' const mockWorkletInstance = { @@ -553,5 +555,70 @@ describe('WorkletLifecycleService', () => { config: JSON.stringify(defaultNetworkConfigs), }) }) + + it('bumps the worklet epoch - reset() alone does not cover the gap before the new seed actually loads', async () => { + mockInitializeWDK.mockClear() + mockInitializeWDK.mockResolvedValue({ status: 'success' }) + + const initPromise = createResolvablePromise() + const startPromise = createResolvablePromise() + startPromise.resolve(true) + initPromise.resolve(true) + + ;(mockStore as any).getState = jest.fn(() => ({ + isWorkletStarted: true, + isInitialized: false, + isLoading: false, + worklet: mockWorkletInstance, + hrpc: mockHRPCInstance, + ipc: mockWorkletInstance.IPC, + error: null, + wdkConfigs: defaultNetworkConfigs, + workletStartResult: null, + wdkInitResult: null, + isWorkletInitializedPromise: initPromise, + isWorkletStartedPromise: startPromise, + })) + + const before = getEpoch() + + await WorkletLifecycleService.initializeWDK({ + encryptionKey: 'key', + encryptedSeed: 'seed', + }) + + expect(getEpoch()).toBe(before + 1) + }) + }) + + describe('reset', () => { + it('clears worklet init state and addresses', () => { + getWalletStore().setState({ addresses: { 'wallet-1': { ethereum: { 0: '0xabc' } } } }) + + WorkletLifecycleService.reset() + + expect(mockSharedStore.setState).toHaveBeenCalledWith( + expect.objectContaining({ isInitialized: false, wdkInitResult: null }), + ) + expect(getWalletStore().getState().addresses).toEqual({}) + }) + + it('bumps the worklet epoch, so any fetch already in flight for the previous wallet can detect it went stale', () => { + const before = getEpoch() + + WorkletLifecycleService.reset() + + expect(getEpoch()).toBe(before + 1) + }) + + it('bumps the epoch again on every call - delete, lock, and switch all route through reset()', () => { + const before = getEpoch() + + WorkletLifecycleService.reset() + WorkletLifecycleService.reset() + WorkletLifecycleService.reset() + + expect(getEpoch()).toBe(before + 3) + }) }) }) diff --git a/tests/utils/requestCoordinator.test.ts b/tests/utils/requestCoordinator.test.ts new file mode 100644 index 0000000..e959f7b --- /dev/null +++ b/tests/utils/requestCoordinator.test.ts @@ -0,0 +1,190 @@ +// Copyright 2026 Tether Operations Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * Tests for RequestCoordinator + * + * The underlying worklet epoch is a module-level singleton shared across the + * whole file (see workletEpoch.test.ts), so assertions here compare + * before/after transitions rather than hardcoded values. + */ + +import { RequestCoordinator } from '../../src/utils/requestCoordinator' +import { bumpEpoch } from '../../src/utils/workletEpoch' + +function deferred() { + let resolve!: (value: T) => void + let reject!: (reason?: unknown) => void + const promise = new Promise((res, rej) => { + resolve = res + reject = rej + }) + return { promise, resolve, reject } +} + +describe('RequestCoordinator', () => { + describe('dedupe', () => { + it('shares a single in-flight fetch for the same key', async () => { + const coordinator = new RequestCoordinator() + const fetch = jest.fn().mockResolvedValue('value') + const commit = jest.fn() + + const [a, b] = await Promise.all([ + coordinator.run('key-1', fetch, commit), + coordinator.run('key-1', fetch, commit), + ]) + + expect(a).toBe('value') + expect(b).toBe('value') + expect(fetch).toHaveBeenCalledTimes(1) + expect(commit).toHaveBeenCalledTimes(1) + expect(commit).toHaveBeenCalledWith('value') + }) + + it('does not dedupe different keys', async () => { + const coordinator = new RequestCoordinator() + const fetch = jest.fn().mockResolvedValue('value') + + await Promise.all([ + coordinator.run('key-1', fetch, () => {}), + coordinator.run('key-2', fetch, () => {}), + ]) + + expect(fetch).toHaveBeenCalledTimes(2) + }) + + it('starts a fresh call once the previous one has settled', async () => { + const coordinator = new RequestCoordinator() + const fetch = jest.fn().mockResolvedValue('value') + + await coordinator.run('key-1', fetch, () => {}) + await coordinator.run('key-1', fetch, () => {}) + + expect(fetch).toHaveBeenCalledTimes(2) + }) + + it('cleans up the in-flight entry even when the fetch rejects', async () => { + const coordinator = new RequestCoordinator() + const error = new Error('boom') + + await expect( + coordinator.run('key-1', () => Promise.reject(error), () => {}), + ).rejects.toThrow('boom') + + const fetch = jest.fn().mockResolvedValue('fresh') + const result = await coordinator.run('key-1', fetch, () => {}) + + expect(result).toBe('fresh') + expect(fetch).toHaveBeenCalledTimes(1) + }) + + it('propagates rejection to every waiter sharing the in-flight fetch', async () => { + const coordinator = new RequestCoordinator() + const error = new Error('boom') + const { promise, reject } = deferred() + + const p1 = coordinator.run('key-1', () => promise, () => {}) + const p2 = coordinator.run('key-1', () => promise, () => {}) + + reject(error) + + await expect(p1).rejects.toThrow('boom') + await expect(p2).rejects.toThrow('boom') + }) + + it('does not call commit when the fetch rejects', async () => { + const coordinator = new RequestCoordinator() + const commit = jest.fn() + + await expect( + coordinator.run('key-1', () => Promise.reject(new Error('boom')), commit), + ).rejects.toThrow('boom') + + expect(commit).not.toHaveBeenCalled() + }) + }) + + describe('epoch fencing', () => { + it('commits when the epoch has not changed since the fetch started', async () => { + const coordinator = new RequestCoordinator() + const commit = jest.fn() + + await coordinator.run('key-1', () => Promise.resolve('value'), commit) + + expect(commit).toHaveBeenCalledWith('value') + }) + + it('does not commit if the epoch was bumped while the fetch was in flight', async () => { + const coordinator = new RequestCoordinator() + const commit = jest.fn() + const pending = deferred() + + const call = coordinator.run('key-1', () => pending.promise, commit) + + bumpEpoch() + pending.resolve('value') + + const result = await call + expect(result).toBe('value') + expect(commit).not.toHaveBeenCalled() + }) + + it('still returns the fetched value even when the result is discarded as stale', async () => { + const coordinator = new RequestCoordinator() + const pending = deferred() + + const call = coordinator.run('key-1', () => pending.promise, () => {}) + + bumpEpoch() + pending.resolve('stale-but-still-returned') + + expect(await call).toBe('stale-but-still-returned') + }) + + it('does not reuse an in-flight fetch from a previous epoch - starts a fresh one instead', async () => { + const coordinator = new RequestCoordinator() + const stale = deferred() + const fetch = jest + .fn() + .mockReturnValueOnce(stale.promise) + .mockResolvedValueOnce('fresh') + + const staleCall = coordinator.run('wallet-a:eth:0', fetch, () => {}) + bumpEpoch() + const freshCall = coordinator.run('wallet-a:eth:0', fetch, () => {}) + + stale.resolve('stale') + + expect(await freshCall).toBe('fresh') + expect(await staleCall).toBe('stale') + expect(fetch).toHaveBeenCalledTimes(2) + }) + }) + + describe('multiple instances', () => { + it('does not affect in-flight entries under a different coordinator instance', async () => { + const addressCoordinator = new RequestCoordinator() + const balanceCoordinator = new RequestCoordinator() + const other = deferred() + + const otherCall = balanceCoordinator.run('wallet-b:eth:0', () => other.promise, () => {}) + // Using the same key on a different instance shouldn't disturb it. + void addressCoordinator.run('wallet-b:eth:0', () => Promise.resolve('unrelated'), () => {}) + + other.resolve('original') + + expect(await otherCall).toBe('original') + }) + }) +}) diff --git a/tests/utils/workletEpoch.test.ts b/tests/utils/workletEpoch.test.ts new file mode 100644 index 0000000..b88b26d --- /dev/null +++ b/tests/utils/workletEpoch.test.ts @@ -0,0 +1,61 @@ +// Copyright 2026 Tether Operations Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * Tests for workletEpoch + * + * The epoch is a true module-level singleton (intentionally - it's the one + * global invalidation signal shared by every RequestCoordinator instance), + * so these tests assert relative transitions rather than absolute values. + */ + +import { getEpoch, bumpEpoch } from '../../src/utils/workletEpoch' + +describe('workletEpoch', () => { + it('stays the same across repeated reads without a bump', () => { + const epoch = getEpoch() + + expect(getEpoch()).toBe(epoch) + expect(getEpoch()).toBe(epoch) + }) + + it('strictly increases by one on each bumpEpoch call', () => { + const before = getEpoch() + + bumpEpoch() + + expect(getEpoch()).toBe(before + 1) + }) + + it('accumulates across multiple bumps', () => { + const before = getEpoch() + + bumpEpoch() + bumpEpoch() + bumpEpoch() + + expect(getEpoch()).toBe(before + 3) + }) + + it('never decreases or wraps back to a previously seen value', () => { + const seen = new Set() + + for (let i = 0; i < 5; i++) { + seen.add(getEpoch()) + bumpEpoch() + } + + expect(seen.size).toBe(5) + }) +})