[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 { hasRuntimePrefetchInLoaderTree } from './prefetch-validation'
 
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
@@ -2702,11 +2703,20 @@ async function renderWithRestartOnCacheMissInDev(
   getPayload: (requestStore: RequestStore) => Promise<RSCPayload>,
   onError: (error: unknown) => void
 ) {
-  const { renderOpts } = ctx
+  const {
+    renderOpts,
+    componentMod: {
+      routeModule: {
+        userland: { loaderTree },
+      },
+    },
+  } = ctx
   const { clientReferenceManifest, ComponentMod, setReactDebugChannel } =
     renderOpts
   assertClientReferenceManifest(clientReferenceManifest)
 
+  const hasRuntimePrefetch = await hasRuntimePrefetchInLoaderTree(loaderTree)
+
   // If the render is restarted, we'll recreate a fresh request store
   let requestStore: RequestStore = initialRequestStore
 
@@ -2716,8 +2726,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:
@@ -2736,6 +2748,7 @@ async function renderWithRestartOnCacheMissInDev(
   // This render might end up being used as a prospective render (if there's cache misses),
   // so we need to set it up for filling caches.
   const cacheSignal = new CacheSignal()
+  const hangingCacheAbortController = new AbortController()
 
   // If we encounter async modules that delay rendering, we'll also need to restart.
   // TODO(restart-on-cache-miss): technically, we only need to wait for pending *server* modules here,
@@ -2745,9 +2758,9 @@ async function renderWithRestartOnCacheMissInDev(
   const prerenderResumeDataCache = createPrerenderResumeDataCache()
 
   const initialReactController = new AbortController()
-  const initialDataController = new AbortController() // Controls hanging promises we create
+  const initialHangingPromiseController = new AbortController()
   const initialStageController = new StagedRenderingController(
-    initialDataController.signal
+    initialHangingPromiseController.signal
   )
 
   requestStore.prerenderResumeDataCache = prerenderResumeDataCache
@@ -2756,10 +2769,15 @@ async function renderWithRestartOnCacheMissInDev(
   requestStore.renderResumeDataCache = null
   requestStore.stagedRendering = initialStageController
   requestStore.cacheSignal = cacheSignal
+  requestStore.hangingCacheAbortSignal = hangingCacheAbortController.signal
+  requestStore.hangingPromiseAbortSignal =
+    initialHangingPromiseController.signal
 
   let debugChannel = setReactDebugChannel && createDebugChannel()
 
   const initialRscPayload = await getPayload(requestStore)
+
+  let hadCacheMissInPreviousStages = false
   const maybeInitialServerStream = await workUnitAsyncStorage.run(
     requestStore,
     () =>
@@ -2781,36 +2799,58 @@ async function renderWithRestartOnCacheMissInDev(
           // 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)
+            const { reason } = initialReactController.signal
+            initialHangingPromiseController.abort(reason)
+
+            // We likely aborted hanging caches already (before waiting for cacheReady()),
+            // But if we haven't, we might as well settle the promises that are waiting for this signal.
+            if (!hangingCacheAbortController.signal.aborted) {
+              hangingCacheAbortController.abort(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()) {
+          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
         }
@@ -2829,11 +2869,13 @@ 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.
+  // Signal to caches that are still blocked on some stage that we're not reaching it
+  // and the cache reads should end. This avoids a deadlock where reads have started
+  // but never end because a cache is waiting for us to reach a certain stage
+  if (initialStageController.currentStage !== RenderStage.Dynamic) {
+    // TODO(restart-on-cache-miss): What if we got advanced to dynamic because of sync IO?
+    hangingCacheAbortController.abort()
+  }
 
   await cacheSignal.cacheReady()
   initialReactController.abort()

packages/next/src/server/app-render/prefetch-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 hasRuntimePrefetchInLoaderTree(
+  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 hasRuntimePrefetchInLoaderTree(parallelRoute)
+    if (hasChildRuntimePrefetch) {
+      return true
+    }
+  }
+
+  return false
+}

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

@@ -52,23 +52,27 @@ export class StagedRenderingController {
     }
   }
 
-  delayUntilStage<T>(stage: NonStaticRenderStage, resolvedValue: T) {
-    let stagePromise: Promise<void>
+  private getStagePromise(stage: NonStaticRenderStage): Promise<void> {
     switch (stage) {
       case RenderStage.Runtime: {
-        stagePromise = this.runtimeStagePromise.promise
-        break
+        return this.runtimeStagePromise.promise
       }
       case RenderStage.Dynamic: {
-        stagePromise = this.dynamicStagePromise.promise
-        break
+        return this.dynamicStagePromise.promise
       }
       default: {
         stage satisfies never
         throw new InvariantError(`Invalid render stage: ${stage}`)
       }
     }
+  }
 
+  waitForStage(stage: NonStaticRenderStage) {
+    return this.getStagePromise(stage)
+  }
+
+  delayUntilStage<T>(stage: NonStaticRenderStage, resolvedValue: T) {
+    const stagePromise = this.getStagePromise(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.

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

@@ -70,8 +70,21 @@ export interface RequestStore extends CommonWorkUnitStore {
   usedDynamic?: boolean
   devFallbackParams?: OpaqueFallbackRouteParams | null
   stagedRendering?: StagedRenderingController | null
+  // Used in the cache warmup dev render (all of these should be present)
   cacheSignal?: CacheSignal | null
   prerenderResumeDataCache?: PrerenderResumeDataCache | null
+  /**
+   * Fired if we're not going to advance the render to the end (i.e. we remain in an earlier stage).
+   * This lets us end reads for caches that are waiting for a stage and would otherwise
+   * never finish their cache reads and thus block `cacheSignal.cacheReady()`.
+   * If it fires, it should always be before aborting the render or `hangingPromiseAbortSignal`.
+   * */
+  hangingCacheAbortSignal?: AbortSignal | null
+  /**
+   * Fired after a render has been aborted to reject hanging promises.
+   * Analogous to `PrerenderStoreModern.renderSignal`.
+   */
+  hangingPromiseAbortSignal?: AbortSignal | null
 }
 
 /**

packages/next/src/server/request/search-params.ts

@@ -454,9 +454,20 @@ function makeUntrackedSearchParamsWithDevWarnings(
     requestStore,
     RenderStage.Runtime
   )
-  promise.then(() => {
-    promiseInitialized = true
-  })
+  promise.then(
+    () => {
+      promiseInitialized = true
+    },
+    // If we're in staged rendering, this promise will reject if the render
+    // is aborted before it can reach the runtime stage.
+    // In that case, we have to prevent an unhandled rejection from the promise
+    // created by this `.then()` call.
+    // This does not affect the `promiseInitialized` logic above,
+    // because `proxiedUnderlying` will not be used to resolve the promise,
+    // so there's no risk of any of its properties being accessed and triggering
+    // an undesireable warning.
+    ignoreReject
+  )
 
   Object.keys(underlyingSearchParams).forEach((prop) => {
     if (wellKnownProperties.has(prop)) {
@@ -525,6 +536,8 @@ function makeUntrackedSearchParamsWithDevWarnings(
   return proxiedPromise
 }
 
+function ignoreReject() {}
+
 const warnForSyncAccess = createDedupedByCallsiteServerErrorLoggerDev(
   createSearchAccessError
 )