From 7813f8906d6a12e2bf72386fa5d79e6bf6f07932 Mon Sep 17 00:00:00 2001 From: Frederik Bolding Date: Thu, 9 Jul 2026 14:18:06 +0200 Subject: [PATCH] feat(base-data-service): Support persisting the cache --- packages/base-data-service/package.json | 1 + .../src/BaseDataService.test.ts | 167 +++++++++++++++++- .../base-data-service/src/BaseDataService.ts | 88 ++++++++- .../tests/ExampleDataService.ts | 3 + .../base-data-service/tsconfig.build.json | 3 +- packages/base-data-service/tsconfig.json | 6 +- yarn.lock | 1 + 7 files changed, 264 insertions(+), 5 deletions(-) diff --git a/packages/base-data-service/package.json b/packages/base-data-service/package.json index 4fb523fa97..c8e32d2b4b 100644 --- a/packages/base-data-service/package.json +++ b/packages/base-data-service/package.json @@ -55,6 +55,7 @@ }, "dependencies": { "@metamask/messenger": "^2.0.0", + "@metamask/storage-service": "^1.0.2", "@metamask/utils": "^11.11.0", "@tanstack/query-core": "^4.43.0", "cockatiel": "^3.1.2", diff --git a/packages/base-data-service/src/BaseDataService.test.ts b/packages/base-data-service/src/BaseDataService.test.ts index c4578e6ae4..2c783fb8f9 100644 --- a/packages/base-data-service/src/BaseDataService.test.ts +++ b/packages/base-data-service/src/BaseDataService.test.ts @@ -1,4 +1,4 @@ -import { Messenger } from '@metamask/messenger'; +import { MOCK_ANY_NAMESPACE, Messenger } from '@metamask/messenger'; import { hashQueryKey } from '@tanstack/query-core'; import { BrokenCircuitError } from 'cockatiel'; import { cleanAll } from 'nock'; @@ -280,4 +280,169 @@ describe('BaseDataService', () => { ); }); }); + + describe('persistence', () => { + it('persists the cache using the StorageService', async () => { + const rootMessenger = new Messenger({ + namespace: MOCK_ANY_NAMESPACE, + captureException: console.error, + }); + + const setItem = jest.fn(); + rootMessenger.registerActionHandler('StorageService:setItem', setItem); + + const messenger = rootMessenger.buildChild({ + namespace: serviceName, + actions: ['StorageService:getItem', 'StorageService:setItem'], + }); + const service = new ExampleDataService(messenger); + + mockAssets(); + + await service.getAssets(MOCK_ASSETS); + + expect(setItem).toHaveBeenCalledWith(serviceName, 'cache', { + state: { + queries: [ + { + queryHash: + '["ExampleDataService:getAssets",["eip155:1/slip44:60","bip122:000000000019d6689c085ae165831e93/slip44:0","eip155:1/erc20:0x6b175474e89094c44da98b954eedeac495271d0f"]]', + queryKey: [ + 'ExampleDataService:getAssets', + [ + 'eip155:1/slip44:60', + 'bip122:000000000019d6689c085ae165831e93/slip44:0', + 'eip155:1/erc20:0x6b175474e89094c44da98b954eedeac495271d0f', + ], + ], + state: { + data: [ + { + assetId: + 'eip155:1/erc20:0x6b175474e89094c44da98b954eedeac495271d0f', + decimals: 18, + name: 'Dai Stablecoin', + symbol: 'DAI', + }, + { + assetId: 'bip122:000000000019d6689c085ae165831e93/slip44:0', + decimals: 8, + name: 'Bitcoin', + symbol: 'BTC', + }, + { + assetId: 'eip155:1/slip44:60', + decimals: 18, + name: 'Ethereum', + symbol: 'ETH', + }, + ], + dataUpdateCount: 1, + dataUpdatedAt: expect.any(Number), + error: null, + errorUpdateCount: 0, + errorUpdatedAt: 0, + fetchFailureCount: 0, + fetchFailureReason: null, + fetchMeta: null, + fetchStatus: 'idle', + isInvalidated: false, + status: 'success', + }, + }, + ], + mutations: [], + }, + timestamp: expect.any(Number), + }); + }); + + it('rehydrates the cache using the StorageService', async () => { + const rootMessenger = new Messenger({ + namespace: MOCK_ANY_NAMESPACE, + captureException: console.error, + }); + + rootMessenger.registerActionHandler('StorageService:setItem', jest.fn()); + rootMessenger.registerActionHandler('StorageService:getItem', () => { + return { + result: { + state: { + queries: [ + { + queryHash: + '["ExampleDataService:getAssets",["eip155:1/slip44:60","bip122:000000000019d6689c085ae165831e93/slip44:0","eip155:1/erc20:0x6b175474e89094c44da98b954eedeac495271d0f"]]', + queryKey: [ + 'ExampleDataService:getAssets', + [ + 'eip155:1/slip44:60', + 'bip122:000000000019d6689c085ae165831e93/slip44:0', + 'eip155:1/erc20:0x6b175474e89094c44da98b954eedeac495271d0f', + ], + ], + state: { + data: [ + { + assetId: + 'eip155:1/erc20:0x6b175474e89094c44da98b954eedeac495271d0f', + decimals: 18, + name: 'Dai Stablecoin', + symbol: 'DAI', + }, + { + assetId: + 'bip122:000000000019d6689c085ae165831e93/slip44:0', + decimals: 8, + name: 'Bitcoin', + symbol: 'BTC', + }, + { + assetId: 'eip155:1/slip44:60', + decimals: 18, + name: 'Ethereum', + symbol: 'ETH', + }, + ], + dataUpdateCount: 1, + dataUpdatedAt: Date.now(), + error: null, + errorUpdateCount: 0, + errorUpdatedAt: 0, + fetchFailureCount: 0, + fetchFailureReason: null, + fetchMeta: null, + fetchStatus: 'idle', + isInvalidated: false, + status: 'success', + }, + }, + ], + mutations: [], + }, + timestamp: Date.now(), + }, + }; + }); + + const messenger = rootMessenger.buildChild({ + namespace: serviceName, + actions: ['StorageService:getItem', 'StorageService:setItem'], + }); + const spy = jest.spyOn(messenger, 'call'); + const service = new ExampleDataService(messenger); + service.init(); + + mockAssets({ status: 500 }); + + const result = await service.getAssets(MOCK_ASSETS); + + expect(result).toHaveLength(3); + + expect(spy).toHaveBeenCalledWith( + 'StorageService:getItem', + serviceName, + 'cache', + ); + }); + }); }); diff --git a/packages/base-data-service/src/BaseDataService.ts b/packages/base-data-service/src/BaseDataService.ts index 888144e6fa..643aac9b61 100644 --- a/packages/base-data-service/src/BaseDataService.ts +++ b/packages/base-data-service/src/BaseDataService.ts @@ -3,6 +3,10 @@ import { ActionConstraint, EventConstraint, } from '@metamask/messenger'; +import type { + StorageServiceGetItemAction, + StorageServiceSetItemAction, +} from '@metamask/storage-service'; import { Duration, inMilliseconds } from '@metamask/utils'; import type { Json } from '@metamask/utils'; import { @@ -18,6 +22,7 @@ import { QueryClientConfig, WithRequired, dehydrate, + hydrate, } from '@tanstack/query-core'; import deepEqual from 'fast-deep-equal'; @@ -55,6 +60,10 @@ export type DataServiceInvalidateQueriesAction = { type DataServiceActions = DataServiceInvalidateQueriesAction; +type DataServiceAllowedActions = + | StorageServiceGetItemAction + | StorageServiceSetItemAction; + export type DataServiceCacheUpdatedEvent = { type: `${ServiceName}:cacheUpdated`; payload: [DataServiceCacheUpdatedPayload]; @@ -77,6 +86,17 @@ const QUERY_CLIENT_DEFAULTS: DefaultOptions = { }, }; +const STORAGE_SERVICE_KEY = 'cache'; + +type PersistConfiguration = { + maxAge: number; +}; + +type PersistedCache = { + state: DehydratedState; + timestamp: number; +}; + export class BaseDataService< ServiceName extends string, ServiceMessenger extends Messenger< @@ -93,7 +113,7 @@ export class BaseDataService< readonly #messenger: Messenger< ServiceName, - DataServiceActions, + DataServiceActions | DataServiceAllowedActions, DataServiceEvents >; @@ -105,16 +125,20 @@ export class BaseDataService< readonly #queryCacheUnsubscribe: () => void; + readonly #persistConfig?: PersistConfiguration; + constructor({ name, messenger, queryClientConfig = {}, policyOptions, + persistConfig, }: { name: ServiceName; messenger: ServiceMessenger; queryClientConfig?: QueryClientConfig; policyOptions?: CreateServicePolicyOptions; + persistConfig?: PersistConfiguration; }) { this.name = name; @@ -122,7 +146,7 @@ export class BaseDataService< // and a generic public one that is typed using the generic parameters and accessible to implementations. this.#messenger = messenger as unknown as Messenger< ServiceName, - DataServiceActions, + DataServiceActions | DataServiceAllowedActions, DataServiceEvents >; this.messenger = messenger; @@ -138,6 +162,8 @@ export class BaseDataService< }, }); + this.#persistConfig = persistConfig; + this.#policy = createServicePolicy(policyOptions); this.#queryCacheUnsubscribe = this.#queryClient @@ -148,6 +174,11 @@ export class BaseDataService< event.query.queryHash, event.type as CacheUpdatedType, ); + + // TODO: Debounce + this.#persistCache().catch((error) => + this.#messenger.captureException?.(error), + ); } }); @@ -265,6 +296,15 @@ export class BaseDataService< return this.#queryClient.invalidateQueries(filters, options); } + /** + * Initialize the service, rehydrating the cache with persisted data if possible. + */ + init(): void { + this.#rehydrateCache().catch((error) => + this.#messenger.captureException?.(error), + ); + } + /** * Prepares the service for garbage collection. This should be extended * by any subclasses to clean up any additional connections or events. @@ -306,4 +346,48 @@ export class BaseDataService< } as DataServiceGranularCacheUpdatedPayload, ); } + + async #persistCache(): Promise { + if (!this.#persistConfig) { + return; + } + + const state = dehydrate(this.#queryClient); + + if (state.queries.length === 0 && state.mutations.length === 0) { + return; + } + + const cache: PersistedCache = { + timestamp: Date.now(), + state, + }; + + await this.#messenger.call( + 'StorageService:setItem', + this.name, + STORAGE_SERVICE_KEY, + cache as unknown as Json, + ); + } + + async #rehydrateCache(): Promise { + if (!this.#persistConfig) { + return; + } + + const { result: persisted } = await this.#messenger.call( + 'StorageService:getItem', + this.name, + STORAGE_SERVICE_KEY, + ); + + const cache = persisted as PersistedCache; + + if (Date.now() - cache.timestamp >= this.#persistConfig.maxAge) { + return; + } + + hydrate(this.#queryClient, cache.state); + } } diff --git a/packages/base-data-service/tests/ExampleDataService.ts b/packages/base-data-service/tests/ExampleDataService.ts index f3bbdc01a6..1df4fa8bb2 100644 --- a/packages/base-data-service/tests/ExampleDataService.ts +++ b/packages/base-data-service/tests/ExampleDataService.ts @@ -69,6 +69,9 @@ export class ExampleDataService extends BaseDataService< maxConsecutiveFailures: 3, backoff: new ConstantBackoff(0), }, + persistConfig: { + maxAge: inMilliseconds(1, Duration.Day), + }, }); this.messenger.registerMethodActionHandlers( diff --git a/packages/base-data-service/tsconfig.build.json b/packages/base-data-service/tsconfig.build.json index f83c71b4af..1c38dc1ad5 100644 --- a/packages/base-data-service/tsconfig.build.json +++ b/packages/base-data-service/tsconfig.build.json @@ -7,7 +7,8 @@ }, "references": [ { "path": "../messenger/tsconfig.build.json" }, - { "path": "../controller-utils/tsconfig.build.json" } + { "path": "../controller-utils/tsconfig.build.json" }, + { "path": "../storage-service/tsconfig.build.json" } ], "include": ["../../types", "./src"] } diff --git a/packages/base-data-service/tsconfig.json b/packages/base-data-service/tsconfig.json index 3dbaffd259..1237a83172 100644 --- a/packages/base-data-service/tsconfig.json +++ b/packages/base-data-service/tsconfig.json @@ -3,6 +3,10 @@ "compilerOptions": { "baseUrl": "./" }, - "references": [{ "path": "../messenger" }, { "path": "../controller-utils" }], + "references": [ + { "path": "../messenger" }, + { "path": "../controller-utils" }, + { "path": "../storage-service" } + ], "include": ["../../types", "./src", "./tests"] } diff --git a/yarn.lock b/yarn.lock index e44825fb42..d68408fe83 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5889,6 +5889,7 @@ __metadata: dependencies: "@metamask/auto-changelog": "npm:^6.1.0" "@metamask/messenger": "npm:^2.0.0" + "@metamask/storage-service": "npm:^1.0.2" "@metamask/utils": "npm:^11.11.0" "@tanstack/query-core": "npm:^4.43.0" "@ts-bridge/cli": "npm:^0.6.4"