[Cache Components] set environment labels if caches are disabled

Author: lubieowoce

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

@@ -168,7 +168,10 @@ import {
   prerenderAndAbortInSequentialTasks,
 } from './app-render-prerender-utils'
 import { printDebugThrownValueForProspectiveRender } from './prospective-render-utils'
-import { pipelineInSequentialTasks } from './app-render-render-utils'
+import {
+  pipelineInSequentialTasks,
+  scheduleInSequentialTasks,
+} from './app-render-render-utils'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import {
   workUnitAsyncStorage,
@@ -604,52 +607,149 @@ 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) {
     setReactDebugChannel(debugChannel.clientSide, htmlRequestId, requestId)
   }
 
-  // For app dir, use the bundled version of Flight server renderer (renderToReadableStream)
-  // which contains the subset React.
-  const flightReadableStream = workUnitAsyncStorage.run(
-    requestStore,
-    renderToReadableStream,
-    RSCPayload,
-    clientReferenceManifest.clientModules,
-    {
-      onError,
-      temporaryReferences: options?.temporaryReferences,
-      filterStackFrame,
-      debugChannel: debugChannel?.serverSide,
+  let stream: ReadableStream<Uint8Array<ArrayBufferLike>>
+  if (process.env.NODE_ENV === 'development' && renderOpts.cacheComponents) {
+    // 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
+    ) => {
+      const rscPayload: RSCPayload & RSCPayloadDevProperties =
+        await workUnitAsyncStorage.run(
+          requestStore,
+          generateDynamicRSCPayload,
+          ctx,
+          options
+        )
+      if (isBypassingCachesInDev(renderOpts, requestStore)) {
+        rscPayload._bypassCachesInDev = createElement(
+          WarnForBypassCachesInDev,
+          {
+            route: workStore.route,
+          }
+        )
+      }
+      return rscPayload
     }
-  )
 
-  return new FlightRenderResult(flightReadableStream, {
+    stream = await stagedRenderToReadableStreamWithoutCachesInDev(
+      ctx,
+      requestStore,
+      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,
+      clientReferenceManifest.clientModules,
+      {
+        onError,
+        temporaryReferences: options?.temporaryReferences,
+        filterStackFrame,
+        debugChannel: debugChannel?.serverSide,
+      }
+    )
+  }
+
+  return new FlightRenderResult(stream, {
     fetchMetrics: workStore.fetchMetrics,
   })
 }
 
+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)
@@ -2192,6 +2292,12 @@ function applyMetadataFromPrerenderResult(
   }
 }
 
+type RSCPayloadDevProperties = {
+  /** Only available during cacheComponents development builds. Used for logging errors. */
+  _validation?: Promise<ReactNode>
+  _bypassCachesInDev?: ReactNode
+}
+
 async function renderToStream(
   requestStore: RequestStore,
   req: BaseNextRequest,
@@ -2334,26 +2440,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
-      ): Promise<RSCPayloadWithValidation> => {
-        const payload: RSCPayloadWithValidation =
+      ) => {
+        const payload: InitialRSCPayload & RSCPayloadDevProperties =
           await workUnitAsyncStorage.run(
             requestStore,
             getRSCPayload,
@@ -2366,23 +2461,61 @@ async function renderToStream(
         // because we're not going to wait for the stream to complete,
         // so leaving the validation unresolved is fine.
         payload._validation = validationOutlet
+
+        if (isBypassingCachesInDev(renderOpts, requestStore)) {
+          // Mark the RSC payload to indicate that caches were bypassed in dev.
+          // This lets the client know not to cache anything based on this render.
+          payload._bypassCachesInDev = createElement(WarnForBypassCachesInDev, {
+            route: workStore.route,
+          })
+        }
+
         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
+        debugChannel = returnedDebugChannel
+      } else {
+        // We're either bypassing caches or we can't restart the render.
+        // Do a dynamic render, but with (basic) environment labels.
 
-      reactServerResult = new ReactServerResult(serverStream)
-      requestStore = finalRequestStore
+        debugChannel = setReactDebugChannel && 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] =
@@ -2413,15 +2546,14 @@ async function renderToStream(
       )
     } else {
       // This is a dynamic render. We don't do dynamic tracking because we're not prerendering
-      const RSCPayload: RSCPayload & {
-        _bypassCachesInDev?: React.ReactNode
-      } = await workUnitAsyncStorage.run(
-        requestStore,
-        getRSCPayload,
-        tree,
-        ctx,
-        res.statusCode === 404
-      )
+      const RSCPayload: RSCPayload & RSCPayloadDevProperties =
+        await workUnitAsyncStorage.run(
+          requestStore,
+          getRSCPayload,
+          tree,
+          ctx,
+          res.statusCode === 404
+        )
 
       if (isBypassingCachesInDev(renderOpts, requestStore)) {
         // Mark the RSC payload to indicate that caches were bypassed in dev.