feat: add monitored cache wrapper and use

create an internal cache package with
using adapter pattern to wrap LRU or TTL
caches.

use it for refactoring existing TTL pool
cache and JWT/S3 creds LRU caches.

use it also for tenant JWKs and configs
as LRUs.

add metrics and a few panels to show

fix double destroy call in pool TTL cache

Signed-off-by: ferhat elmas <elmas.ferhat@gmail.com>

Author: ferhatelmas

monitoring/grafana/dashboards/storage-otel.json

@@ -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 = 2048
+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/jwks/manager.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'
 
@@ -113,12 +118,21 @@ function getJWTAlgorithms(jwks: JwksConfig | null) {
   return algorithms
 }
 
-const jwtCache = new LRUCache<string, { token: string; payload: JWTPayload }>({
+function getJWTCacheKey(token: string, secret: string, jwks?: { keys: JwksConfigKey[] } | null) {
+  return createHash('sha256')
+    .update(token)
+    .update('\0')
+    .update(secret)
+    .update('\0')
+    .update(JSON.stringify(jwks?.keys ?? null))
+    .digest('base64url')
+}
+
+const jwtCache = createLruCache<string, JWTPayload>(JWT_CACHE_NAME, {
   maxSize: 1024 * 1024 * 50, // 50MB
-  sizeCalculation: (value) => {
-    return objectSizeOf(value)
-  },
+  sizeCalculation: (value) => objectSizeOf(value),
   ttlResolution: 5000, // 5 seconds
+  purgeStaleIntervalMs: DEFAULT_CACHE_PURGE_STALE_INTERVAL_MS,
 })
 
 /**
@@ -133,13 +147,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 +159,9 @@ export async function verifyJWTWithCache(
       return payload
     }
 
-    jwtCache.set(
-      token,
-      { token, payload },
-      {
-        ttl: payload.exp * 1000 - Date.now(),
-      }
-    )
+    jwtCache.set(cacheKey, payload, {
+      ttl: payload.exp * 1000 - Date.now(),
+    })
     return payload
   } catch (e) {
     throw e

src/internal/auth/jwt.ts

@@ -0,0 +1,30 @@
+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>
+}

src/internal/cache/adapter.ts

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

src/internal/cache/index.ts

@@ -0,0 +1,93 @@
+import { LRUCache as BaseLruCache } from 'lru-cache'
+import { CacheLookupOptions, CacheLookupOutcome, OutcomeAwareCache } 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 OutcomeAwareCache<K, V, LruCacheSetOptions<K, V>>
+{
+  private readonly cache: BaseLruCache<K, V>
+
+  constructor(options: LruCacheOptions<K, V>) {
+    const { purgeStaleIntervalMs, ...cacheOptions } = options
+
+    this.cache = new BaseLruCache<K, V>({
+      ...cacheOptions,
+    })
+
+    if (purgeStaleIntervalMs) {
+      const timer = setInterval(() => {
+        this.cache.purgeStale()
+      }, purgeStaleIntervalMs)
+      timer.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()
+  }
+}
+
+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>
+): OutcomeAwareCache<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()
+    },
+  })
+}

src/internal/cache/lru.ts

@@ -0,0 +1,104 @@
+import {
+  cacheEntries,
+  cacheEvictionsTotal,
+  cacheRequestsTotal,
+  cacheSizeBytes,
+  isMetricEnabled,
+  meter,
+} from '@internal/monitoring/metrics'
+import { Attributes } from '@opentelemetry/api'
+import { CacheLookupOptions, OutcomeAwareCache } from './adapter'
+import { CacheName } from './names'
+
+type CacheDisposeHandler<K, V, R extends string> = (value: V, key: K, reason: R) => void
+
+type MonitorCacheOptions = {
+  purgeStale?: () => void
+}
+
+function cacheAttrs(
+  cache: CacheName,
+  attrs?: Record<string, string | number | boolean>
+): Attributes {
+  return attrs ? { cache, ...attrs } : { cache }
+}
+
+export function withCacheEvictionMetrics<K, V, R extends string>(
+  cacheName: CacheName,
+  dispose?: CacheDisposeHandler<K, V, R>
+): CacheDisposeHandler<K, V, R> {
+  return (value, key, reason) => {
+    if (reason === 'evict') {
+      cacheEvictionsTotal.add(1, cacheAttrs(cacheName))
+    }
+
+    dispose?.(value, key, reason)
+  }
+}
+
+class MonitoredCache<K, V, SetOptions = undefined> implements OutcomeAwareCache<K, V, SetOptions> {
+  constructor(
+    private readonly name: CacheName,
+    private readonly cache: OutcomeAwareCache<K, V, SetOptions>,
+    options?: MonitorCacheOptions
+  ) {
+    meter.addBatchObservableCallback(
+      (observer) => {
+        const cacheEntriesEnabled = isMetricEnabled('cache_entries')
+        const cacheSizeBytesEnabled = isMetricEnabled('cache_size_bytes')
+
+        if (!cacheEntriesEnabled && !cacheSizeBytesEnabled) {
+          return
+        }
+
+        options?.purgeStale?.()
+        const stats = this.cache.getStats()
+        const attrs = cacheAttrs(this.name)
+
+        if (cacheEntriesEnabled) {
+          observer.observe(cacheEntries, stats.entries, attrs)
+        }
+
+        if (cacheSizeBytesEnabled) {
+          observer.observe(cacheSizeBytes, stats.sizeBytes, attrs)
+        }
+      },
+      [cacheEntries, cacheSizeBytes]
+    )
+  }
+
+  get(key: K, options?: CacheLookupOptions): V | undefined {
+    if (options?.recordMetrics === false) {
+      return this.cache.get(key, options)
+    }
+
+    const { value, outcome } = this.cache.getWithOutcome(key)
+    cacheRequestsTotal.add(1, cacheAttrs(this.name, { outcome }))
+
+    return value
+  }
+
+  getWithOutcome(key: K) {
+    return this.cache.getWithOutcome(key)
+  }
+
+  set(key: K, value: V, options?: SetOptions): void {
+    this.cache.set(key, value, options)
+  }
+
+  delete(key: K): boolean {
+    return this.cache.delete(key)
+  }
+
+  getStats() {
+    return this.cache.getStats()
+  }
+}
+
+export function monitorCache<K, V, SetOptions = undefined>(
+  cacheName: CacheName,
+  cache: OutcomeAwareCache<K, V, SetOptions>,
+  options?: MonitorCacheOptions
+): OutcomeAwareCache<K, V, SetOptions> {
+  return new MonitoredCache(cacheName, cache, options)
+}

src/internal/cache/monitoring.ts

@@ -0,0 +1,10 @@
+export const JWT_CACHE_NAME = 'jwt' as const
+export const TENANT_CONFIG_CACHE_NAME = 'tenant_config' as const
+export const TENANT_JWKS_CACHE_NAME = 'tenant_jwks' as const
+export const TENANT_S3_CREDENTIALS_CACHE_NAME = 'tenant_s3_credentials' as const
+
+export type CacheName =
+  | typeof JWT_CACHE_NAME
+  | typeof TENANT_CONFIG_CACHE_NAME
+  | typeof TENANT_JWKS_CACHE_NAME
+  | typeof TENANT_S3_CREDENTIALS_CACHE_NAME