[Cache Components] separate runtime stage in dev render

Author: lubieowoce

packages/next/errors.json

@@ -878,5 +878,6 @@
   "877": "Config options `experimental.externalProxyRewritesResolve` and `experimental.externalMiddlewareRewritesResolve` cannot be set at the same time. Please use `experimental.externalProxyRewritesResolve` instead.",
   "878": "Config options `skipProxyUrlNormalize` and `skipMiddlewareUrlNormalize` cannot be set at the same time. Please use `skipProxyUrlNormalize` instead.",
   "879": "Config options `experimental.proxyClientMaxBodySize` and `experimental.middlewareClientMaxBodySize` cannot be set at the same time. Please use `experimental.proxyClientMaxBodySize` instead.",
-  "880": "Config options `experimental.proxyPrefetch` and `experimental.middlewarePrefetch` cannot be set at the same time. Please use `experimental.proxyPrefetch` instead."
+  "880": "Config options `experimental.proxyPrefetch` and `experimental.middlewarePrefetch` cannot be set at the same time. Please use `experimental.proxyPrefetch` instead.",
+  "881": "Invalid render stage: %s"
 }

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

@@ -35,30 +35,45 @@ export function scheduleInSequentialTasks<R>(
  * We schedule on the same queue (setTimeout) at the same time to ensure no other events can sneak in between.
  * The function that runs in the second task gets access to the first tasks's result.
  */
-export function pipelineInSequentialTasks<A, B>(
-  render: () => A,
-  followup: (a: A) => B | Promise<B>
-): Promise<B> {
+export function pipelineInSequentialTasks<A, B, C>(
+  one: () => A,
+  two: (a: A) => B,
+  three: (b: B) => C | Promise<C>
+): Promise<C> {
   if (process.env.NEXT_RUNTIME === 'edge') {
     throw new InvariantError(
       '`pipelineInSequentialTasks` should not be called in edge runtime.'
     )
   } else {
     return new Promise((resolve, reject) => {
-      let renderResult: A | undefined = undefined
+      let oneResult: A | undefined = undefined
       setTimeout(() => {
         try {
-          renderResult = render()
+          oneResult = one()
         } catch (err) {
-          clearTimeout(followupId)
+          clearTimeout(twoId)
+          clearTimeout(threeId)
           reject(err)
         }
       }, 0)
-      const followupId = setTimeout(() => {
-        // if `render` threw, then the `followup` timeout would've been cleared,
-        // so if we got here, we're guaranteed to have a `renderResult`.
+
+      let twoResult: B | undefined = undefined
+      const twoId = setTimeout(() => {
+        // if `one` threw, then this timeout would've been cleared,
+        // so if we got here, we're guaranteed to have a value.
+        try {
+          twoResult = two(oneResult!)
+        } catch (err) {
+          clearTimeout(threeId)
+          reject(err)
+        }
+      }, 0)
+
+      const threeId = setTimeout(() => {
+        // if `two` threw, then this timeout would've been cleared,
+        // so if we got here, we're guaranteed to have a value.
         try {
-          resolve(followup(renderResult!))
+          resolve(three(twoResult!))
         } catch (err) {
           reject(err)
         }

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

@@ -168,10 +168,7 @@ import {
   prerenderAndAbortInSequentialTasks,
 } from './app-render-prerender-utils'
 import { printDebugThrownValueForProspectiveRender } from './prospective-render-utils'
-import {
-  pipelineInSequentialTasks,
-  scheduleInSequentialTasks,
-} from './app-render-render-utils'
+import { pipelineInSequentialTasks } from './app-render-render-utils'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import {
   workUnitAsyncStorage,
@@ -214,6 +211,7 @@ import type { Params } from '../request/params'
 import { createPromiseWithResolvers } from '../../shared/lib/promise-with-resolvers'
 import { ImageConfigContext } from '../../shared/lib/image-config-context.shared-runtime'
 import { imageConfigDefault } from '../../shared/lib/image-config'
+import { RenderStage, StagedRenderingController } from './staged-rendering'
 
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
@@ -2678,8 +2676,21 @@ async function renderWithRestartOnCacheMissInDev(
   // If the render is restarted, we'll recreate a fresh request store
   let requestStore: RequestStore = initialRequestStore
 
-  const environmentName = () =>
-    requestStore.prerenderPhase === true ? 'Prerender' : 'Server'
+  const environmentName = () => {
+    const currentStage = requestStore.stagedRendering!.currentStage
+    switch (currentStage) {
+      case RenderStage.Static:
+        return 'Prerender'
+      case RenderStage.Runtime:
+        // TODO: only label as "Prefetch" if the page has a `prefetch` config.
+        return 'Prefetch'
+      case RenderStage.Dynamic:
+        return 'Server'
+      default:
+        currentStage satisfies never
+        throw new InvariantError(`Invalid render stage: ${currentStage}`)
+    }
+  }
 
   //===============================================
   // Initial render
@@ -2699,14 +2710,19 @@ async function renderWithRestartOnCacheMissInDev(
 
   const prerenderResumeDataCache = createPrerenderResumeDataCache()
 
+  const initialReactController = new AbortController()
+  const initialDataController = new AbortController() // Controls hanging promises we create
+  const initialStageController = new StagedRenderingController(
+    initialDataController.signal
+  )
+
   requestStore.prerenderResumeDataCache = prerenderResumeDataCache
   // `getRenderResumeDataCache` will fall back to using `prerenderResumeDataCache` as `renderResumeDataCache`,
   // so not having a resume data cache won't break any expectations in case we don't need to restart.
   requestStore.renderResumeDataCache = null
+  requestStore.stagedRendering = initialStageController
   requestStore.cacheSignal = cacheSignal
 
-  const initialReactController = new AbortController()
-
   let debugChannel = setReactDebugChannel && createDebugChannel()
 
   const initialRscPayload = await getPayload(requestStore)
@@ -2716,8 +2732,7 @@ async function renderWithRestartOnCacheMissInDev(
       pipelineInSequentialTasks(
         () => {
           // Static stage
-          requestStore.prerenderPhase = true
-          return ComponentMod.renderToReadableStream(
+          const stream = ComponentMod.renderToReadableStream(
             initialRscPayload,
             clientReferenceManifest.clientModules,
             {
@@ -2728,25 +2743,42 @@ async function renderWithRestartOnCacheMissInDev(
               signal: initialReactController.signal,
             }
           )
+          // If we abort the render, we want to reject the stage-dependent promises as well.
+          // Note that we want to install this listener after the render is started
+          // so that it runs after react is finished running its abort code.
+          initialReactController.signal.addEventListener('abort', () => {
+            initialDataController.abort(initialReactController.signal.reason)
+          })
+          return stream
+        },
+        (stream) => {
+          // Runtime stage
+          initialStageController.advanceStage(RenderStage.Runtime)
+
+          // If we had a cache miss in the static stage, we'll have to disard this stream
+          // and render again once the caches are warm.
+          if (cacheSignal.hasPendingReads()) {
+            return null
+          }
+
+          // If there's no cache misses, we'll continue rendering,
+          // and see if there's any cache misses in the runtime stage.
+          return stream
         },
-        async (stream) => {
+        async (maybeStream) => {
           // Dynamic stage
-          // Note: if we had cache misses, things that would've happened statically otherwise
-          // may be marked as dynamic instead.
-          requestStore.prerenderPhase = false
-
-          // If all cache reads initiated in the static stage have completed,
-          // then all of the necessary caches have to be warm (or there's no caches on the page).
-          // On the other hand, if we still have pending cache reads, then we had a cache miss,
-          // and the static stage didn't render all the content that it normally would have.
-          const hadCacheMiss = cacheSignal.hasPendingReads()
-          if (!hadCacheMiss) {
-            // No cache misses. We can use the stream as is.
-            return stream
-          } else {
-            // Cache miss. We'll discard this stream, and render again.
+
+          // If we had cache misses in either of the previous stages,
+          // then we'll only use this render for filling caches.
+          // We won't advance the stage, and thus leave dynamic APIs hanging,
+          // because they won't be cached anyway, so it'd be wasted work.
+          if (maybeStream === null || cacheSignal.hasPendingReads()) {
             return null
           }
+
+          // If there's no cache misses, we'll use this render, so let it advance to the dynamic stage.
+          initialStageController.advanceStage(RenderStage.Dynamic)
+          return maybeStream
         }
       )
   )
@@ -2779,40 +2811,48 @@ async function renderWithRestartOnCacheMissInDev(
   // The initial render acted as a prospective render to warm the caches.
   requestStore = createRequestStore()
 
+  const finalStageController = new StagedRenderingController()
+
   // We've filled the caches, so now we can render as usual,
   // without any cache-filling mechanics.
   requestStore.prerenderResumeDataCache = null
   requestStore.renderResumeDataCache = createRenderResumeDataCache(
     prerenderResumeDataCache
   )
+  requestStore.stagedRendering = finalStageController
   requestStore.cacheSignal = null
 
   // The initial render already wrote to its debug channel.
   // We're not using it, so we need to create a new one.
   debugChannel = setReactDebugChannel && createDebugChannel()
 
   const finalRscPayload = await getPayload(requestStore)
-  const finalServerStream = await workUnitAsyncStorage.run(
-    requestStore,
-    scheduleInSequentialTasks,
-    () => {
-      // Static stage
-      requestStore.prerenderPhase = true
-      return ComponentMod.renderToReadableStream(
-        finalRscPayload,
-        clientReferenceManifest.clientModules,
-        {
-          onError,
-          environmentName,
-          filterStackFrame,
-          debugChannel: debugChannel?.serverSide,
-        }
-      )
-    },
-    () => {
-      // Dynamic stage
-      requestStore.prerenderPhase = false
-    }
+  const finalServerStream = await workUnitAsyncStorage.run(requestStore, () =>
+    pipelineInSequentialTasks(
+      () => {
+        // Static stage
+        return ComponentMod.renderToReadableStream(
+          finalRscPayload,
+          clientReferenceManifest.clientModules,
+          {
+            onError,
+            environmentName,
+            filterStackFrame,
+            debugChannel: debugChannel?.serverSide,
+          }
+        )
+      },
+      (stream) => {
+        // Runtime stage
+        finalStageController.advanceStage(RenderStage.Runtime)
+        return stream
+      },
+      (stream) => {
+        // Dynamic stage
+        finalStageController.advanceStage(RenderStage.Dynamic)
+        return stream
+      }
+    )
   )
 
   return {

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

@@ -50,6 +50,7 @@ import {
 import { scheduleOnNextTick } from '../../lib/scheduler'
 import { BailoutToCSRError } from '../../shared/lib/lazy-dynamic/bailout-to-csr'
 import { InvariantError } from '../../shared/lib/invariant-error'
+import { RenderStage } from './staged-rendering'
 
 const hasPostpone = typeof React.unstable_postpone === 'function'
 
@@ -298,8 +299,12 @@ export function trackSynchronousPlatformIOAccessInDev(
   requestStore: RequestStore
 ): void {
   // We don't actually have a controller to abort but we do the semantic equivalent by
-  // advancing the request store out of prerender mode
-  requestStore.prerenderPhase = false
+  // advancing the request store out of the prerender stage
+  if (requestStore.stagedRendering) {
+    // TODO: error for sync IO in the runtime stage
+    // (which is not currently covered by the validation render in `spawnDynamicValidationInDev`)
+    requestStore.stagedRendering.advanceStage(RenderStage.Dynamic)
+  }
 }
 
 /**

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

@@ -0,0 +1,90 @@
+import { InvariantError } from '../../shared/lib/invariant-error'
+import { createPromiseWithResolvers } from '../../shared/lib/promise-with-resolvers'
+
+export enum RenderStage {
+  Static = 1,
+  Runtime = 2,
+  Dynamic = 3,
+}
+
+export type NonStaticRenderStage = RenderStage.Runtime | RenderStage.Dynamic
+
+export class StagedRenderingController {
+  currentStage: RenderStage = RenderStage.Static
+
+  private runtimeStagePromise = createPromiseWithResolvers<void>()
+  private dynamicStagePromise = createPromiseWithResolvers<void>()
+
+  constructor(private abortSignal: AbortSignal | null = null) {
+    if (abortSignal) {
+      abortSignal.addEventListener(
+        'abort',
+        () => {
+          const { reason } = abortSignal
+          if (this.currentStage < RenderStage.Runtime) {
+            this.runtimeStagePromise.promise.catch(ignoreReject) // avoid unhandled rejections
+            this.runtimeStagePromise.reject(reason)
+          }
+          if (this.currentStage < RenderStage.Dynamic) {
+            this.dynamicStagePromise.promise.catch(ignoreReject) // avoid unhandled rejections
+            this.dynamicStagePromise.reject(reason)
+          }
+        },
+        { once: true }
+      )
+    }
+  }
+
+  advanceStage(stage: NonStaticRenderStage) {
+    // If we're already at the target stage or beyond, do nothing.
+    // (this can happen e.g. if sync IO advanced us to the dynamic stage)
+    if (this.currentStage >= stage) {
+      return
+    }
+    this.currentStage = stage
+    // Note that we might be going directly from Static to Dynamic,
+    // so we need to resolve the runtime stage as well.
+    if (stage >= RenderStage.Runtime) {
+      this.runtimeStagePromise.resolve()
+    }
+    if (stage >= RenderStage.Dynamic) {
+      this.dynamicStagePromise.resolve()
+    }
+  }
+
+  delayUntilStage<T>(stage: NonStaticRenderStage, resolvedValue: T) {
+    let stagePromise: Promise<void>
+    switch (stage) {
+      case RenderStage.Runtime: {
+        stagePromise = this.runtimeStagePromise.promise
+        break
+      }
+      case RenderStage.Dynamic: {
+        stagePromise = this.dynamicStagePromise.promise
+        break
+      }
+      default: {
+        stage satisfies never
+        throw new InvariantError(`Invalid render stage: ${stage}`)
+      }
+    }
+
+    // FIXME: this seems to be the only form that leads to correct API names
+    // being displayed in React Devtools (in the "suspended by" section).
+    // If we use `promise.then(() => resolvedValue)`, the names are lost.
+    // It's a bit strange that only one of those works right.
+    const promise = new Promise<T>((resolve, reject) => {
+      stagePromise.then(resolve.bind(null, resolvedValue), reject)
+    })
+
+    // Analogously to `makeHangingPromise`, we might reject this promise if the signal is invoked.
+    // (e.g. in the case where we don't want want the render to proceed to the dynamic stage and abort it).
+    // We shouldn't consider this an unhandled rejection, so we attach a noop catch handler here to suppress this warning.
+    if (this.abortSignal) {
+      promise.catch(ignoreReject)
+    }
+    return promise
+  }
+}
+
+function ignoreReject() {}