feat(base-data-service): Support persisting the cache

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",

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',
+      );
+    });
+  });
 });

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<ServiceName extends string> = {
 type DataServiceActions<ServiceName extends string> =
   DataServiceInvalidateQueriesAction<ServiceName>;
 
+type DataServiceAllowedActions =
+  | StorageServiceGetItemAction
+  | StorageServiceSetItemAction;
+
 export type DataServiceCacheUpdatedEvent<ServiceName extends string> = {
   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<ServiceName>,
+    DataServiceActions<ServiceName> | DataServiceAllowedActions,
     DataServiceEvents<ServiceName>
   >;
 
@@ -105,24 +125,28 @@ 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;
 
     // We are storing a separately typed messenger for known actions and events provided by data services
     // 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<ServiceName>,
+      DataServiceActions<ServiceName> | DataServiceAllowedActions,
       DataServiceEvents<ServiceName>
     >;
     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<void> {
+    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<void> {
+    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);
+  }
 }

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(

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"]
 }

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"]
 }

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"