fix: use in-memory cache in price service (#427)

* Swap `StateCache` to `InMemoryCache` in `PriceApiClient`
* Implement TTL logic for the `InMemoryCache`

Author: aganglada

packages/snap/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snap-solana-wallet.git"
   },
   "source": {
-    "shasum": "fYdnpG929p3sfCbwUNQQe4P+LSPyGQdZtXgdMUtt97E=",
+    "shasum": "XiO/mFkK3oitUU/FZOqAC3nBmYGcP2ZIPt0notbfeaY=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/snap/src/core/caching/InMemoryCache.ts

@@ -1,69 +1,182 @@
+import { assert } from '@metamask/utils';
+
 import type { Serializable } from '../serialization/types';
+import type { ILogger } from '../utils/logger';
 import type { ICache } from './ICache';
+import type { CacheEntry } from './types';
 
 /**
- * A simple in-memory cache implementation, primarily used for testing purposes.
+ * A simple in-memory cache implementation supporting TTL (Time To Live) functionality.
  *
  * WARNINGS:
  * - This cache is not persistent and will be lost when the process is restarted.
- * - It does not support TTL.
  */
 export class InMemoryCache implements ICache<Serializable> {
-  #cache: Map<string, Serializable> = new Map();
+  #cache: Map<string, CacheEntry> = new Map();
+
+  public readonly logger: ILogger;
+
+  constructor(logger: ILogger) {
+    this.logger = logger;
+  }
+
+  #validateTtlOrThrow(ttlMilliseconds?: number): void {
+    if (ttlMilliseconds === undefined) {
+      return;
+    }
+
+    if (typeof ttlMilliseconds !== 'number') {
+      throw new Error('TTL must be a number');
+    }
+
+    if (ttlMilliseconds < 0) {
+      throw new Error('TTL must be positive');
+    }
+
+    if (ttlMilliseconds > Number.MAX_SAFE_INTEGER) {
+      throw new Error('TTL must be less than 2^53 - 1');
+    }
+  }
+
+  #isExpired(cacheEntry: CacheEntry): boolean {
+    return cacheEntry.expiresAt < Date.now();
+  }
+
+  async #cleanupExpiredEntries(): Promise<void> {
+    const expiredKeys: string[] = [];
+    for (const [key, entry] of this.#cache.entries()) {
+      if (this.#isExpired(entry)) {
+        expiredKeys.push(key);
+      }
+    }
+    await this.mdelete(expiredKeys);
+  }
 
   async get(key: string): Promise<Serializable | undefined> {
-    return this.#cache.get(key);
+    const result = await this.mget([key]);
+    return result[key];
   }
 
   async set(
     key: string,
     value: Serializable,
-    _ttlMilliseconds?: number,
+    ttlMilliseconds = Number.MAX_SAFE_INTEGER,
   ): Promise<void> {
-    this.#cache.set(key, value);
+    this.#validateTtlOrThrow(ttlMilliseconds);
+
+    this.#cache.set(key, {
+      value,
+      expiresAt: Math.min(
+        Date.now() + (ttlMilliseconds ?? Number.MAX_SAFE_INTEGER),
+        Number.MAX_SAFE_INTEGER,
+      ),
+    });
   }
 
   async delete(key: string): Promise<boolean> {
-    return this.#cache.delete(key);
+    const result = await this.mdelete([key]);
+    return result[key] ?? false;
   }
 
   async clear(): Promise<void> {
     this.#cache.clear();
   }
 
   async has(key: string): Promise<boolean> {
-    return this.#cache.has(key);
+    const cacheEntry = this.#cache.get(key);
+    if (!cacheEntry) {
+      return false;
+    }
+
+    if (this.#isExpired(cacheEntry)) {
+      this.#cache.delete(key);
+      return false;
+    }
+
+    return true;
   }
 
   async keys(): Promise<string[]> {
+    await this.#cleanupExpiredEntries();
     return Array.from(this.#cache.keys());
   }
 
   async size(): Promise<number> {
+    await this.#cleanupExpiredEntries();
     return this.#cache.size;
   }
 
   async peek(key: string): Promise<Serializable | undefined> {
-    return this.#cache.get(key);
+    const cacheEntry = this.#cache.get(key);
+    if (!cacheEntry) {
+      return undefined;
+    }
+
+    if (this.#isExpired(cacheEntry)) {
+      this.#cache.delete(key);
+      return undefined;
+    }
+
+    return cacheEntry.value;
   }
 
-  async mget(keys: string[]): Promise<Record<string, Serializable>> {
-    return Object.fromEntries(
-      keys.map((key) => [key, this.#cache.get(key)]),
-    ) as Record<string, Serializable>;
+  async mget(
+    keys: string[],
+  ): Promise<Record<string, Serializable | undefined>> {
+    await this.#cleanupExpiredEntries();
+
+    const result: Record<string, Serializable | undefined> = {};
+
+    for (const key of keys) {
+      const cacheEntry = this.#cache.get(key);
+      if (!cacheEntry) {
+        this.logger.info(`[InMemoryCache] ❌ Cache miss for key "${key}"`);
+        result[key] = undefined;
+        continue;
+      }
+
+      this.logger.info(`[InMemoryCache] 🎉 Cache hit for key "${key}"`);
+      result[key] = cacheEntry.value;
+    }
+
+    return result;
   }
 
   async mset(
-    entries: { key: string; value: Serializable; _ttlMilliseconds?: number }[],
+    entries: { key: string; value: Serializable; ttlMilliseconds?: number }[],
   ): Promise<void> {
-    entries.forEach(({ key, value }) => {
-      this.#cache.set(key, value);
+    if (entries.length === 0) {
+      return;
+    }
+
+    if (entries.length === 1) {
+      assert(entries[0]); // Enforce type narrowing as TS cannot infer that entries[0] is defined
+      const { key, value, ttlMilliseconds } = entries[0];
+      await this.set(key, value, ttlMilliseconds);
+      return;
+    }
+
+    entries.forEach(({ ttlMilliseconds }) => {
+      this.#validateTtlOrThrow(ttlMilliseconds);
+    });
+
+    entries.forEach(({ key, value, ttlMilliseconds }) => {
+      if (value === undefined) {
+        return;
+      }
+      this.#cache.set(key, {
+        value,
+        expiresAt: Math.min(
+          Date.now() + (ttlMilliseconds ?? Number.MAX_SAFE_INTEGER),
+          Number.MAX_SAFE_INTEGER,
+        ),
+      });
     });
   }
 
   async mdelete(keys: string[]): Promise<Record<string, boolean>> {
     return Object.fromEntries(
       keys.map((key) => [key, this.#cache.delete(key)]),
-    ) as Record<string, boolean>;
+    );
   }
 }

packages/snap/src/core/caching/StateCache.ts

@@ -6,16 +6,7 @@ import type { Serializable } from '../serialization/types';
 import type { IStateManager } from '../services/state/IStateManager';
 import type { ILogger } from '../utils/logger';
 import type { ICache } from './ICache';
-
-export type TimestampMilliseconds = number;
-
-/**
- * A single cache entry.
- */
-export type CacheEntry = {
-  value: Serializable;
-  expiresAt: TimestampMilliseconds;
-};
+import type { CacheEntry } from './types';
 
 /**
  * The whole cache store.

packages/snap/src/core/caching/types.ts

@@ -0,0 +1,11 @@
+import type { Serializable } from '../serialization/types';
+
+export type TimestampMilliseconds = number;
+
+/**
+ * A single cache entry.
+ */
+export type CacheEntry = {
+  value: Serializable;
+  expiresAt: TimestampMilliseconds;
+};

packages/snap/src/core/clients/price-api/PriceApiClient.test.ts

@@ -31,7 +31,7 @@ describe('PriceApiClient', () => {
       }),
     } as unknown as ConfigProvider;
 
-    mockCache = new InMemoryCache();
+    mockCache = new InMemoryCache(mockLogger);
 
     client = new PriceApiClient(
       mockConfigProvider,

packages/snap/src/core/services/assets/AssetsService.test.ts

@@ -64,7 +64,7 @@ describe('AssetsService', () => {
 
     mockState = new InMemoryState(DEFAULT_UNENCRYPTED_STATE);
 
-    mockCache = new InMemoryCache();
+    mockCache = new InMemoryCache(logger);
 
     stateUpdateSpy = jest.spyOn(mockState, 'update');
 

packages/snap/src/snapContext.ts

@@ -1,4 +1,5 @@
 import type { ICache } from './core/caching/ICache';
+import { InMemoryCache } from './core/caching/InMemoryCache';
 import { StateCache } from './core/caching/StateCache';
 import { PriceApiClient } from './core/clients/price-api/PriceApiClient';
 import { SecurityAlertsApiClient } from './core/clients/security-alerts-api/SecurityAlertsApiClient';
@@ -56,7 +57,8 @@ const state = new State({
   defaultState: DEFAULT_UNENCRYPTED_STATE,
 });
 
-const cache = new StateCache(state, logger);
+const stateCache = new StateCache(state, logger);
+const inMemoryCache = new InMemoryCache(logger);
 
 const connection = new SolanaConnection(configProvider);
 const transactionHelper = new TransactionHelper(connection, logger);
@@ -67,7 +69,7 @@ const sendSplTokenBuilder = new SendSplTokenBuilder(
   logger,
 );
 const tokenMetadataClient = new TokenMetadataClient(configProvider);
-const priceApiClient = new PriceApiClient(configProvider, cache);
+const priceApiClient = new PriceApiClient(configProvider, inMemoryCache);
 
 const tokenMetadataService = new TokenMetadataService({
   tokenMetadataClient,
@@ -80,7 +82,7 @@ const assetsService = new AssetsService({
   configProvider,
   state,
   tokenMetadataService,
-  cache,
+  cache: inMemoryCache,
 });
 
 const transactionsService = new TransactionsService({
@@ -128,7 +130,7 @@ const snapContext: SnapExecutionContext = {
   keyring,
   priceApiClient,
   state,
-  cache,
+  cache: stateCache,
   /* Services */
   assetsService,
   tokenPricesService,