Backport: Use provided waitUntil for pending revalidates (#74164) (#74573)

This backports #74164 to leverage
built-in waitUntil if available instead of using the approach that keeps
the stream open until the waitUntil promise resolves.

x-ref: [slack
thread](https://vercel.slack.com/archives/C02K2HCH5V4/p1736211642221149?thread_ts=1734707275.666089&cid=C02K2HCH5V4)

Author: ijjk

packages/next/src/server/app-render/app-render.tsx

@@ -148,6 +148,7 @@ export type AppRenderContext = AppRenderBaseContext & {
   serverComponentsErrorHandler: ErrorHandler
   isNotFoundPath: boolean
   res: ServerResponse
+  builtInWaitUntil: RenderOpts['builtInWaitUntil']
 }
 
 function createNotFoundLoaderTree(loaderTree: LoaderTree): LoaderTree {
@@ -387,12 +388,23 @@ async function generateFlight(
     ctx.staticGenerationStore.pendingRevalidates ||
     ctx.staticGenerationStore.revalidatedTags
   ) {
-    resultOptions.waitUntil = Promise.all([
+    const pendingPromise = Promise.all([
       ctx.staticGenerationStore.incrementalCache?.revalidateTag(
         ctx.staticGenerationStore.revalidatedTags || []
       ),
       ...Object.values(ctx.staticGenerationStore.pendingRevalidates || {}),
-    ])
+    ]).finally(() => {
+      if (process.env.NEXT_PRIVATE_DEBUG_CACHE) {
+        console.log('pending revalidates promise finished for:', urlPathname)
+      }
+    })
+
+    // use built-in waitUntil if available
+    if (ctx.builtInWaitUntil) {
+      ctx.builtInWaitUntil(pendingPromise)
+    } else {
+      resultOptions.waitUntil = pendingPromise
+    }
   }
 
   return new FlightRenderResult(flightReadableStream, resultOptions)
@@ -848,6 +860,7 @@ async function renderToHTMLOrFlightImpl(
 
   const ctx: AppRenderContext = {
     ...baseCtx,
+    builtInWaitUntil: renderOpts.builtInWaitUntil,
     getDynamicParamFromSegment,
     query,
     isPrefetch: isPrefetchRSCRequest,
@@ -1401,12 +1414,23 @@ async function renderToHTMLOrFlightImpl(
     staticGenerationStore.pendingRevalidates ||
     staticGenerationStore.revalidatedTags
   ) {
-    options.waitUntil = Promise.all([
+    const pendingPromise = Promise.all([
       staticGenerationStore.incrementalCache?.revalidateTag(
         staticGenerationStore.revalidatedTags || []
       ),
       ...Object.values(staticGenerationStore.pendingRevalidates || {}),
-    ])
+    ]).finally(() => {
+      if (process.env.NEXT_PRIVATE_DEBUG_CACHE) {
+        console.log('pending revalidates promise finished for:', req.url)
+      }
+    })
+
+    // use built-in waitUntil if available
+    if (renderOpts.builtInWaitUntil) {
+      renderOpts.builtInWaitUntil(pendingPromise)
+    } else {
+      options.waitUntil = pendingPromise
+    }
   }
 
   addImplicitTags(staticGenerationStore)

packages/next/src/server/app-render/types.ts

@@ -10,6 +10,7 @@ import type { LoadingModuleData } from '../../shared/lib/app-router-context.shar
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
 
 import s from 'next/dist/compiled/superstruct'
+import type { WaitUntil } from '../lib/builtin-request-context'
 
 export type DynamicParamTypes =
   | 'catchall'
@@ -170,6 +171,8 @@ export interface RenderOptsPartial {
    */
   isDebugPPRSkeleton?: boolean
   isStaticGeneration?: boolean
+
+  builtInWaitUntil?: WaitUntil
 }
 
 export type RenderOpts = LoadComponentsReturnType<AppPageModule> &

packages/next/src/server/async-storage/static-generation-async-storage-wrapper.ts

@@ -42,6 +42,7 @@ export type StaticGenerationContext = {
     | 'nextExport'
     | 'isDraftMode'
     | 'isDebugPPRSkeleton'
+    | 'builtInWaitUntil'
   >
 }
 

packages/next/src/server/base-server.ts

@@ -20,6 +20,10 @@ import {
   normalizeRepeatedSlashes,
   MissingStaticPage,
 } from '../shared/lib/utils'
+import {
+  getBuiltinRequestContext,
+  type WaitUntil,
+} from './lib/builtin-request-context'
 import type { PreviewData } from 'next/types'
 import type { PagesManifest } from '../build/webpack/plugins/pages-manifest-plugin'
 import type { BaseNextRequest, BaseNextResponse } from './base-http'
@@ -1596,6 +1600,18 @@ export default abstract class Server<ServerOptions extends Options = Options> {
     )
   }
 
+  protected getWaitUntil(): WaitUntil | undefined {
+    const builtinRequestContext = getBuiltinRequestContext()
+    if (builtinRequestContext) {
+      // the platform provided a request context.
+      // use the `waitUntil` from there, whether actually present or not --
+      // if not present, `after` will error.
+
+      // NOTE: if we're in an edge runtime sandbox, this context will be used to forward the outer waitUntil.
+      return builtinRequestContext.waitUntil
+    }
+  }
+
   private async renderImpl(
     req: BaseNextRequest,
     res: BaseNextResponse,
@@ -2198,6 +2214,7 @@ export default abstract class Server<ServerOptions extends Options = Options> {
         isDraftMode: isPreviewMode,
         isServerAction,
         postponed,
+        builtInWaitUntil: this.getWaitUntil(),
       }
 
       if (isDebugPPRSkeleton) {
@@ -2225,6 +2242,7 @@ export default abstract class Server<ServerOptions extends Options = Options> {
               supportsDynamicResponse,
               incrementalCache,
               isRevalidate: isSSG,
+              builtInWaitUntil: this.getWaitUntil(),
             },
           }
 

packages/next/src/server/future/route-modules/app-route/module.ts

@@ -383,14 +383,28 @@ export class AppRouteRouteModule extends RouteModule<
                     context.renderOpts.fetchMetrics =
                       staticGenerationStore.fetchMetrics
 
-                    context.renderOpts.waitUntil = Promise.all([
+                    const pendingPromise = Promise.all([
                       staticGenerationStore.incrementalCache?.revalidateTag(
                         staticGenerationStore.revalidatedTags || []
                       ),
                       ...Object.values(
                         staticGenerationStore.pendingRevalidates || {}
                       ),
-                    ])
+                    ]).finally(() => {
+                      if (process.env.NEXT_PRIVATE_DEBUG_CACHE) {
+                        console.log(
+                          'pending revalidates promise finished for:',
+                          rawRequest.url.toString()
+                        )
+                      }
+                    })
+
+                    // use built-in waitUntil if available
+                    if (context.renderOpts.builtInWaitUntil) {
+                      context.renderOpts.builtInWaitUntil(pendingPromise)
+                    } else {
+                      context.renderOpts.waitUntil = pendingPromise
+                    }
 
                     addImplicitTags(staticGenerationStore)
                     ;(context.renderOpts as any).fetchTags =

packages/next/src/server/lib/builtin-request-context.ts

@@ -0,0 +1,41 @@
+import { createAsyncLocalStorage } from '../../client/components/async-local-storage'
+
+export function getBuiltinRequestContext():
+  | BuiltinRequestContextValue
+  | undefined {
+  const _globalThis = globalThis as GlobalThisWithRequestContext
+  const ctx = _globalThis[NEXT_REQUEST_CONTEXT_SYMBOL]
+  return ctx?.get()
+}
+
+const NEXT_REQUEST_CONTEXT_SYMBOL = Symbol.for('@next/request-context')
+
+type GlobalThisWithRequestContext = typeof globalThis & {
+  [NEXT_REQUEST_CONTEXT_SYMBOL]?: BuiltinRequestContext
+}
+
+/** A request context provided by the platform. */
+export type BuiltinRequestContext = {
+  get(): BuiltinRequestContextValue | undefined
+}
+
+export type RunnableBuiltinRequestContext = BuiltinRequestContext & {
+  run<T>(value: BuiltinRequestContextValue, callback: () => T): T
+}
+
+export type BuiltinRequestContextValue = {
+  waitUntil?: WaitUntil
+}
+export type WaitUntil = (promise: Promise<any>) => void
+
+/** "@next/request-context" has a different signature from AsyncLocalStorage,
+ * matching [AsyncContext.Variable](https://github.com/tc39/proposal-async-context).
+ * We don't need a full AsyncContext adapter here, just having `.get()` is enough
+ */
+export function createLocalRequestContext(): RunnableBuiltinRequestContext {
+  const storage = createAsyncLocalStorage<BuiltinRequestContextValue>()
+  return {
+    get: () => storage.getStore(),
+    run: (value, callback) => storage.run(value, callback),
+  }
+}

test/e2e/app-dir/app-fetch-deduping/app-fetch-deduping.test.ts

@@ -1,4 +1,4 @@
-import { findPort, waitFor } from 'next-test-utils'
+import { findPort, retry } from 'next-test-utils'
 import http from 'http'
 import { outdent } from 'outdent'
 import { isNextDev, isNextStart, nextTestSetup } from 'e2e-utils'
@@ -104,12 +104,10 @@ describe('app-fetch-deduping', () => {
 
         expect(invocation(next.cliOutput)).toBe(1)
 
-        // wait for the revalidation to finish
-        await waitFor(revalidate * 1000 + 1000)
-
-        await next.render('/test')
-
-        expect(invocation(next.cliOutput)).toBe(2)
+        await retry(async () => {
+          await next.render('/test')
+          expect(invocation(next.cliOutput)).toBe(2)
+        }, 10_000)
       })
     })
   } else {