staged render for initial load

Author: lubieowoce

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

@@ -607,26 +607,6 @@ async function generateDynamicFlightRenderResult(
     onFlightDataRenderError
   )
 
-  const RSCPayload: RSCPayload & {
-    /** Only available during cacheComponents development builds. Used for logging errors. */
-    _validation?: Promise<ReactNode>
-    _bypassCachesInDev?: React.ReactNode
-  } = await workUnitAsyncStorage.run(
-    requestStore,
-    generateDynamicRSCPayload,
-    ctx,
-    options
-  )
-
-  if (
-    process.env.NODE_ENV === 'development' &&
-    isBypassingCachesInDev(renderOpts, requestStore)
-  ) {
-    RSCPayload._bypassCachesInDev = createElement(WarnForBypassCachesInDev, {
-      route: workStore.route,
-    })
-  }
-
   const debugChannel = setReactDebugChannel && createDebugChannel()
 
   if (debugChannel) {
@@ -635,63 +615,60 @@ async function generateDynamicFlightRenderResult(
 
   let stream: ReadableStream<Uint8Array<ArrayBufferLike>>
   if (process.env.NODE_ENV === 'development' && renderOpts.cacheComponents) {
-    // If we got here in Cache Components dev, we're rendering while bypassing caches,
-    // so we have no hope of showing a useful runtime stage.
-    // But we still want things like `params` to show up in devtools correctly,
-    // which relies on mechanisms we've set up for staged rendering,
-    // so we do a 2-task version (Static -> Dynamic) instead.
-
-    const stageController = new StagedRenderingController()
-
-    const environmentName = () => {
-      const currentStage = stageController.currentStage
-      switch (currentStage) {
-        case RenderStage.Static:
-          return 'Prerender'
-        case RenderStage.Runtime:
-        case RenderStage.Dynamic:
-          return 'Server'
-        default:
-          currentStage satisfies never
-          throw new InvariantError(`Invalid render stage: ${currentStage}`)
+    // If we got here, it means we didn't go into `generateDynamicFlightRenderResultWithCachesInDev`,
+    // so caches must be disabled. Render dynamically, but include basic environment labels.
+
+    const getPayload = async (
+      // eslint-disable-next-line @typescript-eslint/no-shadow
+      requestStore: RequestStore
+    ): Promise<RSCPayload> => {
+      const rscPayload: RSCPayload & {
+        /** Only available during cacheComponents development builds. Used for logging errors. */
+        _validation?: Promise<ReactNode>
+        _bypassCachesInDev?: React.ReactNode
+      } = await workUnitAsyncStorage.run(
+        requestStore,
+        generateDynamicRSCPayload,
+        ctx,
+        options
+      )
+      if (isBypassingCachesInDev(renderOpts, requestStore)) {
+        rscPayload._bypassCachesInDev = createElement(
+          WarnForBypassCachesInDev,
+          {
+            route: workStore.route,
+          }
+        )
       }
+      return rscPayload
     }
 
-    requestStore.stagedRendering = stageController
-    requestStore.asyncApiPromises = createAsyncApiPromisesInDev(
-      stageController,
-      requestStore.cookies,
-      requestStore.mutableCookies,
-      requestStore.headers
-    )
-
-    stream = await workUnitAsyncStorage.run(
+    stream = await stagedRenderToReadableStreamWithoutCachesInDev(
+      ctx,
       requestStore,
-      scheduleInSequentialTasks,
-      () => {
-        return renderToReadableStream(
-          RSCPayload,
-          clientReferenceManifest.clientModules,
-          {
-            onError,
-            temporaryReferences: options?.temporaryReferences,
-            filterStackFrame,
-            debugChannel: debugChannel?.serverSide,
-            environmentName,
-          }
-        )
-      },
-      () => {
-        stageController.advanceStage(RenderStage.Dynamic)
+      getPayload,
+      clientReferenceManifest,
+      {
+        onError,
+        temporaryReferences: options?.temporaryReferences,
+        filterStackFrame,
+        debugChannel: debugChannel?.serverSide,
       }
     )
   } else {
     // For app dir, use the bundled version of Flight server renderer (renderToReadableStream)
     // which contains the subset React.
+    const rscPayload = await workUnitAsyncStorage.run(
+      requestStore,
+      generateDynamicRSCPayload,
+      ctx,
+      options
+    )
+
     stream = workUnitAsyncStorage.run(
       requestStore,
       renderToReadableStream,
-      RSCPayload,
+      rscPayload,
       clientReferenceManifest.clientModules,
       {
         onError,
@@ -707,6 +684,75 @@ async function generateDynamicFlightRenderResult(
   })
 }
 
+type RenderToReadableStreamServerOptions = NonNullable<
+  Parameters<
+    (typeof import('react-server-dom-webpack/server.node'))['renderToReadableStream']
+  >[2]
+>
+
+async function stagedRenderToReadableStreamWithoutCachesInDev(
+  ctx: AppRenderContext,
+  requestStore: RequestStore,
+  getPayload: (requestStore: RequestStore) => Promise<RSCPayload>,
+  clientReferenceManifest: NonNullable<RenderOpts['clientReferenceManifest']>,
+  options: Omit<RenderToReadableStreamServerOptions, 'environmentName'>
+) {
+  const {
+    componentMod: { renderToReadableStream },
+  } = ctx
+  // We're rendering while bypassing caches,
+  // so we have no hope of showing a useful runtime stage.
+  // But we still want things like `params` to show up in devtools correctly,
+  // which relies on mechanisms we've set up for staged rendering,
+  // so we do a 2-task version (Static -> Dynamic) instead.
+
+  const stageController = new StagedRenderingController()
+  console.log('rendering RSC payload without caches')
+  const environmentName = () => {
+    const currentStage = stageController.currentStage
+    switch (currentStage) {
+      case RenderStage.Static:
+        return 'Prerender'
+      case RenderStage.Runtime:
+      case RenderStage.Dynamic:
+        return 'Server'
+      default:
+        currentStage satisfies never
+        throw new InvariantError(`Invalid render stage: ${currentStage}`)
+    }
+  }
+
+  requestStore.stagedRendering = stageController
+  requestStore.asyncApiPromises = createAsyncApiPromisesInDev(
+    stageController,
+    requestStore.cookies,
+    requestStore.mutableCookies,
+    requestStore.headers
+  )
+
+  const rscPayload = await getPayload(requestStore)
+
+  return await workUnitAsyncStorage.run(
+    requestStore,
+    scheduleInSequentialTasks,
+    () => {
+      const stream = renderToReadableStream(
+        rscPayload,
+        clientReferenceManifest.clientModules,
+        {
+          ...options,
+          environmentName,
+        }
+      )
+      stageController.advanceStage(RenderStage.Dynamic)
+      return stream
+    },
+    () => {
+      stageController.advanceStage(RenderStage.Dynamic)
+    }
+  )
+}
+
 /**
  * Fork of `generateDynamicFlightRenderResult` that renders using `renderWithRestartOnCacheMissInDev`
  * to ensure correct separation of environments Prerender/Server (for use in Cache Components)
@@ -2375,21 +2421,15 @@ async function renderToStream(
       // Edge routes never prerender so we don't have a Prerender environment for anything in edge runtime
       process.env.NEXT_RUNTIME !== 'edge' &&
       // We only have a Prerender environment for projects opted into cacheComponents
-      cacheComponents &&
-      // We only do this flow if we can safely recreate the store from scratch
-      // (which is not the case for renders after an action)
-      createRequestStore &&
-      // We only do this flow if we're not bypassing caches in dev using
-      // "disable cache" in devtools or a hard refresh (cache-control: "no-store")
-      !isBypassingCachesInDev(renderOpts, requestStore)
+      cacheComponents
     ) {
       type RSCPayloadWithValidation = InitialRSCPayload & {
         /** Only available during cacheComponents development builds. Used for logging errors. */
         _validation?: Promise<ReactNode>
       }
 
       const [resolveValidation, validationOutlet] = createValidationOutlet()
-
+      let debugChannel: DebugChannelPair | undefined
       const getPayload = async (
         // eslint-disable-next-line @typescript-eslint/no-shadow
         requestStore: RequestStore
@@ -2410,20 +2450,47 @@ async function renderToStream(
         return payload
       }
 
-      const {
-        stream: serverStream,
-        debugChannel,
-        requestStore: finalRequestStore,
-      } = await renderWithRestartOnCacheMissInDev(
-        ctx,
-        requestStore,
-        createRequestStore,
-        getPayload,
-        serverComponentsErrorHandler
-      )
+      if (
+        // We only do this flow if we can safely recreate the store from scratch
+        // (which is not the case for renders after an action)
+        createRequestStore &&
+        // We only do this flow if we're not bypassing caches in dev using
+        // "disable cache" in devtools or a hard refresh (cache-control: "no-store")
+        !isBypassingCachesInDev(renderOpts, requestStore)
+      ) {
+        const {
+          stream: serverStream,
+          debugChannel: returnedDebugChannel,
+          requestStore: finalRequestStore,
+        } = await renderWithRestartOnCacheMissInDev(
+          ctx,
+          requestStore,
+          createRequestStore,
+          getPayload,
+          serverComponentsErrorHandler
+        )
 
-      reactServerResult = new ReactServerResult(serverStream)
-      requestStore = finalRequestStore
+        reactServerResult = new ReactServerResult(serverStream)
+        requestStore = finalRequestStore
+        debugChannel = returnedDebugChannel
+      } else {
+        // We're either bypassing caches or we can't restart the render.
+        // Do a dynamic render, but with (basic) environment labels.
+        debugChannel = createDebugChannel && createDebugChannel()
+        const serverStream =
+          await stagedRenderToReadableStreamWithoutCachesInDev(
+            ctx,
+            requestStore,
+            getPayload,
+            clientReferenceManifest,
+            {
+              onError: serverComponentsErrorHandler,
+              filterStackFrame,
+              debugChannel: debugChannel?.serverSide,
+            }
+          )
+        reactServerResult = new ReactServerResult(serverStream)
+      }
 
       if (debugChannel && setReactDebugChannel) {
         const [readableSsr, readableBrowser] =