feat: add monitored cache wrapper and use

src/internal/auth/jwks/manager.ts

@@ -1,7 +1,13 @@
 import { decrypt, encrypt, generateHS512JWK } from '@internal/auth'
+import {
+  createLruCache,
+  DEFAULT_CACHE_PURGE_STALE_INTERVAL_MS,
+  TENANT_JWKS_CACHE_NAME,
+} from '@internal/cache'
 import { createMutexByKey } from '@internal/concurrency'
 import { PubSubAdapter } from '@internal/pubsub'
 import { Knex } from 'knex'
+import objectSizeOf from 'object-sizeof'
 import { JwksConfig, JwksConfigKeyOCT } from '../../../config'
 import { JWKSManagerStore } from './store'
 
@@ -10,7 +16,23 @@ const JWK_KIND_STORAGE_URL_SIGNING = 'storage-url-signing-key'
 const JWK_KID_SEPARATOR = '_'
 
 const tenantJwksMutex = createMutexByKey<JwksConfig>()
-const tenantJwksConfigCache = new Map<string, JwksConfig>()
+export const TENANT_JWKS_CACHE_MAX_ITEMS = 16384
+export const TENANT_JWKS_CACHE_MAX_SIZE_BYTES = 1024 * 1024 * 50 // 50 MiB
+export const TENANT_JWKS_CACHE_TTL_MS = 1000 * 60 * 60 // 1h
+
+const tenantJwksConfigCache = createLruCache<string, JwksConfig>(TENANT_JWKS_CACHE_NAME, {
+  max: TENANT_JWKS_CACHE_MAX_ITEMS,
+  maxSize: TENANT_JWKS_CACHE_MAX_SIZE_BYTES,
+  ttl: TENANT_JWKS_CACHE_TTL_MS,
+  sizeCalculation: (value) => objectSizeOf(value),
+  updateAgeOnGet: true,
+  allowStale: false,
+  purgeStaleIntervalMs: DEFAULT_CACHE_PURGE_STALE_INTERVAL_MS,
+})
+
+export function deleteTenantJwksConfig(tenantId: string): void {
+  tenantJwksConfigCache.delete(tenantId)
+}
 
 function createJwkKid({ kind, id }: { id: string; kind: string }): string {
   return kind + JWK_KID_SEPARATOR + id
@@ -72,14 +94,14 @@ export class JWKSManager {
   async getJwksTenantConfig(tenantId: string): Promise<JwksConfig> {
     const cachedJwks = tenantJwksConfigCache.get(tenantId)
 
-    if (cachedJwks) {
+    if (cachedJwks !== undefined) {
       return cachedJwks
     }
 
     return tenantJwksMutex(tenantId, async () => {
-      const cachedJwks = tenantJwksConfigCache.get(tenantId)
+      const cachedJwks = tenantJwksConfigCache.get(tenantId, { recordMetrics: false })
 
-      if (cachedJwks) {
+      if (cachedJwks !== undefined) {
         return cachedJwks
       }
 

src/internal/auth/jwt.ts

@@ -1,3 +1,9 @@
+import { createHash } from 'node:crypto'
+import {
+  createLruCache,
+  DEFAULT_CACHE_PURGE_STALE_INTERVAL_MS,
+  JWT_CACHE_NAME,
+} from '@internal/cache'
 import { ERRORS } from '@internal/errors'
 import {
   exportJWK,
@@ -9,7 +15,6 @@ import {
   jwtVerify,
   SignJWT,
 } from 'jose'
-import { LRUCache } from 'lru-cache'
 import objectSizeOf from 'object-sizeof'
 import { getConfig, JwksConfig, JwksConfigKey, JwksConfigKeyOCT } from '../../config'
 
@@ -33,6 +38,8 @@ export type SignedUploadToken = {
   exp: number
 }
 
+const jwtJwksFingerprintCache = new WeakMap<object, string>()
+
 async function findJWKFromHeader(
   header: JWTHeaderParameters,
   secret: string,
@@ -113,12 +120,46 @@ function getJWTAlgorithms(jwks: JwksConfig | null) {
   return algorithms
 }
 
-const jwtCache = new LRUCache<string, { token: string; payload: JWTPayload }>({
-  maxSize: 1024 * 1024 * 50, // 50MB
-  sizeCalculation: (value) => {
-    return objectSizeOf(value)
-  },
-  ttlResolution: 5000, // 5 seconds
+function getJWTJwksFingerprint(jwks?: { keys: JwksConfigKey[] } | null): string {
+  if (!jwks) {
+    return 'null'
+  }
+
+  const cachedFingerprint = jwtJwksFingerprintCache.get(jwks)
+  if (cachedFingerprint) {
+    return cachedFingerprint
+  }
+
+  const fingerprint = createHash('sha256')
+    .update(JSON.stringify(jwks.keys ?? null))
+    .digest('base64url')
+  jwtJwksFingerprintCache.set(jwks, fingerprint)
+  return fingerprint
+}
+
+function getJWTCacheKey(token: string, secret: string, jwks?: { keys: JwksConfigKey[] } | null) {
+  const hash = createHash('sha256')
+    .update(token)
+    .update('\0')
+    .update(secret)
+    .update('\0')
+    .update(getJWTJwksFingerprint(jwks))
+
+  return hash.digest('base64url')
+}
+
+// JWT payloads are comparatively small and high-churn, so keep a higher
+// cardinality guardrail than the longer-lived config-style caches.
+export const JWT_CACHE_MAX_ITEMS = 65536
+export const JWT_CACHE_MAX_SIZE_BYTES = 1024 * 1024 * 50 // 50 MiB
+export const JWT_CACHE_TTL_RESOLUTION_MS = 5000 // 5 seconds
+
+const jwtCache = createLruCache<string, JWTPayload>(JWT_CACHE_NAME, {
+  max: JWT_CACHE_MAX_ITEMS,
+  maxSize: JWT_CACHE_MAX_SIZE_BYTES,
+  sizeCalculation: (value) => objectSizeOf(value),
+  ttlResolution: JWT_CACHE_TTL_RESOLUTION_MS,
+  purgeStaleIntervalMs: DEFAULT_CACHE_PURGE_STALE_INTERVAL_MS,
 })
 
 /**
@@ -133,13 +174,10 @@ export async function verifyJWTWithCache(
   secret: string,
   jwks?: { keys: JwksConfigKey[] } | null
 ) {
-  const cachedVerification = jwtCache.get(token)
-  if (
-    cachedVerification &&
-    cachedVerification.payload.exp &&
-    cachedVerification.payload.exp * 1000 > Date.now()
-  ) {
-    return Promise.resolve(cachedVerification.payload)
+  const cacheKey = getJWTCacheKey(token, secret, jwks)
+  const cachedPayload = jwtCache.get(cacheKey)
+  if (cachedPayload && cachedPayload.exp && cachedPayload.exp * 1000 > Date.now()) {
+    return Promise.resolve(cachedPayload)
   }
 
   try {
@@ -148,13 +186,10 @@ export async function verifyJWTWithCache(
       return payload
     }
 
-    jwtCache.set(
-      token,
-      { token, payload },
-      {
-        ttl: payload.exp * 1000 - Date.now(),
-      }
-    )
+    const ttl = payload.exp * 1000 - Date.now()
+    if (ttl > 0) {
+      jwtCache.set(cacheKey, payload, { ttl })
+    }
     return payload
   } catch (e) {
     throw e

src/internal/cache/adapter.ts

@@ -0,0 +1,38 @@
+export type CacheLookupOptions = {
+  recordMetrics?: boolean
+}
+
+export type CacheLookupOutcome = 'hit' | 'miss' | 'stale'
+
+export type CacheLookupResult<V> = {
+  value: V | undefined
+  outcome: CacheLookupOutcome
+}
+
+export type CacheStats = {
+  entries: number
+  sizeBytes: number
+}
+
+export interface Cache<K, V, SetOptions = undefined> {
+  get(key: K, options?: CacheLookupOptions): V | undefined
+  set(key: K, value: V, options?: SetOptions): void
+  delete(key: K): boolean
+}
+
+export interface InspectableCache<K, V, SetOptions = undefined> extends Cache<K, V, SetOptions> {
+  getStats(): CacheStats
+}
+
+export interface OutcomeAwareCache<K, V, SetOptions = undefined>
+  extends InspectableCache<K, V, SetOptions> {
+  getWithOutcome(key: K): CacheLookupResult<V>
+}
+
+export interface Disposable {
+  dispose(): void
+}
+
+export interface DisposableCache<K, V, SetOptions = undefined>
+  extends OutcomeAwareCache<K, V, SetOptions>,
+    Disposable {}

src/internal/cache/index.ts

@@ -0,0 +1,4 @@
+export * from './adapter'
+export * from './lru'
+export * from './names'
+export * from './ttl'

src/internal/cache/lru.ts

@@ -0,0 +1,100 @@
+import { LRUCache as BaseLruCache } from 'lru-cache'
+import { CacheLookupOptions, CacheLookupOutcome, DisposableCache } from './adapter'
+import { monitorCache, withCacheEvictionMetrics } from './monitoring'
+import { CacheName } from './names'
+
+export type LruCacheSetOptions<K extends {}, V extends {}> = BaseLruCache.SetOptions<K, V, unknown>
+
+export type LruCacheOptions<K extends {}, V extends {}> = BaseLruCache.Options<K, V, unknown> & {
+  purgeStaleIntervalMs?: number
+}
+
+export const DEFAULT_CACHE_PURGE_STALE_INTERVAL_MS = 1000 * 60 // 1 minute
+
+export class LruCache<K extends {}, V extends {}>
+  implements DisposableCache<K, V, LruCacheSetOptions<K, V>>
+{
+  private readonly cache: BaseLruCache<K, V>
+  private readonly purgeStaleTimer?: ReturnType<typeof setInterval>
+
+  constructor(options: LruCacheOptions<K, V>) {
+    const { purgeStaleIntervalMs, ...cacheOptions } = options
+
+    this.cache = new BaseLruCache<K, V>({
+      ...cacheOptions,
+    })
+
+    if (purgeStaleIntervalMs) {
+      this.purgeStaleTimer = setInterval(() => {
+        this.cache.purgeStale()
+      }, purgeStaleIntervalMs)
+      this.purgeStaleTimer.unref?.()
+    }
+  }
+
+  get(key: K, options?: CacheLookupOptions): V | undefined {
+    return this.getWithOutcome(key).value
+  }
+
+  getWithOutcome(key: K) {
+    const status: BaseLruCache.Status<V> = {}
+    const value = this.cache.get(key, { status })
+    const outcome = (status.get || (value === undefined ? 'miss' : 'hit')) as CacheLookupOutcome
+
+    return { value, outcome }
+  }
+
+  set(key: K, value: V, options?: LruCacheSetOptions<K, V>): void {
+    this.cache.set(key, value, options)
+  }
+
+  delete(key: K): boolean {
+    return this.cache.delete(key)
+  }
+
+  getStats() {
+    return {
+      entries: this.cache.size,
+      sizeBytes: this.cache.calculatedSize,
+    }
+  }
+
+  purgeStale(): boolean {
+    return this.cache.purgeStale()
+  }
+
+  dispose(): void {
+    if (this.purgeStaleTimer) {
+      clearInterval(this.purgeStaleTimer)
+    }
+  }
+}
+
+export function createLruCache<K extends {}, V extends {}>(
+  options: LruCacheOptions<K, V>
+): LruCache<K, V>
+export function createLruCache<K extends {}, V extends {}>(
+  name: CacheName,
+  options: LruCacheOptions<K, V>
+): DisposableCache<K, V, LruCacheSetOptions<K, V>>
+export function createLruCache<K extends {}, V extends {}>(
+  nameOrOptions: CacheName | LruCacheOptions<K, V>,
+  maybeOptions?: LruCacheOptions<K, V>
+) {
+  if (typeof nameOrOptions !== 'string') {
+    return new LruCache(nameOrOptions)
+  }
+
+  const cacheName = nameOrOptions
+  const options = maybeOptions as LruCacheOptions<K, V>
+  const cache = new LruCache<K, V>({
+    ...options,
+    disposeAfter: withCacheEvictionMetrics(cacheName, options.disposeAfter),
+  })
+
+  return monitorCache(cacheName, cache, {
+    purgeStale: () => {
+      cache.purgeStale()
+    },
+  })
+}