use already existing trackBackgroundWork helper from request context to handle next/after

Author: pieh

src/build/templates/handler-monorepo.tmpl.js

@@ -16,7 +16,7 @@ export default async function (req, context) {
     tracing.start()
   }
 
-  const requestContext = createRequestContext(req)
+  const requestContext = createRequestContext(req, context)
   const tracer = getTracer()
 
   const handlerResponse = await runWithRequestContext(requestContext, () => {

src/build/templates/handler.tmpl.js

@@ -13,7 +13,7 @@ export default async function handler(req, context) {
   if (process.env.NETLIFY_OTLP_TRACE_EXPORTER_URL) {
     tracing.start()
   }
-  const requestContext = createRequestContext(req)
+  const requestContext = createRequestContext(req, context)
   const tracer = getTracer()
 
   const handlerResponse = await runWithRequestContext(requestContext, () => {

src/run/handlers/request-context.cts

@@ -1,11 +1,18 @@
 import { AsyncLocalStorage } from 'node:async_hooks'
 
+import type { Context } from '@netlify/functions'
 import { LogLevel, systemLogger } from '@netlify/functions/internal'
 
 import type { NetlifyCachedRouteValue } from '../../shared/cache-types.cjs'
 
 type SystemLogger = typeof systemLogger
 
+// TODO: remove once https://github.com/netlify/serverless-functions-api/pull/219
+// is released and public types are updated
+export interface FutureContext extends Context {
+  waitUntil?: (promise: Promise<unknown>) => void
+}
+
 export type RequestContext = {
   captureServerTiming: boolean
   responseCacheGetLastModified?: number
@@ -28,13 +35,17 @@ export type RequestContext = {
 
 type RequestContextAsyncLocalStorage = AsyncLocalStorage<RequestContext>
 
-export function createRequestContext(request?: Request): RequestContext {
+export function createRequestContext(request?: Request, context?: FutureContext): RequestContext {
   const backgroundWorkPromises: Promise<unknown>[] = []
 
   return {
     captureServerTiming: request?.headers.has('x-next-debug-logging') ?? false,
     trackBackgroundWork: (promise) => {
-      backgroundWorkPromises.push(promise)
+      if (context?.waitUntil) {
+        context.waitUntil(promise)
+      } else {
+        backgroundWorkPromises.push(promise)
+      }
     },
     get backgroundWorkPromise() {
       return Promise.allSettled(backgroundWorkPromises)

src/run/handlers/server.ts

@@ -1,7 +1,6 @@
 import type { OutgoingHttpHeaders } from 'http'
 
 import { ComputeJsOutgoingMessage, toComputeResponse, toReqRes } from '@fastly/http-compute-js'
-import { Context } from '@netlify/functions'
 import type { NextConfigComplete } from 'next/dist/server/config-shared.js'
 import type { WorkerRequestHandler } from 'next/dist/server/lib/types.js'
 
@@ -16,10 +15,12 @@ import { nextResponseProxy } from '../revalidate.js'
 
 import { createRequestContext, getLogger, getRequestContext } from './request-context.cjs'
 import { getTracer } from './tracer.cjs'
-import { setWaitUntil } from './wait-until.cjs'
+import { setupWaitUntil } from './wait-until.cjs'
 
 const nextImportPromise = import('../next.cjs')
 
+setupWaitUntil()
+
 let nextHandler: WorkerRequestHandler, nextConfig: NextConfigComplete
 
 /**
@@ -45,15 +46,9 @@ const disableFaultyTransferEncodingHandling = (res: ComputeJsOutgoingMessage) =>
   }
 }
 
-// TODO: remove once https://github.com/netlify/serverless-functions-api/pull/219
-// is released and public types are updated
-interface FutureContext extends Context {
-  waitUntil?: (promise: Promise<unknown>) => void
-}
-
-export default async (request: Request, context: FutureContext) => {
+export default async (request: Request) => {
   const tracer = getTracer()
-  setWaitUntil(context)
+
   if (!nextHandler) {
     await tracer.withActiveSpan('initialize next server', async () => {
       // set the server config
@@ -129,19 +124,19 @@ export default async (request: Request, context: FutureContext) => {
       return new Response(body || null, response)
     }
 
-    if (context.waitUntil) {
-      context.waitUntil(requestContext.backgroundWorkPromise)
-    }
-
     const keepOpenUntilNextFullyRendered = new TransformStream({
       async flush() {
         // it's important to keep the stream open until the next handler has finished
         await nextHandlerPromise
-        if (!context.waitUntil) {
-          // if waitUntil is not available, we have to keep response stream open until background promises are resolved
-          // to ensure that all background work executes
-          await requestContext.backgroundWorkPromise
-        }
+
+        // Next.js relies on `close` event emitted by response to trigger running callback variant of `next/after`
+        // however @fastly/http-compute-js never actually emits that event - so we have to emit it ourselves,
+        // otherwise Next would never run the callback variant of `next/after`
+        res.emit('close')
+
+        // if waitUntil is not available, we have to keep response stream open until background promises are resolved
+        // to ensure that all background work executes
+        await requestContext.backgroundWorkPromise
       },
     })
 

src/run/handlers/wait-until.cts

@@ -1,3 +1,5 @@
+import { getRequestContext } from './request-context.cjs'
+
 /**
  *  @see https://github.com/vercel/next.js/blob/canary/packages/next/src/server/after/builtin-request-context.ts
  */
@@ -14,15 +16,11 @@ type GlobalThisWithRequestContext = typeof globalThis & {
 /**
  * Registers a `waitUntil` to be used by Next.js for next/after
  */
-
-export function setWaitUntil({ waitUntil }: { waitUntil?: (promise: Promise<unknown>) => void }) {
-  if (!waitUntil) {
-    return
-  }
+export function setupWaitUntil() {
   // eslint-disable-next-line @typescript-eslint/no-extra-semi
   ;(globalThis as GlobalThisWithRequestContext)[NEXT_REQUEST_CONTEXT_SYMBOL] = {
     get() {
-      return { waitUntil }
+      return { waitUntil: getRequestContext()?.trackBackgroundWork }
     },
   }
 }