[Cache Components] separate runtime stage in dev render

Author: lubieowoce

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 pipelineInSequentialTasks3<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 { pipelineInSequentialTasks3 } from './app-render-render-utils'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import {
   workUnitAsyncStorage,
@@ -209,6 +206,7 @@ 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'
+import { RenderStage, StagedRenderingController } from './staged-rendering'
 
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
@@ -2175,8 +2173,21 @@ async function renderToStream(
         )
       }
 
-      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}`)
+        }
+      }
 
       // Try to render the page and see if there's any cache misses.
       // If there are, wait for caches to finish and restart the render.
@@ -2192,26 +2203,30 @@ async function renderToStream(
 
       const prerenderResumeDataCache = createPrerenderResumeDataCache()
 
+      const initialRenderReactController = new AbortController() // Controls the react render
+      const initialRenderDataController = new AbortController() // Controls hanging promises we create
+      const initialRenderStageController = new StagedRenderingController(
+        initialRenderDataController.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 = initialRenderStageController
       requestStore.cacheSignal = cacheSignal
 
-      const initialRenderReactController = new AbortController()
-
       const intialRenderDebugChannel =
         setReactDebugChannel && createDebugChannel()
 
       const initialRscPayload = await getPayload()
       const maybeInitialServerStream = await workUnitAsyncStorage.run(
         requestStore,
         () =>
-          pipelineInSequentialTasks(
+          pipelineInSequentialTasks3(
             () => {
               // Static stage
-              requestStore.prerenderPhase = true
-              return ComponentMod.renderToReadableStream(
+              const stream = ComponentMod.renderToReadableStream(
                 initialRscPayload,
                 clientReferenceManifest.clientModules,
                 {
@@ -2222,25 +2237,57 @@ async function renderToStream(
                   signal: initialRenderReactController.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.
+              initialRenderReactController.signal.addEventListener(
+                'abort',
+                () => {
+                  initialRenderDataController.abort(
+                    initialRenderReactController.signal.reason
+                  )
+                }
+              )
+              return stream
             },
-            async (stream) => {
-              // Dynamic stage
-              // Note: if we had cache misses, things that would've happened statically otherwise
-              // may be marked as dynamic instead.
-              requestStore.prerenderPhase = false
+            (stream) => {
+              // Runtime stage
+              initialRenderStageController.advanceStage(RenderStage.Runtime)
 
               // 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) {
+              if (!cacheSignal.hasPendingReads()) {
                 // No cache misses. We can use the stream as is.
                 return stream
               } else {
                 // Cache miss. We'll discard this stream, and render again.
                 return null
               }
+            },
+            async (maybeStream) => {
+              // Dynamic stage
+
+              if (maybeStream === null) {
+                // 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.
+                return null
+              }
+
+              // Note: if we had cache misses, things that would've happened statically otherwise
+              // may be marked as dynamic instead.
+              initialRenderStageController.advanceStage(RenderStage.Dynamic)
+
+              // Analogous to the previous stage.
+              if (!cacheSignal.hasPendingReads()) {
+                // No cache misses. We can use the stream as is.
+                return maybeStream
+              } else {
+                // Cache miss. We'll discard this stream, and render again.
+                return null
+              }
             }
           )
       )
@@ -2273,11 +2320,14 @@ async function renderToStream(
         // Now, we need to do another render.
         requestStore = createRequestStore()
 
+        const finalRenderStageController = new StagedRenderingController()
+
         // We've filled the caches, so now we can render as usual.
         requestStore.prerenderResumeDataCache = null
         requestStore.renderResumeDataCache = createRenderResumeDataCache(
           prerenderResumeDataCache
         )
+        requestStore.stagedRendering = finalRenderStageController
         requestStore.cacheSignal = null
 
         // The initial render already wrote to its debug channel. We're not using it,
@@ -2292,25 +2342,32 @@ async function renderToStream(
         const finalRscPayload = await getPayload()
         const finalServerStream = await workUnitAsyncStorage.run(
           requestStore,
-          scheduleInSequentialTasks,
-          () => {
-            // Static stage
-            requestStore.prerenderPhase = true
-            return ComponentMod.renderToReadableStream(
-              finalRscPayload,
-              clientReferenceManifest.clientModules,
-              {
-                onError: serverComponentsErrorHandler,
-                environmentName,
-                filterStackFrame,
-                debugChannel: finalRenderDebugChannel?.serverSide,
+          () =>
+            pipelineInSequentialTasks3(
+              () => {
+                // Static stage
+                return ComponentMod.renderToReadableStream(
+                  finalRscPayload,
+                  clientReferenceManifest.clientModules,
+                  {
+                    onError: serverComponentsErrorHandler,
+                    environmentName,
+                    filterStackFrame,
+                    debugChannel: finalRenderDebugChannel?.serverSide,
+                  }
+                )
+              },
+              (stream) => {
+                // Runtime stage
+                finalRenderStageController.advanceStage(RenderStage.Runtime)
+                return stream
+              },
+              (stream) => {
+                // Dynamic stage
+                finalRenderStageController.advanceStage(RenderStage.Dynamic)
+                return stream
               }
             )
-          },
-          () => {
-            // Dynamic stage
-            requestStore.prerenderPhase = false
-          }
         )
 
         reactServerResult = new ReactServerResult(finalServerStream)

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,11 @@ 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: this doesn't seem like it'll actually do what we need?
+    requestStore.stagedRendering.advanceStage(RenderStage.Dynamic)
+  }
 }
 
 /**

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

@@ -0,0 +1,67 @@
+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(abortSignal?: AbortSignal) {
+    if (abortSignal) {
+      abortSignal.addEventListener(
+        'abort',
+        () => {
+          const { reason } = abortSignal
+          if (this.currentStage < RenderStage.Runtime) {
+            this.runtimeStagePromise.reject(reason)
+          }
+          if (this.currentStage < RenderStage.Dynamic) {
+            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()
+    }
+  }
+
+  waitForStage(stage: NonStaticRenderStage) {
+    switch (stage) {
+      case RenderStage.Runtime: {
+        return this.runtimeStagePromise.promise
+      }
+      case RenderStage.Dynamic: {
+        return this.dynamicStagePromise.promise
+      }
+      default: {
+        stage satisfies never
+        throw new InvariantError(`Invalid render stage: ${stage}`)
+      }
+    }
+  }
+}

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

@@ -19,6 +19,7 @@ import type { ImplicitTags } from '../lib/implicit-tags'
 import type { WorkStore } from './work-async-storage.external'
 import { NEXT_HMR_REFRESH_HASH_COOKIE } from '../../client/components/app-router-headers'
 import { InvariantError } from '../../shared/lib/invariant-error'
+import type { StagedRenderingController } from './staged-rendering'
 
 export type WorkUnitPhase = 'action' | 'render' | 'after'
 
@@ -67,8 +68,8 @@ export interface RequestStore extends CommonWorkUnitStore {
 
   // DEV-only
   usedDynamic?: boolean
-  prerenderPhase?: boolean
   devFallbackParams?: OpaqueFallbackRouteParams | null
+  stagedRendering?: StagedRenderingController | null
   cacheSignal?: CacheSignal | null
   prerenderResumeDataCache?: PrerenderResumeDataCache | null
 }