Move from tiny-lru to lru-cache in the usage service (#6362)

packages/services/usage/package.json

@@ -23,7 +23,6 @@
     "kafkajs": "2.2.4",
     "lru-cache": "11.0.2",
     "pino-pretty": "11.3.0",
-    "tiny-lru": "8.0.2",
     "zod": "3.24.1"
   }
 }

packages/services/usage/src/metrics.ts

@@ -1,10 +1,5 @@
 import { metrics } from '@hive/service-common';
 
-export const tokenCacheHits = new metrics.Counter({
-  name: 'usage_tokens_cache_hits',
-  help: 'Number of cache hits',
-});
-
 export const tokenRequests = new metrics.Counter({
   name: 'usage_tokens_requests',
   help: 'Number of requests to Tokens service',

packages/services/usage/src/tokens.ts

@@ -1,8 +1,8 @@
-import LRU from 'tiny-lru';
+import { LRUCache } from 'lru-cache';
 import { ServiceLogger } from '@hive/service-common';
 import type { TokensApi } from '@hive/tokens';
 import { createTRPCProxyClient, httpLink } from '@trpc/client';
-import { tokenCacheHits, tokenRequests } from './metrics';
+import { tokenRequests } from './metrics';
 
 export enum TokenStatus {
   NotFound,
@@ -20,7 +20,6 @@ type Token = TokensResponse | TokenStatus;
 
 export function createTokens(config: { endpoint: string; logger: ServiceLogger }) {
   const endpoint = config.endpoint.replace(/\/$/, '');
-  const tokens = LRU<Promise<Token>>(1000, 30_000);
   const tokensApi = createTRPCProxyClient<TokensApi>({
     links: [
       httpLink({
@@ -32,44 +31,43 @@ export function createTokens(config: { endpoint: string; logger: ServiceLogger }
       }),
     ],
   });
-  async function fetchFreshToken(token: string) {
-    try {
-      const info = await tokensApi.getToken.query({
-        token,
-      });
+  const tokens = new LRUCache<string, Token>({
+    max: 1000,
+    ttl: 30_000,
+    allowStale: false,
+    // If a cache entry is stale or missing, this method is called
+    // to fill the cache with fresh data.
+    // This method is called only once per cache key,
+    // even if multiple requests are waiting for it.
+    async fetchMethod(token) {
+      tokenRequests.inc();
+      try {
+        const info = await tokensApi.getToken.query({
+          token,
+        });
 
-      if (info) {
-        const result = info.scopes.includes('target:registry:write')
-          ? {
-              target: info.target,
-              project: info.project,
-              organization: info.organization,
-              scopes: info.scopes,
-            }
-          : TokenStatus.NoAccess;
-        return result;
+        if (info) {
+          const result = info.scopes.includes('target:registry:write')
+            ? {
+                target: info.target,
+                project: info.project,
+                organization: info.organization,
+                scopes: info.scopes,
+              }
+            : TokenStatus.NoAccess;
+          return result;
+        }
+        return TokenStatus.NotFound;
+      } catch (error) {
+        config.logger.error('Failed to fetch fresh token', error);
+        return TokenStatus.NotFound;
       }
-      return TokenStatus.NotFound;
-    } catch (error) {
-      config.logger.error('Failed to fetch fresh token', error);
-      return TokenStatus.NotFound;
-    }
-  }
+    },
+  });
 
   return {
     async fetch(token: string) {
-      tokenRequests.inc();
-      const tokenInfo = await tokens.get(token);
-
-      if (!tokenInfo) {
-        const result = fetchFreshToken(token);
-        tokens.set(token, result);
-        return result;
-      }
-
-      tokenCacheHits.inc();
-
-      return tokenInfo ?? TokenStatus.NotFound;
+      return (await tokens.fetch(token)) ?? TokenStatus.NotFound;
     },
     isNotFound(token: Token): token is TokenStatus.NotFound {
       return token === TokenStatus.NotFound;

packages/services/usage/src/usage-processor-1.ts

@@ -2,7 +2,7 @@ import { createHash, randomUUID } from 'node:crypto';
 import type { JSONSchemaType } from 'ajv';
 import Ajv from 'ajv';
 import { parse } from 'graphql';
-import LRU from 'tiny-lru';
+import { LRUCache } from 'lru-cache';
 import type { ServiceLogger as Logger } from '@hive/service-common';
 import { RawReport } from '@hive/usage-common';
 import {
@@ -208,7 +208,9 @@ const ajv = new Ajv({
   },
 });
 
-const validOperationBodyCache = LRU<boolean>(5000, 300_000 /* 5 minutes */);
+const validOperationBodyCache = new LRUCache<string, boolean>({
+  max: 20_000,
+});
 
 const operationMapRecordSchema: JSONSchemaType<OperationMapRecord> = {
   type: 'object',
@@ -270,7 +272,8 @@ export function validateOperationMapRecord(record: OperationMapRecord) {
 }
 
 export function isValidOperationBody(operation: string) {
-  const cached = validOperationBodyCache.get(operation);
+  const operationHash = createHash('sha256').update(operation).digest('hex');
+  const cached = validOperationBodyCache.get(operationHash);
 
   if (typeof cached === 'boolean') {
     return cached;
@@ -280,10 +283,10 @@ export function isValidOperationBody(operation: string) {
     parse(operation, {
       noLocation: true,
     });
-    validOperationBodyCache.set(operation, true);
+    validOperationBodyCache.set(operationHash, true);
     return true;
   } catch (error) {
-    validOperationBodyCache.set(operation, false);
+    validOperationBodyCache.set(operationHash, false);
     return false;
   }
 }