perf: memoize blobs requests in the request scope

Author: pieh

src/run/handlers/cache.cts

@@ -31,22 +31,18 @@ import { getTracer } from './tracer.cjs'
 
 type TagManifest = { revalidatedAt: number }
 
-type TagManifestBlobCache = Record<string, Promise<TagManifest | null>>
-
 const purgeCacheUserAgent = `${nextRuntimePkgName}@${nextRuntimePkgVersion}`
 
 export class NetlifyCacheHandler implements CacheHandlerForMultipleVersions {
   options: CacheHandlerContext
   revalidatedTags: string[]
   cacheStore: MemoizedKeyValueStoreBackedByRegionalBlobStore
   tracer = getTracer()
-  tagManifestsFetchedFromBlobStoreInCurrentRequest: TagManifestBlobCache
 
   constructor(options: CacheHandlerContext) {
     this.options = options
     this.revalidatedTags = options.revalidatedTags
     this.cacheStore = getMemoizedKeyValueStoreBackedByRegionalBlobStore({ consistency: 'strong' })
-    this.tagManifestsFetchedFromBlobStoreInCurrentRequest = {}
   }
 
   private getTTL(blob: NetlifyCacheHandlerValue) {
@@ -478,7 +474,7 @@ export class NetlifyCacheHandler implements CacheHandlerForMultipleVersions {
   }
 
   resetRequestCache() {
-    this.tagManifestsFetchedFromBlobStoreInCurrentRequest = {}
+    // no-op
   }
 
   /**
@@ -530,14 +526,10 @@ export class NetlifyCacheHandler implements CacheHandlerForMultipleVersions {
       const tagManifestPromises: Promise<boolean>[] = []
 
       for (const tag of cacheTags) {
-        let tagManifestPromise: Promise<TagManifest | null> =
-          this.tagManifestsFetchedFromBlobStoreInCurrentRequest[tag]
-
-        if (!tagManifestPromise) {
-          tagManifestPromise = this.cacheStore.get<TagManifest>(tag, 'tagManifest.get')
-
-          this.tagManifestsFetchedFromBlobStoreInCurrentRequest[tag] = tagManifestPromise
-        }
+        const tagManifestPromise: Promise<TagManifest | null> = this.cacheStore.get<TagManifest>(
+          tag,
+          'tagManifest.get',
+        )
 
         tagManifestPromises.push(
           tagManifestPromise.then((tagManifest) => {

src/run/regional-blob-store.cts

@@ -1,5 +1,6 @@
 import { getDeployStore, GetWithMetadataOptions, Store } from '@netlify/blobs'
 
+import { getRequestContext, RequestContext } from './handlers/request-context.cjs'
 import { getTracer } from './handlers/tracer.cjs'
 
 const FETCH_BEFORE_NEXT_PATCHED_IT = Symbol.for('nf-not-patched-fetch')
@@ -68,6 +69,44 @@ const encodeBlobKey = async (key: string) => {
   return await encodeBlobKeyImpl(key)
 }
 
+// hold in-memory cache tied to request in memory for as long as request context is alive
+const perRequestInMemoryStore = new WeakMap<
+  RequestContext,
+  Record<string, Promise<unknown> | unknown>
+>()
+
+const getRequestSpecificInMemoryCache = () => {
+  const requestContext = getRequestContext()
+  if (!requestContext) {
+    // Fallback to a no-op store if we can't find request context
+    return {
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      get<T>(key: string): Promise<T | null> | undefined {
+        return undefined
+      },
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      set(key: string, value: unknown) {
+        // no-op
+      },
+    }
+  }
+
+  let perRequestStore = perRequestInMemoryStore.get(requestContext)
+  if (!perRequestStore) {
+    perRequestStore = {}
+    perRequestInMemoryStore.set(requestContext, perRequestStore)
+  }
+
+  return {
+    get<T>(key: string): Promise<T | null> | undefined {
+      return perRequestStore[key] as Promise<T | null> | undefined
+    },
+    set(key: string, value: unknown) {
+      perRequestStore[key] = value
+    },
+  }
+}
+
 export const getMemoizedKeyValueStoreBackedByRegionalBlobStore = (
   args: GetWithMetadataOptions = {},
 ) => {
@@ -76,21 +115,32 @@ export const getMemoizedKeyValueStoreBackedByRegionalBlobStore = (
 
   return {
     async get<T>(key: string, otelSpanTitle: string): Promise<T | null> {
-      const blobKey = await encodeBlobKey(key)
+      const inMemoryCache = getRequestSpecificInMemoryCache()
 
-      return tracer.withActiveSpan(otelSpanTitle, async (span) => {
+      const memoizedValue = inMemoryCache.get<T>(key)
+      if (memoizedValue) {
+        return memoizedValue
+      }
+
+      const blobKey = await encodeBlobKey(key)
+      const getPromise = tracer.withActiveSpan(otelSpanTitle, async (span) => {
         span.setAttributes({ key, blobKey })
-        const blob = await store.get(key, { type: 'json' })
+        const blob = await store.get(blobKey, { type: 'json' })
         span.addEvent(blob ? 'Hit' : 'Miss')
         return blob
       })
+      inMemoryCache.set(key, getPromise)
+      return getPromise
     },
     async set(key: string, value: unknown, otelSpanTitle: string) {
-      const blobKey = await encodeBlobKey(key)
+      const inMemoryCache = getRequestSpecificInMemoryCache()
 
+      inMemoryCache.set(key, value)
+
+      const blobKey = await encodeBlobKey(key)
       return tracer.withActiveSpan(otelSpanTitle, async (span) => {
         span.setAttributes({ key, blobKey })
-        return await store.setJSON(key, value)
+        return await store.setJSON(blobKey, value)
       })
     },
   }