From 7f0921389ce3d25df4f336d967a3166e53af303c Mon Sep 17 00:00:00 2001 From: Aliaksei Shaltykou Date: Fri, 31 Jul 2026 17:06:48 +0200 Subject: [PATCH] feat: implement wallet address invalidation and deduplication logic --- src/hooks/useWalletManager.ts | 5 + src/services/addressService.ts | 131 +++++++++++++++++++++--- src/services/workletLifecycleService.ts | 4 + tests/services/addressService.test.ts | 78 ++++++++++++++ 4 files changed, 206 insertions(+), 12 deletions(-) diff --git a/src/hooks/useWalletManager.ts b/src/hooks/useWalletManager.ts index 8bdd462..1ace82f 100644 --- a/src/hooks/useWalletManager.ts +++ b/src/hooks/useWalletManager.ts @@ -14,6 +14,7 @@ import { produce } from 'immer' import { useMemo, useCallback } from 'react' +import { AddressService } from '../services/addressService' import { WalletSetupService } from '../services/walletSetupService' import { WorkletLifecycleService } from '../services/workletLifecycleService' import { @@ -363,6 +364,10 @@ export function useWalletManager(): UseWalletManagerResult { } try { + // Invalidate before clearing store so in-flight getAddress calls + // skip stale writes under this walletId. + AddressService.invalidateWallet(walletId) + await WalletSetupService.deleteWallet(walletId) walletStore.setState((prev) => diff --git a/src/services/addressService.ts b/src/services/addressService.ts index b45503d..6310603 100644 --- a/src/services/addressService.ts +++ b/src/services/addressService.ts @@ -38,6 +38,71 @@ import { AddressInfoResult } from '../types' * Provides methods for retrieving and caching wallet addresses. */ export class AddressService { + /** + * Dedupes concurrent getAddress calls: walletId:network:accountIndex -> Promise + */ + private static inflight = new Map>() + + /** + * Per-wallet epoch; bumped on invalidate so in-flight writes are skipped + */ + private static walletEpochs = new Map() + + /** + * Global epoch; bumped on invalidateAll (e.g. lock/reset) + */ + private static globalEpoch = 0 + + /** + * Generate a cache key for a specific wallet, network, and account index + */ + private static cacheKey( + walletId: string, + network: string, + accountIndex: number, + ): string { + return `${walletId}:${network}:${accountIndex}` + } + + /** + * Get the epoch for a specific wallet + */ + private static getEpoch(walletId: string): number { + return this.globalEpoch + (this.walletEpochs.get(walletId) ?? 0) + } + + /** + * Clear the inflight map for a specific wallet + */ + private static clearInflightForWallet(walletId: string): void { + const prefix = `${walletId}:` + for (const key of [...this.inflight.keys()]) { + if (key.startsWith(prefix)) { + this.inflight.delete(key) + } + } + } + + /** + * Invalidate in-flight address fetches for a wallet. + * Clears the promise cache for that wallet and bumps its epoch so any + * still-running fetch will not write to the store when it completes. + */ + static invalidateWallet(walletId: string): void { + this.clearInflightForWallet(walletId) + this.walletEpochs.set(walletId, (this.walletEpochs.get(walletId) ?? 0) + 1) + log('[AddressService] Invalidated wallet address fetches', { walletId }) + } + + /** + * Invalidate all in-flight address fetches (e.g. on lock/reset). + */ + static invalidateAll(): void { + this.inflight.clear() + this.globalEpoch += 1 + log('[AddressService] Invalidated all address fetches') + } + /** * Get address for a specific network and account index * Caches the address in walletStore for future use @@ -64,17 +129,48 @@ export class AddressService { return cachedAddress } + const cacheKey = this.cacheKey(targetWalletId, network, accountIndex) + const existing = this.inflight.get(cacheKey) + if (existing) { + return existing + } + + const promise = this.fetchAndCacheAddress( + network, + accountIndex, + targetWalletId, + ) + this.inflight.set(cacheKey, promise) + + try { + return await promise + } finally { + if (this.inflight.get(cacheKey) === promise) { + this.inflight.delete(cacheKey) + } + } + } + + private static async fetchAndCacheAddress( + network: string, + accountIndex: number, + targetWalletId: string, + ): Promise { + const walletStore = getWalletStore() + const epoch = this.getEpoch(targetWalletId) const hrpc = await requireInitialized() const loadingKey = `${network}-${accountIndex}` try { - walletStore.setState((prev) => - produce(prev, (state) => { - state.walletLoading[targetWalletId] ??= {} - state.walletLoading[targetWalletId][loadingKey] = true - }), - ) + if (this.getEpoch(targetWalletId) === epoch) { + walletStore.setState((prev) => + produce(prev, (state) => { + state.walletLoading[targetWalletId] ??= {} + state.walletLoading[targetWalletId][loadingKey] = true + }), + ) + } const response = await hrpc.callMethod({ methodName: 'getAddress', @@ -101,6 +197,15 @@ export class AddressService { ) } + // Skip store writes if wallet was invalidated while this fetch was in flight + if (this.getEpoch(targetWalletId) !== epoch) { + log( + '[AddressService] Skipping stale address write after invalidation', + { walletId: targetWalletId, network, accountIndex }, + ) + return address + } + // Cache the address using helper (per-wallet) walletStore.setState((prev) => produce( @@ -120,12 +225,14 @@ export class AddressService { return address } catch (error) { - walletStore.setState((prev) => - produce(prev, (state) => { - state.walletLoading[targetWalletId] ??= {} - state.walletLoading[targetWalletId][loadingKey] = false - }), - ) + if (this.getEpoch(targetWalletId) === epoch) { + walletStore.setState((prev) => + produce(prev, (state) => { + state.walletLoading[targetWalletId] ??= {} + state.walletLoading[targetWalletId][loadingKey] = false + }), + ) + } handleServiceError(error, 'AddressService', 'getAddress', { network, diff --git a/src/services/workletLifecycleService.ts b/src/services/workletLifecycleService.ts index 2a2b1a1..aeebe73 100644 --- a/src/services/workletLifecycleService.ts +++ b/src/services/workletLifecycleService.ts @@ -21,6 +21,7 @@ import { Worklet } from 'react-native-bare-kit' +import { AddressService } from './addressService' import { getWalletStore } from '../store/walletStore' import { getWorkletStore } from '../store/workletStore' import { DEFAULT_MNEMONIC_WORD_COUNT } from '../utils/constants' @@ -532,6 +533,9 @@ export class WorkletLifecycleService { workletStore.setState({ isWorkletInitializedPromise: createResolvablePromise() }) + // Drop in-flight address fetches so late responses cannot repopulate the store + AddressService.invalidateAll() + // Clear addresses from wallet store walletStore.setState({ addresses: {}, diff --git a/tests/services/addressService.test.ts b/tests/services/addressService.test.ts index 5b792e0..f1a8f6e 100644 --- a/tests/services/addressService.test.ts +++ b/tests/services/addressService.test.ts @@ -39,6 +39,9 @@ describe('AddressService', () => { beforeEach(() => { jest.clearAllMocks() + // Clear static inflight map / bump global epoch so tests don't leak state + AddressService.invalidateAll() + // Setup mock HRPC mockHRPC = { callMethod: jest.fn(), @@ -323,6 +326,81 @@ 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 skip store writes after invalidateWallet', async () => { + const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0' + mockHRPC.callMethod.mockResolvedValue({ + result: JSON.stringify(address), + }) + + const pending = AddressService.getAddress('ethereum', 0, 'test-wallet-1') + AddressService.invalidateWallet('test-wallet-1') + const result = await pending + + expect(result).toBe(address) + expect(mockWalletStore.setState).not.toHaveBeenCalled() + }) + + it('should skip store writes after invalidateAll', async () => { + const address = '0x742d35Cc6634C0532925a3b844Bc9e7595f0bEb0' + mockHRPC.callMethod.mockResolvedValue({ + result: JSON.stringify(address), + }) + + const pending = AddressService.getAddress('ethereum', 0, 'test-wallet-1') + AddressService.invalidateAll() + const result = await pending + + expect(result).toBe(address) + expect(mockWalletStore.setState).not.toHaveBeenCalled() + }) + + it('should start a fresh fetch after invalidate clears inflight', 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') + AddressService.invalidateWallet('test-wallet-1') + + 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) + }) }) })