Add cache scope handling for dynamic IO for dev/build (vercel#70408)

As discussed this adds an in memory cache scope which is leveraged for
seeding during prefetch and then leveraged during non-prefetch requests
in development and during build it shares a cache scope across one build
worker. During production server mode the cache scopes are specific
per-request with no prefetch cache seeding.

packages/next/src/build/index.ts

@@ -1897,6 +1897,7 @@ export default async function build(
               distDir,
               configFileName,
               runtimeEnvConfig,
+              dynamicIO: Boolean(config.experimental.dynamicIO),
               httpAgentOptions: config.httpAgentOptions,
               locales: config.i18n?.locales,
               defaultLocale: config.i18n?.defaultLocale,
@@ -2112,6 +2113,7 @@ export default async function build(
                             pageRuntime,
                             edgeInfo,
                             pageType,
+                            dynamicIO: Boolean(config.experimental.dynamicIO),
                             cacheHandler: config.cacheHandler,
                             isrFlushToDisk: ciEnvironment.hasNextSupport
                               ? false

packages/next/src/build/utils.ts

@@ -1348,6 +1348,7 @@ export async function buildAppStaticPaths({
   dir,
   page,
   distDir,
+  dynamicIO,
   configFileName,
   generateParams,
   isrFlushToDisk,
@@ -1362,6 +1363,7 @@ export async function buildAppStaticPaths({
 }: {
   dir: string
   page: string
+  dynamicIO: boolean
   configFileName: string
   generateParams: GenerateParamsResults
   distDir: string
@@ -1390,6 +1392,7 @@ export async function buildAppStaticPaths({
   const incrementalCache = new IncrementalCache({
     fs: nodeFs,
     dev: true,
+    dynamicIO,
     flushToDisk: isrFlushToDisk,
     serverDistDir: path.join(distDir, 'server'),
     fetchCacheKeyPrefix,
@@ -1581,6 +1584,7 @@ export async function isPageStatic({
   pageRuntime,
   edgeInfo,
   pageType,
+  dynamicIO,
   originalAppPath,
   isrFlushToDisk,
   maxMemoryCacheSize,
@@ -1592,6 +1596,7 @@ export async function isPageStatic({
   dir: string
   page: string
   distDir: string
+  dynamicIO: boolean
   configFileName: string
   runtimeEnvConfig: any
   httpAgentOptions: NextConfigComplete['httpAgentOptions']
@@ -1722,6 +1727,7 @@ export async function isPageStatic({
             await buildAppStaticPaths({
               dir,
               page,
+              dynamicIO,
               configFileName,
               generateParams,
               distDir,

packages/next/src/client/components/static-generation-async-storage.external.ts

@@ -63,6 +63,7 @@ export interface StaticGenerationStore {
 
   isDraftMode?: boolean
   isUnstableNoStore?: boolean
+  isPrefetchRequest?: boolean
 
   requestEndedState?: { ended?: boolean }
 }

packages/next/src/export/helpers/create-incremental-cache.ts

@@ -7,12 +7,14 @@ import { formatDynamicImportPath } from '../../lib/format-dynamic-import-path'
 
 export async function createIncrementalCache({
   cacheHandler,
+  dynamicIO,
   cacheMaxMemorySize,
   fetchCacheKeyPrefix,
   distDir,
   dir,
   flushToDisk,
 }: {
+  dynamicIO: boolean
   cacheHandler?: string
   cacheMaxMemorySize?: number
   fetchCacheKeyPrefix?: string
@@ -34,6 +36,7 @@ export async function createIncrementalCache({
     dev: false,
     requestHeaders: {},
     flushToDisk,
+    dynamicIO,
     fetchCache: true,
     maxMemoryCacheSize: cacheMaxMemorySize,
     fetchCacheKeyPrefix,

packages/next/src/export/worker.ts

@@ -46,6 +46,7 @@ import {
   type FallbackRouteParams,
 } from '../server/request/fallback-params'
 import { needsExperimentalReact } from '../lib/needs-experimental-react'
+import { runWithCacheScope } from '../server/async-storage/cache-scope'
 
 const envConfig = require('../shared/lib/runtime-config.external')
 
@@ -352,6 +353,7 @@ export async function exportPages(
     fetchCacheKeyPrefix,
     distDir,
     dir,
+    dynamicIO: Boolean(nextConfig.experimental.dynamicIO),
     // skip writing to disk in minimal mode for now, pending some
     // changes to better support it
     flushToDisk: !hasNextSupport,
@@ -459,21 +461,26 @@ export async function exportPages(
 
     return { result, path, pageKey }
   }
+  // for each build worker we share one dynamic IO cache scope
+  // this is only leveraged if the flag is enabled
+  const dynamicIOCacheScope = new Map()
 
-  for (let i = 0; i < paths.length; i += maxConcurrency) {
-    const subset = paths.slice(i, i + maxConcurrency)
+  await runWithCacheScope({ cache: dynamicIOCacheScope }, async () => {
+    for (let i = 0; i < paths.length; i += maxConcurrency) {
+      const subset = paths.slice(i, i + maxConcurrency)
 
-    const subsetResults = await Promise.all(
-      subset.map((path) =>
-        exportPageWithRetry(
-          path,
-          nextConfig.experimental.staticGenerationRetryCount ?? 1
+      const subsetResults = await Promise.all(
+        subset.map((path) =>
+          exportPageWithRetry(
+            path,
+            nextConfig.experimental.staticGenerationRetryCount ?? 1
+          )
         )
       )
-    )
 
-    results.push(...subsetResults)
-  }
+      results.push(...subsetResults)
+    }
+  })
 
   return results
 }

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

@@ -1300,6 +1300,7 @@ export const renderToHTMLOrFlight: AppPageRender = (
           fallbackRouteParams,
           renderOpts,
           requestEndedState,
+          isPrefetchRequest: Boolean(req.headers[NEXT_ROUTER_PREFETCH_HEADER]),
         },
         (staticGenerationStore) =>
           renderToHTMLOrFlightImpl(

packages/next/src/server/async-storage/cache-scope.ts

@@ -0,0 +1,27 @@
+import { AsyncLocalStorage } from 'async_hooks'
+
+export interface CacheScopeStore {
+  cache?: Map<string, any>
+}
+
+export const cacheScopeAsyncLocalStorage =
+  new AsyncLocalStorage<CacheScopeStore>()
+
+/**
+ * For dynamic IO handling we want to have a scoped memory
+ * cache which can live either the lifetime of a build worker,
+ * the lifetime of a specific request, or from a prefetch request
+ * to the request for non-prefetch version of a page (with
+ * drop-off after so long to prevent memory inflating)
+ */
+export function runWithCacheScope(
+  store: Partial<CacheScopeStore>,
+  fn: (...args: any[]) => Promise<any>
+) {
+  return cacheScopeAsyncLocalStorage.run(
+    {
+      cache: store.cache || new Map(),
+    },
+    fn
+  )
+}

packages/next/src/server/async-storage/with-static-generation-store.ts

@@ -21,6 +21,7 @@ export type StaticGenerationContext = {
   fallbackRouteParams: FallbackRouteParams | null
 
   requestEndedState?: { ended?: boolean }
+  isPrefetchRequest?: boolean
   renderOpts: {
     incrementalCache?: IncrementalCache
     isOnDemandRevalidate?: boolean
@@ -69,6 +70,7 @@ export const withStaticGenerationStore: WithStore<
     fallbackRouteParams,
     renderOpts,
     requestEndedState,
+    isPrefetchRequest,
   }: StaticGenerationContext,
   callback: (store: StaticGenerationStore) => Result
 ): Result => {
@@ -111,6 +113,7 @@ export const withStaticGenerationStore: WithStore<
     isDraftMode: renderOpts.isDraftMode,
 
     requestEndedState,
+    isPrefetchRequest,
   }
 
   // TODO: remove this when we resolve accessing the store outside the execution context

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

@@ -173,6 +173,11 @@ import type { RouteModule } from './route-modules/route-module'
 import { FallbackMode, parseFallbackField } from '../lib/fallback'
 import { toResponseCacheEntry } from './response-cache/utils'
 import { scheduleOnNextTick } from '../lib/scheduler'
+import { PrefetchCacheScopes } from './lib/prefetch-cache-scopes'
+import {
+  runWithCacheScope,
+  type CacheScopeStore,
+} from './async-storage/cache-scope'
 
 export type FindComponentsResult = {
   components: LoadComponentsReturnType
@@ -454,6 +459,14 @@ export default abstract class Server<
 
   private readonly isAppPPREnabled: boolean
 
+  private readonly prefetchCacheScopesDev = new PrefetchCacheScopes()
+
+  /**
+   * This is used to persist cache scopes across
+   * prefetch -> full route requests for dynamic IO
+   * it's only fully used in dev
+   */
+
   public constructor(options: ServerOptions) {
     const {
       dir = '.',
@@ -2748,7 +2761,7 @@ export default abstract class Server<
       }
     }
 
-    const responseGenerator: ResponseGenerator = async ({
+    let responseGenerator: ResponseGenerator = async ({
       hasResolved,
       previousCacheEntry,
       isRevalidating,
@@ -2999,6 +3012,54 @@ export default abstract class Server<
       }
     }
 
+    if (this.nextConfig.experimental.dynamicIO) {
+      const originalResponseGenerator = responseGenerator
+
+      responseGenerator = async (
+        ...args: Parameters<typeof responseGenerator>
+      ): ReturnType<typeof responseGenerator> => {
+        let cache: CacheScopeStore['cache'] | undefined
+
+        if (this.renderOpts.dev) {
+          cache = this.prefetchCacheScopesDev.get(urlPathname)
+
+          // we need to seed the prefetch cache scope in dev
+          // since we did not have a prefetch cache available
+          // and this is not a prefetch request
+          if (
+            !cache &&
+            !isPrefetchRSCRequest &&
+            routeModule?.definition.kind === RouteKind.APP_PAGE
+          ) {
+            req.headers[RSC_HEADER] = '1'
+            req.headers[NEXT_ROUTER_PREFETCH_HEADER] = '1'
+
+            cache = new Map()
+
+            await runWithCacheScope({ cache }, () =>
+              originalResponseGenerator(...args)
+            )
+            this.prefetchCacheScopesDev.set(urlPathname, cache)
+
+            delete req.headers[RSC_HEADER]
+            delete req.headers[NEXT_ROUTER_PREFETCH_HEADER]
+          }
+        }
+
+        return runWithCacheScope({ cache }, () =>
+          originalResponseGenerator(...args)
+        ).finally(() => {
+          if (this.renderOpts.dev) {
+            if (isPrefetchRSCRequest) {
+              this.prefetchCacheScopesDev.set(urlPathname, cache)
+            } else {
+              this.prefetchCacheScopesDev.del(urlPathname)
+            }
+          }
+        })
+      }
+    }
+
     const cacheEntry = await this.responseCache.get(
       ssgCacheKey,
       responseGenerator,

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

@@ -762,6 +762,7 @@ export default class DevServer extends Server {
             configFileName,
             publicRuntimeConfig,
             serverRuntimeConfig,
+            dynamicIO: Boolean(this.nextConfig.experimental.dynamicIO),
           },
           httpAgentOptions,
           locales,

packages/next/src/server/dev/static-paths-worker.ts

@@ -30,6 +30,7 @@ type RuntimeConfig = {
   configFileName: string
   publicRuntimeConfig: { [key: string]: any }
   serverRuntimeConfig: { [key: string]: any }
+  dynamicIO: boolean
 }
 
 // we call getStaticPaths in a separate process to ensure
@@ -115,6 +116,7 @@ export async function loadStaticPaths({
     return await buildAppStaticPaths({
       dir,
       page: pathname,
+      dynamicIO: config.dynamicIO,
       generateParams,
       configFileName: config.configFileName,
       distDir,