perf: reuse lastModified captured in cache handler to calculate date header (#260)

* perf: reuse lastModified captured in cache handler to calculate date header

* chore: record exception and set warning and severity attributes

* test: spy on blobs.get usage and assert we only do it once, compare produced date headers

* test: add integration test for request-context not leaking for concurrent requests

* chore: remove some commented out code

* fix: there is no requestContext on CacheHandler anymore

Author: pieh

package-lock.json

@@ -69,6 +69,7 @@
     "get-port": "^7.0.0",
     "lambda-local": "^2.1.2",
     "memfs": "^4.6.0",
+    "mock-require": "^3.0.3",
     "msw": "^2.0.7",
     "next": "^14.0.4",
     "os": "^0.1.2",

package.json

@@ -5,7 +5,7 @@ import { Buffer } from 'node:buffer'
 
 import { getDeployStore, Store } from '@netlify/blobs'
 import { purgeCache } from '@netlify/functions'
-import { trace } from '@opentelemetry/api'
+import { trace, type Span } from '@opentelemetry/api'
 import { NEXT_CACHE_TAGS_HEADER } from 'next/dist/lib/constants.js'
 import type {
   CacheHandler,
@@ -14,6 +14,8 @@ import type {
   IncrementalCache,
 } from 'next/dist/server/lib/incremental-cache/index.js'
 
+import { getRequestContext } from './request-context.cjs'
+
 type TagManifest = { revalidatedAt: number }
 
 const fetchBeforeNextPatchedIt = globalThis.fetch
@@ -35,6 +37,54 @@ export class NetlifyCacheHandler implements CacheHandler {
     return await encodeBlobKey(key)
   }
 
+  private captureResponseCacheLastModified(
+    cacheValue: CacheHandlerValue,
+    key: string,
+    getCacheKeySpan: Span,
+  ) {
+    if (cacheValue.value?.kind === 'FETCH') {
+      return
+    }
+
+    const requestContext = getRequestContext()
+
+    if (!requestContext) {
+      // we will not be able to use request context for date header calculation
+      // we will fallback to using blobs
+      getCacheKeySpan.recordException(
+        new Error('CacheHandler was called without a request context'),
+      )
+      getCacheKeySpan.setAttributes({
+        severity: 'alert',
+        warning: true,
+      })
+      return
+    }
+
+    if (requestContext.responseCacheKey && requestContext.responseCacheKey !== key) {
+      // if there are multiple response-cache keys, we don't know which one we should use
+      // so as a safety measure we will not use any of them and let blobs be used
+      // to calculate the date header
+      requestContext.responseCacheGetLastModified = undefined
+      getCacheKeySpan.recordException(
+        new Error(
+          `Multiple response cache keys used in single request: ["${requestContext.responseCacheKey}, "${key}"]`,
+        ),
+      )
+      getCacheKeySpan.setAttributes({
+        severity: 'alert',
+        warning: true,
+      })
+      return
+    }
+
+    requestContext.responseCacheKey = key
+    if (cacheValue.lastModified) {
+      // we store it to use it later when calculating date header
+      requestContext.responseCacheGetLastModified = cacheValue.lastModified
+    }
+  }
+
   async get(...args: Parameters<CacheHandler['get']>): ReturnType<CacheHandler['get']> {
     return this.tracer.startActiveSpan('get cache key', async (span) => {
       const [key, ctx = {}] = args
@@ -61,6 +111,8 @@ export class NetlifyCacheHandler implements CacheHandler {
         return null
       }
 
+      this.captureResponseCacheLastModified(blob, key, span)
+
       switch (blob.value?.kind) {
         case 'FETCH':
           span.addEvent('FETCH', { lastModified: blob.lastModified, revalidate: ctx.revalidate })

src/run/handlers/cache.cts

@@ -0,0 +1,43 @@
+import { AsyncLocalStorage } from 'node:async_hooks'
+
+export type RequestContext = {
+  responseCacheGetLastModified?: number
+  responseCacheKey?: string
+}
+
+type RequestContextAsyncLocalStorage = AsyncLocalStorage<RequestContext>
+
+export function createRequestContext(): RequestContext {
+  return {}
+}
+
+const REQUEST_CONTEXT_GLOBAL_KEY = Symbol.for('nf-request-context-async-local-storage')
+let requestContextAsyncLocalStorage: RequestContextAsyncLocalStorage | undefined
+function getRequestContextAsyncLocalStorage() {
+  if (requestContextAsyncLocalStorage) {
+    return requestContextAsyncLocalStorage
+  }
+  // for cases when there is multiple "copies" of this module, we can't just init
+  // AsyncLocalStorage in the module scope, because it will be different for each
+  // copy - so first time an instance of this module is used, we store AsyncLocalStorage
+  // in global scope and reuse it for all subsequent calls
+  const extendedGlobalThis = globalThis as typeof globalThis & {
+    [REQUEST_CONTEXT_GLOBAL_KEY]?: RequestContextAsyncLocalStorage
+  }
+  if (extendedGlobalThis[REQUEST_CONTEXT_GLOBAL_KEY]) {
+    return extendedGlobalThis[REQUEST_CONTEXT_GLOBAL_KEY]
+  }
+
+  const storage = new AsyncLocalStorage<RequestContext>()
+  // store for future use of this instance of module
+  requestContextAsyncLocalStorage = storage
+  // store for future use of copy of this module
+  extendedGlobalThis[REQUEST_CONTEXT_GLOBAL_KEY] = storage
+  return storage
+}
+
+export const getRequestContext = () => getRequestContextAsyncLocalStorage().getStore()
+
+export function runWithRequestContext<T>(requestContext: RequestContext, fn: () => T): T {
+  return getRequestContextAsyncLocalStorage().run(requestContext, fn)
+}

src/run/handlers/request-context.cts

@@ -14,6 +14,8 @@ import {
 import { nextResponseProxy } from '../revalidate.js'
 import { logger } from '../systemlog.js'
 
+import { createRequestContext, runWithRequestContext } from './request-context.cjs'
+
 let nextHandler: WorkerRequestHandler, nextConfig: NextConfigComplete, tagsManifest: TagsManifest
 
 export default async (request: Request) => {
@@ -51,26 +53,34 @@ export default async (request: Request) => {
 
     const resProxy = nextResponseProxy(res)
 
+    const requestContext = createRequestContext()
+
     // We don't await this here, because it won't resolve until the response is finished.
-    const nextHandlerPromise = nextHandler(req, resProxy).catch((error) => {
-      logger.withError(error).error('next handler error')
-      console.error(error)
-      resProxy.statusCode = 500
-      span.recordException(error)
-      span.setAttribute('http.status_code', 500)
-      span.setStatus({
-        code: SpanStatusCode.ERROR,
-        message: error instanceof Error ? error.message : String(error),
-      })
-      span.end()
-      resProxy.end('Internal Server Error')
-    })
+    const nextHandlerPromise = runWithRequestContext(requestContext, () =>
+      nextHandler(req, resProxy).catch((error) => {
+        logger.withError(error).error('next handler error')
+        console.error(error)
+        resProxy.statusCode = 500
+        span.recordException(error)
+        span.setAttribute('http.status_code', 500)
+        span.setStatus({
+          code: SpanStatusCode.ERROR,
+          message: error instanceof Error ? error.message : String(error),
+        })
+        span.end()
+        resProxy.end('Internal Server Error')
+      }),
+    )
 
     // Contrary to the docs, this resolves when the headers are available, not when the stream closes.
     // See https://github.com/fastly/http-compute-js/blob/main/src/http-compute-js/http-server.ts#L168-L173
     const response = await toComputeResponse(resProxy)
 
-    await adjustDateHeader(response.headers, request, span, tracer)
+    if (requestContext.responseCacheKey) {
+      span.setAttribute('responseCacheKey', requestContext.responseCacheKey)
+    }
+
+    await adjustDateHeader({ headers: response.headers, request, span, tracer, requestContext })
 
     setCacheControlHeaders(response.headers, request)
     setCacheTagsHeaders(response.headers, request, tagsManifest)

src/run/handlers/server.ts

@@ -5,6 +5,7 @@ import type { NextConfigComplete } from 'next/dist/server/config-shared.js'
 import { encodeBlobKey } from '../shared/blobkey.js'
 
 import type { TagsManifest } from './config.js'
+import type { RequestContext } from './handlers/request-context.cjs'
 
 interface NetlifyVaryValues {
   headers: string[]
@@ -84,12 +85,19 @@ const fetchBeforeNextPatchedIt = globalThis.fetch
  * By default, Next.js sets the date header to the current time, even if it's served
  * from the cache, meaning that the CDN will cache it for 10 seconds, which is incorrect.
  */
-export const adjustDateHeader = async (
-  headers: Headers,
-  request: Request,
-  span: Span,
-  tracer: Tracer,
-) => {
+export const adjustDateHeader = async ({
+  headers,
+  request,
+  span,
+  tracer,
+  requestContext,
+}: {
+  headers: Headers
+  request: Request
+  span: Span
+  tracer: Tracer
+  requestContext: RequestContext
+}) => {
   const cacheState = headers.get('x-nextjs-cache')
   const isServedFromCache = cacheState === 'HIT' || cacheState === 'STALE'
 
@@ -102,28 +110,72 @@ export const adjustDateHeader = async (
     return
   }
   const key = new URL(request.url).pathname
-  const blobKey = await encodeBlobKey(key)
-  const blobStore = getDeployStore({ fetch: fetchBeforeNextPatchedIt, consistency: 'strong' })
-
-  // TODO: use metadata for this
-  const { lastModified } = await tracer.startActiveSpan(
-    'get cache to calculate date header',
-    async (getBlobForDateSpan) => {
-      getBlobForDateSpan.setAttributes({
-        key,
-        blobKey,
-      })
-      const blob = (await blobStore.get(blobKey, { type: 'json' })) ?? {}
-
-      getBlobForDateSpan.addEvent(blob ? 'Cache hit' : 'Cache miss')
-      getBlobForDateSpan.end()
-      return blob
-    },
-  )
+
+  let lastModified: number | undefined
+  if (requestContext.responseCacheGetLastModified) {
+    lastModified = requestContext.responseCacheGetLastModified
+  } else {
+    // this is fallback in case requestContext doesn't contain lastModified
+    // using fallback indicate problem in the setup as assumption is that for cached responses
+    // request context would contain lastModified value
+    // this is not fatal as we have fallback,
+    // but we want to know about it happening
+    span.recordException(
+      new Error('lastModified not found in requestContext, falling back to trying blobs'),
+    )
+    span.setAttributes({
+      severity: 'alert',
+      warning: true,
+    })
+
+    const blobKey = await encodeBlobKey(key)
+    const blobStore = getDeployStore({ fetch: fetchBeforeNextPatchedIt, consistency: 'strong' })
+
+    // TODO: use metadata for this
+    lastModified = await tracer.startActiveSpan(
+      'get cache to calculate date header',
+      async (getBlobForDateSpan) => {
+        getBlobForDateSpan.setAttributes({
+          key,
+          blobKey,
+        })
+        const blob = (await blobStore.get(blobKey, { type: 'json' })) ?? {}
+
+        getBlobForDateSpan.addEvent(blob ? 'Cache hit' : 'Cache miss')
+        getBlobForDateSpan.end()
+        return blob.lastModified
+      },
+    )
+  }
 
   if (!lastModified) {
+    // this should never happen as we only execute this code path for cached responses
+    // and those should always have lastModified value
+    span.recordException(
+      new Error(
+        'lastModified not found in either requestContext or blobs, date header for cached response is not set',
+      ),
+    )
+    span.setAttributes({
+      severity: 'alert',
+      warning: true,
+    })
     return
   }
+
+  if (lastModified === 1) {
+    // epoch 1 time seems to not work well with the CDN
+    // it causes to produce "Thu, 01 Jan 1970 00:00:00 GMT" date
+    // header correctly but CDN calculates then AGE to be 0
+    // causing response to be cached for entire max-age period
+    // since the first hit
+    // so instead of using 1 we use 1 year ago
+    // this is done here instead at build time as it's unclear
+    // what exactly is edge case in CDN and setting 1 year ago
+    // from request time seems to work consistently
+    lastModified = Date.now() - 31536000000
+  }
+
   const lastModifiedDate = new Date(lastModified)
   // Show actual date of the function call in the date header
   headers.set('x-nextjs-date', headers.get('date') ?? lastModifiedDate.toUTCString())