[dynamicIO] Implement a warmup prefetch render in dev (vercel#71278)

To seed caches in dev mode we implemented a prefetch if one had not been
done in the recent past. The implementation isn't quite working b/c it
doesn't match the prefetch header properly but the mechanics are there.
One problem however is the prefetch streams so we end up starting the
regular render before the prefetch render is complete.

The ideal cache warmup is to only render until there are no more caches
to fill in the prefetch. Additionally we can disable certain things like
dynamic Request apis and fetches so they stall forever too. This
suggests that there is enough of a difference in the needs of the
prefetch for cache warming that we ought to implement it as it's own
request type. This change implements the mechanics of triggering the
warmup prefetch but does not yet implement the changes to things like
the request or fetch apis behavior. I will follow up with further
changes on top of this commit.

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

@@ -157,6 +157,7 @@ import {
 } from './work-unit-async-storage.external'
 import { CacheSignal } from './cache-signal'
 import { getTracedMetadata } from '../lib/trace/utils'
+import { InvariantError } from '../../shared/lib/invariant-error'
 
 import './clean-async-snapshot.external'
 import { INFINITE_CACHE } from '../../lib/constants'
@@ -196,6 +197,7 @@ export type AppRenderContext = {
 }
 
 interface ParseRequestHeadersOptions {
+  readonly isDevWarmup: undefined | boolean
   readonly isRoutePPREnabled: boolean
 }
 
@@ -210,6 +212,7 @@ interface ParsedRequestHeaders {
    */
   readonly flightRouterState: FlightRouterState | undefined
   readonly isPrefetchRequest: boolean
+  readonly isDevWarmupRequest: boolean
   readonly isHmrRefresh: boolean
   readonly isRSCRequest: boolean
   readonly nonce: string | undefined
@@ -219,13 +222,19 @@ function parseRequestHeaders(
   headers: IncomingHttpHeaders,
   options: ParseRequestHeadersOptions
 ): ParsedRequestHeaders {
+  const isDevWarmupRequest = options.isDevWarmup === true
+
+  // dev warmup requests are treated as prefetch RSC requests
   const isPrefetchRequest =
+    isDevWarmupRequest ||
     headers[NEXT_ROUTER_PREFETCH_HEADER.toLowerCase()] !== undefined
 
   const isHmrRefresh =
     headers[NEXT_HMR_REFRESH_HEADER.toLowerCase()] !== undefined
 
-  const isRSCRequest = headers[RSC_HEADER.toLowerCase()] !== undefined
+  // dev warmup requests are treated as prefetch RSC requests
+  const isRSCRequest =
+    isDevWarmupRequest || headers[RSC_HEADER.toLowerCase()] !== undefined
 
   const shouldProvideFlightRouterState =
     isRSCRequest && (!isPrefetchRequest || !options.isRoutePPREnabled)
@@ -248,6 +257,7 @@ function parseRequestHeaders(
     isPrefetchRequest,
     isHmrRefresh,
     isRSCRequest,
+    isDevWarmupRequest,
     nonce,
   }
 }
@@ -529,6 +539,69 @@ async function generateDynamicFlightRenderResult(
   })
 }
 
+/**
+ * Performs a "warmup" render of the RSC payload for a given route. This function is called by the server
+ * prior to an actual render request in Dev mode only. It's purpose is to fill caches so the actual render
+ * can accurately log activity in the right render context (Prerender vs Render).
+ *
+ * At the moment this implementation is mostly a fork of generateDynamicFlightRenderResult
+ */
+async function warmupDevRender(
+  req: BaseNextRequest,
+  ctx: AppRenderContext,
+  options?: {
+    actionResult: ActionResult
+    skipFlight: boolean
+    componentTree?: CacheNodeSeedData
+    preloadCallbacks?: PreloadCallbacks
+  }
+): Promise<RenderResult> {
+  const renderOpts = ctx.renderOpts
+  if (!renderOpts.dev) {
+    throw new InvariantError(
+      'generateDynamicFlightRenderResult should never be called in `next start` mode.'
+    )
+  }
+
+  function onFlightDataRenderError(err: DigestedError) {
+    return renderOpts.onInstrumentationRequestError?.(
+      err,
+      req,
+      createErrorContext(ctx, 'react-server-components-payload')
+    )
+  }
+  const onError = createFlightReactServerErrorHandler(
+    true,
+    onFlightDataRenderError
+  )
+
+  const rscPayload = await generateDynamicRSCPayload(ctx, options)
+
+  // For app dir, use the bundled version of Flight server renderer (renderToReadableStream)
+  // which contains the subset React.
+  const flightReadableStream = ctx.componentMod.renderToReadableStream(
+    rscPayload,
+    ctx.clientReferenceManifest.clientModules,
+    {
+      onError,
+    }
+  )
+
+  const reader = flightReadableStream.getReader()
+  while (true) {
+    if ((await reader.read()).done) {
+      break
+    }
+  }
+
+  // We don't really want to return a result here but the stack of functions
+  // that calls into renderToHTML... expects a result. We should refactor this to
+  // lift the warmup pathway outside of renderToHTML... but for now this suffices
+  return new FlightRenderResult('', {
+    fetchMetrics: ctx.workStore.fetchMetrics,
+  })
+}
+
 /**
  * Crawlers will inadvertently think the canonicalUrl in the RSC payload should be crawled
  * when our intention is to just seed the router state with the current URL.
@@ -989,8 +1062,13 @@ async function renderToHTMLOrFlightImpl(
   query = { ...query }
   stripInternalQueries(query)
 
-  const { flightRouterState, isPrefetchRequest, isRSCRequest, nonce } =
-    parsedRequestHeaders
+  const {
+    flightRouterState,
+    isPrefetchRequest,
+    isRSCRequest,
+    isDevWarmupRequest,
+    nonce,
+  } = parsedRequestHeaders
 
   /**
    * The metadata items array created in next-app-loader with all relevant information
@@ -1171,7 +1249,9 @@ async function renderToHTMLOrFlightImpl(
     return new RenderResult(await streamToString(response.stream), options)
   } else {
     // We're rendering dynamically
-    if (isRSCRequest) {
+    if (isDevWarmupRequest) {
+      return warmupDevRender(req, ctx)
+    } else if (isRSCRequest) {
       return generateDynamicFlightRenderResult(req, ctx)
     }
 
@@ -1289,10 +1369,11 @@ export const renderToHTMLOrFlight: AppPageRender = (
   // We read these values from the request object as, in certain cases,
   // base-server will strip them to opt into different rendering behavior.
   const parsedRequestHeaders = parseRequestHeaders(req.headers, {
+    isDevWarmup: renderOpts.isDevWarmup,
     isRoutePPREnabled: renderOpts.experimental.isRoutePPREnabled === true,
   })
 
-  const { isHmrRefresh } = parsedRequestHeaders
+  const { isHmrRefresh, isPrefetchRequest } = parsedRequestHeaders
 
   const requestEndedState = { ended: false }
   let postponedState: PostponedState | null = null
@@ -1339,7 +1420,8 @@ export const renderToHTMLOrFlight: AppPageRender = (
         fallbackRouteParams,
         renderOpts,
         requestEndedState,
-        isPrefetchRequest: Boolean(req.headers[NEXT_ROUTER_PREFETCH_HEADER]),
+        // @TODO move to workUnitStore of type Request
+        isPrefetchRequest,
       },
       (workStore) =>
         renderToHTMLOrFlightImpl(

packages/next/src/server/app-render/types.ts

@@ -171,6 +171,7 @@ export interface RenderOptsPartial {
   }
   params?: ParsedUrlQuery
   isPrefetch?: boolean
+  isDevWarmup?: boolean
   experimental: {
     /**
      * When true, it indicates that the current page supports partial

packages/next/src/server/base-server.ts

@@ -2392,12 +2392,21 @@ export default abstract class Server<
        * The unknown route params for this render.
        */
       fallbackRouteParams: FallbackRouteParams | null
+
+      /**
+       * Whether or not this render is a warmup render for dev mode.
+       */
+      isDevWarmup?: boolean
     }
     type Renderer = (
       context: RendererContext
     ) => Promise<ResponseCacheEntry | null>
 
-    const doRender: Renderer = async ({ postponed, fallbackRouteParams }) => {
+    const doRender: Renderer = async ({
+      postponed,
+      fallbackRouteParams,
+      isDevWarmup,
+    }) => {
       // In development, we always want to generate dynamic HTML.
       let supportsDynamicResponse: boolean =
         // If we're in development, we always support dynamic HTML, unless it's
@@ -2470,6 +2479,7 @@ export default abstract class Server<
         isOnDemandRevalidate,
         isDraftMode: isPreviewMode,
         isServerAction,
+        isDevWarmup,
         postponed,
         waitUntil: this.getWaitUntil(),
         onClose: res.onClose.bind(res),
@@ -2790,6 +2800,7 @@ export default abstract class Server<
       hasResolved,
       previousCacheEntry,
       isRevalidating,
+      isDevWarmup,
     }): Promise<ResponseCacheEntry | null> => {
       const isProduction = !this.renderOpts.dev
       const didRespond = hasResolved || res.sent
@@ -3028,6 +3039,7 @@ export default abstract class Server<
       const result = await doRender({
         postponed,
         fallbackRouteParams,
+        isDevWarmup,
       })
       if (!result) return null
 
@@ -3041,7 +3053,7 @@ export default abstract class Server<
       const originalResponseGenerator = responseGenerator
 
       responseGenerator = async (
-        ...args: Parameters<typeof responseGenerator>
+        state: Parameters<ResponseGenerator>[0]
       ): ReturnType<typeof responseGenerator> => {
         if (this.renderOpts.dev) {
           let cache = this.prefetchCacheScopesDev.get(urlPathname)
@@ -3055,23 +3067,17 @@ export default abstract class Server<
             routeModule?.definition.kind === RouteKind.APP_PAGE &&
             !isServerAction
           ) {
-            req.headers[RSC_HEADER] = '1'
-            req.headers[NEXT_ROUTER_PREFETCH_HEADER] = '1'
-
             cache = new Map()
 
             await runWithCacheScope({ cache }, () =>
-              originalResponseGenerator(...args)
+              originalResponseGenerator({ ...state, isDevWarmup: true })
             )
             this.prefetchCacheScopesDev.set(urlPathname, cache)
-
-            delete req.headers[RSC_HEADER]
-            delete req.headers[NEXT_ROUTER_PREFETCH_HEADER]
           }
 
           if (cache) {
             return runWithCacheScope({ cache }, () =>
-              originalResponseGenerator(...args)
+              originalResponseGenerator(state)
             ).finally(() => {
               if (isPrefetchRSCRequest) {
                 this.prefetchCacheScopesDev.set(urlPathname, cache)
@@ -3082,7 +3088,7 @@ export default abstract class Server<
           }
         }
 
-        return originalResponseGenerator(...args)
+        return originalResponseGenerator(state)
       }
     }
 

packages/next/src/server/lib/incremental-cache/index.ts

@@ -398,6 +398,25 @@ export class IncrementalCache implements IncrementalCacheType {
       isFallback: boolean | undefined
     }
   ): Promise<IncrementalCacheEntry | null> {
+    // unlike other caches if we have a cacheScope we use it even if
+    // testmode would normally disable it or if requestHeaders say 'no-cache'.
+    if (this.hasDynamicIO && ctx.kind === IncrementalCacheKind.FETCH) {
+      const cacheScope = cacheScopeAsyncLocalStorage.getStore()
+
+      if (cacheScope) {
+        const memoryCacheData = cacheScope.cache.get(cacheKey)
+
+        if (memoryCacheData?.kind === CachedRouteKind.FETCH) {
+          return {
+            isStale: false,
+            value: memoryCacheData,
+            revalidateAfter: false,
+            isFallback: false,
+          }
+        }
+      }
+    }
+
     // we don't leverage the prerender cache in dev mode
     // so that getStaticProps is always called for easier debugging
     if (
@@ -418,23 +437,6 @@ export class IncrementalCache implements IncrementalCacheType {
     let entry: IncrementalCacheEntry | null = null
     let revalidate = ctx.revalidate
 
-    if (this.hasDynamicIO && ctx.kind === IncrementalCacheKind.FETCH) {
-      const cacheScope = cacheScopeAsyncLocalStorage.getStore()
-
-      if (cacheScope) {
-        const memoryCacheData = cacheScope.cache.get(cacheKey)
-
-        if (memoryCacheData?.kind === CachedRouteKind.FETCH) {
-          return {
-            isStale: false,
-            value: memoryCacheData,
-            revalidateAfter: false,
-            isFallback: false,
-          }
-        }
-      }
-    }
-
     const cacheData = await this.cacheHandler?.get(cacheKey, ctx)
 
     if (cacheData?.value?.kind === CachedRouteKind.FETCH) {
@@ -538,10 +540,10 @@ export class IncrementalCache implements IncrementalCacheType {
       isFallback?: boolean
     }
   ) {
-    if (this.disableForTestmode || (this.dev && !ctx.fetchCache)) return
-
-    pathname = this._getPathname(pathname, ctx.fetchCache)
-
+    // Even if we otherwise disable caching for testMode or if no fetchCache is configured
+    // we still always stash results in the cacheScope if one exists. This is because this
+    // is a transient in memory cache that populates caches ahead of a dynamic render in dev mode
+    // to allow the RSC debug info to have the right environment associated to it.
     if (this.hasDynamicIO && data?.kind === CachedRouteKind.FETCH) {
       const cacheScope = cacheScopeAsyncLocalStorage.getStore()
 
@@ -550,6 +552,10 @@ export class IncrementalCache implements IncrementalCacheType {
       }
     }
 
+    if (this.disableForTestmode || (this.dev && !ctx.fetchCache)) return
+
+    pathname = this._getPathname(pathname, ctx.fetchCache)
+
     // FetchCache has upper limit of 2MB per-entry currently
     const itemSize = JSON.stringify(data).length
     if (

packages/next/src/server/response-cache/types.ts

@@ -172,6 +172,7 @@ export type ResponseGenerator = (state: {
   hasResolved: boolean
   previousCacheEntry?: IncrementalCacheItem
   isRevalidating?: boolean
+  isDevWarmup?: boolean
 }) => Promise<ResponseCacheEntry | null>
 
 export type IncrementalCacheItem = {

test/development/app-dir/dynamic-io-dev-warmup/app/data-fetching.ts

@@ -0,0 +1,10 @@
+export async function fetchCached(url: string) {
+  const response = await fetch(url, { cache: 'force-cache' })
+  return response.text()
+}
+
+export async function getCachedData(_key: string) {
+  'use cache'
+  await new Promise((r) => setTimeout(r))
+  return Math.random()
+}