[Runtime prefetch] resolve runtime APIs in a separate task (#82475)

In a runtime prefetch, sync IO is dangerous, because it makes us abort
the whole prerender. If this abort happens early, a runtime prefetch can
end up giving us *less* data than a static prefetch. To prevent this, we
can split the final prerender into two separate tasks (which i'm calling
"stages" in the code):
1. The static stage (task 1). This is meant to be equivalent to what
we'd do during build, so runtime APIs like `cookies()` aren't resolved
yet. This lets us render everything that's reachable statically. We
should not hit a sync IO here -- if we had, it would've also caused an
error during build.
2. The runtime stage (task 2). Here, we allow `cookies()` (and all other
runtime APIs available in a runtime prerender) to resolve. We might
encounter sync IO errors here (e.g. `await cookies(); Date.now()`),
which will make the result suboptimal, but still usable.

Author: lubieowoce

packages/next/errors.json

@@ -775,5 +775,6 @@
   "774": "Route %s used %s outside of a Server Component. This is not allowed.",
   "775": "Node.js instrumentation extensions should not be loaded in the Edge runtime.",
   "776": "`unstable_isUnrecognizedActionError` can only be used on the client.",
-  "777": "Invariant: failed to find source route %s for prerender %s"
+  "777": "Invariant: failed to find source route %s for prerender %s",
+  "778": "`prerenderAndAbortInSequentialTasksWithStages` should not be called in edge runtime."
 }

packages/next/src/client/components/segment-cache-impl/cache.ts

@@ -78,6 +78,7 @@ import {
   doesExportedHtmlMatchBuildId,
 } from '../../../shared/lib/segment-cache/output-export-prefetch-encoding'
 import { FetchStrategy } from '../segment-cache'
+import { createPromiseWithResolvers } from '../../../shared/lib/promise-with-resolvers'
 
 // A note on async/await when working in the prefetch cache:
 //
@@ -2050,17 +2051,6 @@ function addSegmentPathToUrlInOutputExportMode(
   return url
 }
 
-function createPromiseWithResolvers<T>(): PromiseWithResolvers<T> {
-  // Shim of Stage 4 Promise.withResolvers proposal
-  let resolve: (value: T | PromiseLike<T>) => void
-  let reject: (reason: any) => void
-  const promise = new Promise<T>((res, rej) => {
-    resolve = res
-    reject = rej
-  })
-  return { resolve: resolve!, reject: reject!, promise }
-}
-
 /**
  * Checks whether the new fetch strategy is likely to provide more content than the old one.
  *

packages/next/src/server/app-render/app-render-prerender-utils.ts

@@ -31,6 +31,41 @@ export function prerenderAndAbortInSequentialTasks<R>(
   }
 }
 
+/**
+ * Like `prerenderAndAbortInSequentialTasks`, but with another task between `prerender` and `abort`,
+ * which allows us to move a part of the render into a separate task.
+ */
+export function prerenderAndAbortInSequentialTasksWithStages<R>(
+  prerender: () => Promise<R>,
+  advanceStage: () => void,
+  abort: () => void
+): Promise<R> {
+  if (process.env.NEXT_RUNTIME === 'edge') {
+    throw new InvariantError(
+      '`prerenderAndAbortInSequentialTasksWithStages` should not be called in edge runtime.'
+    )
+  } else {
+    return new Promise((resolve, reject) => {
+      let pendingResult: Promise<R>
+      setImmediate(() => {
+        try {
+          pendingResult = prerender()
+          pendingResult.catch(() => {})
+        } catch (err) {
+          reject(err)
+        }
+      })
+      setImmediate(() => {
+        advanceStage()
+      })
+      setImmediate(() => {
+        abort()
+        resolve(pendingResult)
+      })
+    })
+  }
+}
+
 // React's RSC prerender function will emit an incomplete flight stream when using `prerender`. If the connection
 // closes then whatever hanging chunks exist will be errored. This is because prerender (an experimental feature)
 // has not yet implemented a concept of resume. For now we will simulate a paused connection by wrapping the stream

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

@@ -155,7 +155,10 @@ import { createMutableActionQueue } from '../../client/components/app-router-ins
 import { getRevalidateReason } from '../instrumentation/utils'
 import { PAGE_SEGMENT_KEY } from '../../shared/lib/segment'
 import type { FallbackRouteParams } from '../request/fallback-params'
-import { processPrelude } from './app-render-prerender-utils'
+import {
+  prerenderAndAbortInSequentialTasksWithStages,
+  processPrelude,
+} from './app-render-prerender-utils'
 import {
   type ReactServerPrerenderResult,
   ReactServerResult,
@@ -200,6 +203,7 @@ import { getRequestMeta } from '../request-meta'
 import { getDynamicParam } from '../../shared/lib/router/utils/get-dynamic-param'
 import type { ExperimentalConfig } from '../config-shared'
 import type { Params } from '../request/params'
+import { createPromiseWithResolvers } from '../../shared/lib/promise-with-resolvers'
 
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
@@ -711,6 +715,8 @@ async function prospectiveRuntimeServerPrerender(
     prerenderResumeDataCache,
     hmrRefreshHash: undefined,
     captureOwnerStack: undefined,
+    // We only need task sequencing in the final prerender.
+    runtimeStagePromise: null,
     // These are not present in regular prerenders, but allowed in a runtime prerender.
     cookies,
     draftMode,
@@ -821,6 +827,9 @@ async function finalRuntimeServerPrerender(
     isDebugDynamicAccesses
   )
 
+  const { promise: runtimeStagePromise, resolve: resolveBlockedRuntimeAPIs } =
+    createPromiseWithResolvers<void>()
+
   const finalServerPrerenderStore: PrerenderStoreModernRuntime = {
     type: 'prerender-runtime',
     phase: 'render',
@@ -841,6 +850,8 @@ async function finalRuntimeServerPrerender(
     renderResumeDataCache,
     hmrRefreshHash: undefined,
     captureOwnerStack: undefined,
+    // Used to separate the "Static" stage from the "Runtime" stage.
+    runtimeStagePromise,
     // These are not present in regular prerenders, but allowed in a runtime prerender.
     cookies,
     draftMode,
@@ -852,8 +863,9 @@ async function finalRuntimeServerPrerender(
   )
 
   let prerenderIsPending = true
-  const result = await prerenderAndAbortInSequentialTasks(
+  const result = await prerenderAndAbortInSequentialTasksWithStages(
     async () => {
+      // Static stage
       const prerenderResult = await workUnitAsyncStorage.run(
         finalServerPrerenderStore,
         ComponentMod.prerender,
@@ -869,6 +881,17 @@ async function finalRuntimeServerPrerender(
       return prerenderResult
     },
     () => {
+      // Advance to the runtime stage.
+      //
+      // We make runtime APIs hang during the first task (above), and unblock them in the following task (here).
+      // This makes sure that, at this point, we'll have finished all the static parts (what we'd prerender statically).
+      // We know that they don't contain any incorrect sync IO, because that'd have caused a build error.
+      // After we unblock Runtime APIs, if we encounter sync IO (e.g. `await cookies(); Date.now()`),
+      // we'll abort, but we'll produce at least as much output as a static prerender would.
+      resolveBlockedRuntimeAPIs()
+    },
+    () => {
+      // Abort.
       if (finalServerController.signal.aborted) {
         // If the server controller is already aborted we must have called something
         // that required aborting the prerender synchronously such as with new Date()

packages/next/src/server/app-render/dynamic-rendering.ts

@@ -26,14 +26,18 @@ import type {
   RequestStore,
   PrerenderStoreLegacy,
   PrerenderStoreModern,
+  PrerenderStoreModernRuntime,
 } from '../app-render/work-unit-async-storage.external'
 
 // Once postpone is in stable we should switch to importing the postpone export directly
 import React from 'react'
 
 import { DynamicServerError } from '../../client/components/hooks-server-context'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
-import { workUnitAsyncStorage } from './work-unit-async-storage.external'
+import {
+  getRuntimeStagePromise,
+  workUnitAsyncStorage,
+} from './work-unit-async-storage.external'
 import { workAsyncStorage } from '../app-render/work-async-storage.external'
 import { makeHangingPromise } from '../dynamic-rendering-utils'
 import {
@@ -548,12 +552,25 @@ export function createHangingInputAbortSignal(
         })
       } else {
         // Otherwise we're in the final render and we should already have all
-        // our caches filled. We might still be waiting on some microtasks so we
+        // our caches filled.
+        // If the prerender uses stages, we have wait until the runtime stage,
+        // at which point all runtime inputs will be resolved.
+        // (otherwise, a runtime prerender might consider `cookies()` hanging
+        //  even though they'd resolve in the next task.)
+        //
+        // We might still be waiting on some microtasks so we
         // wait one tick before giving up. When we give up, we still want to
         // render the content of this cache as deeply as we can so that we can
         // suspend as deeply as possible in the tree or not at all if we don't
         // end up waiting for the input.
-        scheduleOnNextTick(() => controller.abort())
+        const runtimeStagePromise = getRuntimeStagePromise(workUnitStore)
+        if (runtimeStagePromise) {
+          runtimeStagePromise.then(() =>
+            scheduleOnNextTick(() => controller.abort())
+          )
+        } else {
+          scheduleOnNextTick(() => controller.abort())
+        }
       }
 
       return controller.signal
@@ -824,3 +841,13 @@ export function throwIfDisallowedDynamic(
     }
   }
 }
+
+export function delayUntilRuntimeStage<T>(
+  prerenderStore: PrerenderStoreModernRuntime,
+  result: Promise<T>
+): Promise<T> {
+  if (prerenderStore.runtimeStagePromise) {
+    return prerenderStore.runtimeStagePromise.then(() => result)
+  }
+  return result
+}

packages/next/src/server/app-render/work-unit-async-storage.external.ts

@@ -106,6 +106,18 @@ export interface PrerenderStoreModernRuntime
   extends PrerenderStoreModernCommon {
   readonly type: 'prerender-runtime'
 
+  /**
+   * A runtime prerender resolves APIs in two tasks:
+   *
+   * 1. Static data (available in a static prerender)
+   * 2. Runtime data (available in a runtime prerender)
+   *
+   * This separation is achieved by awaiting this promise in "runtime" APIs.
+   * In the final prerender, the promise will be resolved during the second task,
+   * and the render will be aborted in the task that follows it.
+   */
+  readonly runtimeStagePromise: Promise<void> | null
+
   readonly cookies: RequestStore['cookies']
   readonly draftMode: RequestStore['draftMode']
 }
@@ -268,6 +280,18 @@ export interface PublicUseCacheStore extends CommonUseCacheStore {
 export interface PrivateUseCacheStore extends CommonUseCacheStore {
   readonly type: 'private-cache'
 
+  /**
+   * A runtime prerender resolves APIs in two tasks:
+   *
+   * 1. Static data (available in a static prerender)
+   * 2. Runtime data (available in a runtime prerender)
+   *
+   * This separation is achieved by awaiting this promise in "runtime" APIs.
+   * In the final prerender, the promise will be resolved during the second task,
+   * and the render will be aborted in the task that follows it.
+   */
+  readonly runtimeStagePromise: Promise<void> | null
+
   /**
    * As opposed to the public cache store, the private cache store is allowed to
    * access the request cookies.
@@ -486,3 +510,23 @@ export function getCacheSignal(
       return workUnitStore satisfies never
   }
 }
+
+export function getRuntimeStagePromise(
+  workUnitStore: WorkUnitStore
+): Promise<void> | null {
+  switch (workUnitStore.type) {
+    case 'prerender-runtime':
+    case 'private-cache':
+      return workUnitStore.runtimeStagePromise
+    case 'prerender':
+    case 'prerender-client':
+    case 'prerender-ppr':
+    case 'prerender-legacy':
+    case 'request':
+    case 'cache':
+    case 'unstable-cache':
+      return null
+    default:
+      return workUnitStore satisfies never
+  }
+}

packages/next/src/server/request/cookies.ts

@@ -15,6 +15,7 @@ import {
   type PrerenderStoreModern,
 } from '../app-render/work-unit-async-storage.external'
 import {
+  delayUntilRuntimeStage,
   postponeWithTracking,
   throwToInterruptStaticGeneration,
   trackDynamicDataInDynamicRender,
@@ -118,6 +119,10 @@ export function cookies(): Promise<ReadonlyRequestCookies> {
             workUnitStore
           )
         case 'prerender-runtime':
+          return delayUntilRuntimeStage(
+            workUnitStore,
+            makeUntrackedExoticCookies(workUnitStore.cookies)
+          )
         case 'private-cache':
           return makeUntrackedExoticCookies(workUnitStore.cookies)
         case 'request':

packages/next/src/server/request/draft-mode.ts

@@ -12,6 +12,7 @@ import {
 import { workUnitAsyncStorage } from '../app-render/work-unit-async-storage.external'
 import {
   abortAndThrowOnSynchronousRequestDataAccess,
+  delayUntilRuntimeStage,
   postponeWithTracking,
   trackDynamicDataInDynamicRender,
   trackSynchronousRequestDataAccessInDev,
@@ -56,6 +57,11 @@ export function draftMode(): Promise<DraftMode> {
 
   switch (workUnitStore.type) {
     case 'prerender-runtime':
+      // TODO(runtime-ppr): does it make sense to delay this? normally it's always microtasky
+      return delayUntilRuntimeStage(
+        workUnitStore,
+        createOrGetCachedDraftMode(workUnitStore.draftMode, workStore)
+      )
     case 'request':
       return createOrGetCachedDraftMode(workUnitStore.draftMode, workStore)
 

packages/next/src/server/request/params.ts

@@ -9,6 +9,7 @@ import {
   throwToInterruptStaticGeneration,
   postponeWithTracking,
   trackSynchronousRequestDataAccessInDev,
+  delayUntilRuntimeStage,
 } from '../app-render/dynamic-rendering'
 
 import {
@@ -114,6 +115,10 @@ export function createServerParamsForRoute(
           'createServerParamsForRoute should not be called in cache contexts.'
         )
       case 'prerender-runtime':
+        return delayUntilRuntimeStage(
+          workUnitStore,
+          createRenderParams(underlyingParams, workStore)
+        )
       case 'request':
         break
       default:
@@ -142,6 +147,10 @@ export function createServerParamsForServerSegment(
           'createServerParamsForServerSegment should not be called in cache contexts.'
         )
       case 'prerender-runtime':
+        return delayUntilRuntimeStage(
+          workUnitStore,
+          createRenderParams(underlyingParams, workStore)
+        )
       case 'request':
         break
       default:

packages/next/src/server/request/pathname.ts

@@ -1,6 +1,7 @@
 import type { WorkStore } from '../app-render/work-async-storage.external'
 
 import {
+  delayUntilRuntimeStage,
   postponeWithTracking,
   type DynamicTrackingState,
 } from '../app-render/dynamic-rendering'
@@ -37,6 +38,10 @@ export function createServerPathnameForMetadata(
         )
 
       case 'prerender-runtime':
+        return delayUntilRuntimeStage(
+          workUnitStore,
+          createRenderPathname(underlyingPathname)
+        )
       case 'request':
         break
       default: