[Cache Components] Less cache warming in dev if runtime prefetch not enabled

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

@@ -212,6 +212,7 @@ import { createPromiseWithResolvers } from '../../shared/lib/promise-with-resolv
 import { ImageConfigContext } from '../../shared/lib/image-config-context.shared-runtime'
 import { imageConfigDefault } from '../../shared/lib/image-config'
 import { RenderStage, StagedRenderingController } from './staged-rendering'
+import { anySegmentHasRuntimePrefetchEnabled } from './staged-validation'
 
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
@@ -2768,7 +2769,16 @@ async function renderWithRestartOnCacheMissInDev(
   getPayload: (requestStore: RequestStore) => Promise<RSCPayload>,
   onError: (error: unknown) => void
 ) {
-  const { htmlRequestId, renderOpts, requestId } = ctx
+  const {
+    htmlRequestId,
+    renderOpts,
+    requestId,
+    componentMod: {
+      routeModule: {
+        userland: { loaderTree },
+      },
+    },
+  } = ctx
   const {
     clientReferenceManifest,
     ComponentMod,
@@ -2777,6 +2787,9 @@ async function renderWithRestartOnCacheMissInDev(
   } = renderOpts
   assertClientReferenceManifest(clientReferenceManifest)
 
+  const hasRuntimePrefetch =
+    await anySegmentHasRuntimePrefetchEnabled(loaderTree)
+
   // If the render is restarted, we'll recreate a fresh request store
   let requestStore: RequestStore = initialRequestStore
 
@@ -2786,8 +2799,10 @@ async function renderWithRestartOnCacheMissInDev(
       case RenderStage.Static:
         return 'Prerender'
       case RenderStage.Runtime:
-        // TODO: only label as "Prefetch" if the page has a `prefetch` config.
-        return 'Prefetch'
+        // If we're not warming caches reachable in the runtime phase,
+        // we can't trust the "Prefetch" labelling to be correct
+        // due to potential delays caused by cache misses.
+        return hasRuntimePrefetch ? 'Prefetch' : 'Server'
       case RenderStage.Dynamic:
         return 'Server'
       default:
@@ -2836,6 +2851,8 @@ async function renderWithRestartOnCacheMissInDev(
   let debugChannel = setReactDebugChannel && createDebugChannel()
 
   const initialRscPayload = await getPayload(requestStore)
+
+  let hadCacheMissInPreviousStages = false
   const maybeInitialServerStream = await workUnitAsyncStorage.run(
     requestStore,
     () =>
@@ -2863,30 +2880,45 @@ async function renderWithRestartOnCacheMissInDev(
         },
         (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()) {
+          hadCacheMissInPreviousStages = cacheSignal.hasPendingReads()
+
+          // If runtime prefetching isn't enabled for any segment in this page,
+          // then we don't need to validate anything in the runtime phase.
+          // Thus, there's no need for us to warm runtime caches.
+          // We can avoid advancing to the runtime stage and unblocking runtime APIs during the warmup,
+          // which'll make it faster (because we won't wait for any caches hidden behind `await cookies()` etc)
+          if (!hasRuntimePrefetch && hadCacheMissInPreviousStages) {
             return null
           }
 
-          // If there's no cache misses, we'll continue rendering,
-          // and see if there's any cache misses in the runtime stage.
+          // If there's no cache misses, we'll continue rendering.
+          initialStageController.advanceStage(RenderStage.Runtime)
           return stream
         },
-        async (maybeStream) => {
+        (maybeStream) => {
           // Dynamic stage
 
-          // If we had cache misses in either of the previous stages,
-          // then we'll only use this render for filling caches.
+          // If the previous stage bailed out of the render due to a cache miss,
+          // we shouldn't do anything more.
+          if (maybeStream === null) {
+            return null
+          }
+
+          hadCacheMissInPreviousStages ||= cacheSignal.hasPendingReads()
+
+          // If runtime prefetching is enabled for any segment in this page,
+          // then we need a proper runtime phase for validation.
+          // Thus, if we had cache misses in either of the previous stages,
+          // We have to bail out and warm all the caches before retrying.
           // 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()) {
+          if (hasRuntimePrefetch && hadCacheMissInPreviousStages) {
             return null
           }
 
-          // If there's no cache misses, we'll use this render, so let it advance to the dynamic stage.
+          // If we didn't bail out earlier, that means there's no cache misses,
+          // and we use this render. Let it advance to the dynamic stage.
           initialStageController.advanceStage(RenderStage.Dynamic)
           return maybeStream
         }
@@ -2909,12 +2941,7 @@ async function renderWithRestartOnCacheMissInDev(
   // Cache miss. We will use the initial render to fill caches, and discard its result.
   // Then, we can render again with warm caches.
 
-  // TODO(restart-on-cache-miss):
-  // This might end up waiting for more caches than strictly necessary,
-  // because we can't abort the render yet, and we'll let runtime/dynamic APIs resolve.
-  // Ideally we'd only wait for caches that are needed in the static stage.
-  // This will be optimized in the future by not allowing runtime/dynamic APIs to resolve.
-
+  // TODO: potential deadlock if we started reads for caches delayed until runtime/dynamic
   await cacheSignal.cacheReady()
   initialReactController.abort()
 

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

@@ -0,0 +1,32 @@
+import { getLayoutOrPageModule } from '../lib/app-dir-module'
+import type { LoaderTree } from '../lib/app-dir-module'
+import { parseLoaderTree } from '../../shared/lib/router/utils/parse-loader-tree'
+import type { AppSegmentConfig } from '../../build/segment-config/app/app-segment-config'
+
+export async function anySegmentHasRuntimePrefetchEnabled(
+  tree: LoaderTree
+): Promise<boolean> {
+  const { mod: layoutOrPageMod } = await getLayoutOrPageModule(tree)
+
+  // TODO(restart-on-cache-miss): Does this work correctly for client page/layout modules?
+  const prefetchConfig = layoutOrPageMod
+    ? (layoutOrPageMod as AppSegmentConfig).unstable_prefetch
+    : undefined
+  /** Whether this segment should use a runtime prefetch instead of a static prefetch. */
+  const hasRuntimePrefetch = prefetchConfig?.mode === 'runtime'
+  if (hasRuntimePrefetch) {
+    return true
+  }
+
+  const { parallelRoutes } = parseLoaderTree(tree)
+  for (const parallelRouteKey in parallelRoutes) {
+    const parallelRoute = parallelRoutes[parallelRouteKey]
+    const hasChildRuntimePrefetch =
+      await anySegmentHasRuntimePrefetchEnabled(parallelRoute)
+    if (hasChildRuntimePrefetch) {
+      return true
+    }
+  }
+
+  return false
+}