Multi-Tier caching in the tokens service

packages/services/api/src/shared/entities.ts

@@ -318,7 +318,7 @@ export interface Token {
   project: string;
   organization: string;
   date: string;
-  lastUsedAt: string;
+  lastUsedAt: string | null;
   scopes: readonly string[];
 }
 

packages/services/storage/src/tokens.ts

@@ -2,6 +2,20 @@ import { Interceptor, sql } from 'slonik';
 import { getPool, toDate, tokens } from './db';
 import type { Slonik } from './shared';
 
+function transformToken(row: tokens) {
+  return {
+    token: row.token,
+    tokenAlias: row.token_alias,
+    name: row.name,
+    date: row.created_at as unknown as string,
+    lastUsedAt: row.last_used_at as unknown as string | null,
+    organization: row.organization_id,
+    project: row.project_id,
+    target: row.target_id,
+    scopes: row.scopes || [],
+  };
+}
+
 export async function createTokenStorage(
   connection: string,
   maximumPoolSize: number,
@@ -33,19 +47,21 @@ export async function createTokenStorage(
         `,
       );
 
-      return result.rows;
+      return result.rows.map(transformToken);
     },
     async getToken({ token }: { token: string }) {
-      return pool.maybeOne<Slonik<tokens>>(
+      const row = await pool.maybeOne<Slonik<tokens>>(
         sql`
           SELECT *
           FROM tokens
           WHERE token = ${token} AND deleted_at IS NULL
           LIMIT 1
         `,
       );
+
+      return row ? transformToken(row) : null;
     },
-    createToken({
+    async createToken({
       token,
       tokenAlias,
       target,
@@ -62,7 +78,7 @@ export async function createTokenStorage(
       organization: string;
       scopes: readonly string[];
     }) {
-      return pool.one<Slonik<tokens>>(
+      const row = await pool.one<Slonik<tokens>>(
         sql`
           INSERT INTO tokens
             (name, token, token_alias, target_id, project_id, organization_id, scopes)
@@ -73,14 +89,20 @@ export async function createTokenStorage(
             )})
           RETURNING *
         `,
-      );
+      )
+
+      return transformToken(row);
     },
-    async deleteToken({ token }: { token: string }) {
-      await pool.query(
-        sql`
-          UPDATE tokens SET deleted_at = NOW() WHERE token = ${token}
-        `,
-      );
+    async deleteToken(params: { token: string; postDeletionTransaction: () => Promise<void> }) {
+      await pool.transaction(async t => {
+        await t.query(sql`
+          UPDATE tokens
+          SET deleted_at = NOW()
+          WHERE token = ${params.token}
+        `);
+
+        await params.postDeletionTransaction();
+      });
     },
     async touchTokens({ tokens }: { tokens: Array<{ token: string; date: Date }> }) {
       await pool.query(sql`

packages/services/tokens/.env.template

@@ -8,3 +8,4 @@ REDIS_PORT="6379"
 REDIS_PASSWORD=""
 PORT=6001
 OPENTELEMETRY_COLLECTOR_ENDPOINT="<sync>"
+LOG_LEVEL="debug"

packages/services/tokens/package.json

@@ -17,6 +17,7 @@
     "dotenv": "16.4.5",
     "fastify": "4.28.1",
     "ioredis": "5.4.1",
+    "lru-cache": "11.0.2",
     "ms": "2.1.3",
     "p-timeout": "6.1.3",
     "pino-pretty": "11.3.0",

packages/services/tokens/src/api.ts

@@ -5,8 +5,8 @@ import { z } from 'zod';
 import { createErrorHandler, handleTRPCError, maskToken, metrics } from '@hive/service-common';
 import type { inferRouterInputs, inferRouterOutputs } from '@trpc/server';
 import { initTRPC } from '@trpc/server';
-import { useCache } from './cache';
-import { cacheHits, cacheMisses } from './metrics';
+import { recordTokenRead } from './metrics';
+import { Storage } from './multi-tier-storage';
 
 const httpRequests = new metrics.Counter({
   name: 'tokens_http_requests',
@@ -43,7 +43,7 @@ function generateToken() {
 export type Context = {
   req: FastifyRequest;
   errorHandler: ReturnType<typeof createErrorHandler>;
-  getStorage: ReturnType<typeof useCache>['getStorage'];
+  storage: Storage;
   tokenReadFailuresCache: LruType<string>;
 };
 
@@ -72,9 +72,7 @@ export const tokensApiRouter = t.router({
     )
     .query(async ({ ctx, input }) => {
       try {
-        const storage = await ctx.getStorage();
-
-        return await storage.readTarget(input.targetId);
+        return await ctx.storage.readTarget(input.targetId);
       } catch (error) {
         ctx.errorHandler('Failed to get tokens of a target', error as Error);
 
@@ -91,8 +89,7 @@ export const tokensApiRouter = t.router({
     )
     .mutation(async ({ ctx, input }) => {
       try {
-        const storage = await ctx.getStorage();
-        await storage.invalidateTokens(input.tokens);
+        await ctx.storage.invalidateTokens(input.tokens);
 
         return true;
       } catch (error) {
@@ -116,9 +113,8 @@ export const tokensApiRouter = t.router({
     .mutation(async ({ ctx, input }) => {
       try {
         const { target, project, organization, name, scopes } = input;
-        const storage = await ctx.getStorage();
         const token = generateToken();
-        const result = await storage.writeToken({
+        const result = await ctx.storage.writeToken({
           name,
           target,
           project,
@@ -149,8 +145,7 @@ export const tokensApiRouter = t.router({
     .mutation(async ({ ctx, input }) => {
       try {
         const hashed_token = input.token;
-        const storage = await ctx.getStorage();
-        await storage.deleteToken(hashed_token);
+        await ctx.storage.deleteToken(hashed_token);
 
         return true;
       } catch (error) {
@@ -175,14 +170,12 @@ export const tokensApiRouter = t.router({
       const cachedFailure = ctx.tokenReadFailuresCache.get(hash);
 
       if (cachedFailure) {
-        cacheHits.inc(1);
         throw new Error(cachedFailure);
       }
 
       try {
-        const storage = await ctx.getStorage();
-        const result = await storage.readToken(hash);
-
+        const result = await ctx.storage.readToken(hash, alias);
+        recordTokenRead(result ? 200 : 404);
         // removes the token from the failures cache (in case the value expired)
         ctx.tokenReadFailuresCache.delete(hash);
 
@@ -193,8 +186,8 @@ export const tokensApiRouter = t.router({
         // set token read as failure
         // so we don't try to read it again for next X minutes
         ctx.tokenReadFailuresCache.set(hash, (error as Error).message);
-        cacheMisses.inc(1);
 
+        recordTokenRead(500);
         throw error;
       }
     }),