Unify RequestStore, PrerenderStore and CacheStore (#70819)

Stacked on #70889.

This merges these three stores into a single AsyncLocalStorage context
containing a disjoint union named WorkUnitStore. That way they're
mutually exclusive. This prevents us from leaking RequestStore
information in either PrerenderStore or CacheStore.

---------

Co-authored-by: Janka Uryga <lolzatu2@gmail.com>

packages/next/src/build/templates/app-route.ts

@@ -32,25 +32,19 @@ const routeModule = new AppRouteRouteModule({
 // Pull out the exports that we need to expose from the module. This should
 // be eliminated when we've moved the other routes to the new format. These
 // are used to hook into the route.
-const {
-  requestAsyncStorage,
-  workAsyncStorage,
-  prerenderAsyncStorage,
-  serverHooks,
-} = routeModule
+const { workAsyncStorage, workUnitAsyncStorage, serverHooks } = routeModule
 
 function patchFetch() {
   return _patchFetch({
     workAsyncStorage,
-    requestAsyncStorage,
-    prerenderAsyncStorage,
+    workUnitAsyncStorage,
   })
 }
 
 export {
   routeModule,
-  requestAsyncStorage,
   workAsyncStorage,
+  workUnitAsyncStorage,
   serverHooks,
   patchFetch,
 }

packages/next/src/build/webpack-config.ts

@@ -131,7 +131,7 @@ const browserNonTranspileModules = [
 const precompileRegex = /[\\/]next[\\/]dist[\\/]compiled[\\/]/
 
 const asyncStoragesRegex =
-  /next[\\/]dist[\\/](esm[\\/])?client[\\/]components[\\/](work-async-storage|action-async-storage|request-async-storage)/
+  /next[\\/]dist[\\/](esm[\\/])?(client[\\/]components|server[\\/]app-render|)[\\/](work-async-storage|action-async-storage|work-unit-async-storage)/
 
 // Support for NODE_PATH
 const nodePathList = (process.env.NODE_PATH || '')

packages/next/src/server/after/after-context.ts

@@ -1,8 +1,9 @@
 import PromiseQueue from 'next/dist/compiled/p-queue'
 import {
-  requestAsyncStorage,
+  workUnitAsyncStorage,
   type RequestStore,
-} from '../../client/components/request-async-storage.external'
+  type WorkUnitStore,
+} from '../../server/app-render/work-unit-async-storage.external'
 import { ResponseCookies } from '../web/spec-extension/cookies'
 import type { RequestLifecycleOpts } from '../base-server'
 import type { AfterCallback, AfterTask } from './after'
@@ -20,7 +21,7 @@ export class AfterContext {
   private waitUntil: RequestLifecycleOpts['waitUntil'] | undefined
   private onClose: RequestLifecycleOpts['onClose'] | undefined
 
-  private requestStore: RequestStore | undefined
+  private workUnitStore: WorkUnitStore | undefined
 
   private runCallbacksOnClosePromise: Promise<void> | undefined
   private callbackQueue: PromiseQueue
@@ -55,11 +56,11 @@ export class AfterContext {
     if (!this.waitUntil) {
       errorWaitUntilNotAvailable()
     }
-    if (!this.requestStore) {
+    if (!this.workUnitStore) {
       // We just stash the first request store we have but this is not sufficient.
       // TODO: We should store a request store per callback since each callback might
       // be inside a different store. E.g. inside different batched actions, prerenders or caches.
-      this.requestStore = requestAsyncStorage.getStore()
+      this.workUnitStore = workUnitAsyncStorage.getStore()
     }
     if (!this.onClose) {
       throw new InvariantError(
@@ -72,7 +73,7 @@ export class AfterContext {
       // NOTE: We're creating a promise here, which means that
       // we will propagate any AsyncLocalStorage contexts we're currently in
       // to the callbacks that'll execute later.
-      // This includes e.g. `requestAsyncStorage` and React's `requestStorage` (which backs `React.cache()`).
+      // This includes e.g. `workUnitAsyncStorage` and React's `requestStorage` (which backs `React.cache()`).
       this.runCallbacksOnClosePromise = this.runCallbacksOnClose()
       this.waitUntil(this.runCallbacksOnClosePromise)
     }
@@ -94,27 +95,27 @@ export class AfterContext {
 
   private async runCallbacksOnClose() {
     await new Promise<void>((resolve) => this.onClose!(resolve))
-    return this.runCallbacks(this.requestStore)
+    return this.runCallbacks(this.workUnitStore)
   }
 
   private async runCallbacks(
-    requestStore: undefined | RequestStore
+    workUnitStore: undefined | WorkUnitStore
   ): Promise<void> {
     if (this.callbackQueue.size === 0) return
 
-    const readonlyRequestStore: undefined | RequestStore =
-      requestStore === undefined
+    const readonlyRequestStore: undefined | WorkUnitStore =
+      workUnitStore === undefined || workUnitStore.type !== 'request'
         ? undefined
         : // TODO: This is not sufficient. It should just be the same store that mutates.
-          wrapRequestStoreForAfterCallbacks(requestStore)
+          wrapRequestStoreForAfterCallbacks(workUnitStore)
 
     const workStore = workAsyncStorage.getStore()
 
     return withExecuteRevalidates(workStore, () =>
       // Clearing it out or running the first request store.
       // TODO: This needs to be the request store that was active at the time the
       // callback was scheduled but p-queue makes this hard so need further refactoring.
-      requestAsyncStorage.run(readonlyRequestStore as any, async () => {
+      workUnitAsyncStorage.run(readonlyRequestStore as any, async () => {
         this.callbackQueue.start()
         await this.callbackQueue.onIdle()
       })
@@ -133,6 +134,7 @@ function wrapRequestStoreForAfterCallbacks(
   requestStore: RequestStore
 ): RequestStore {
   return {
+    type: 'request',
     url: requestStore.url,
     get headers() {
       return requestStore.headers

packages/next/src/server/after/after.ts

@@ -1,5 +1,5 @@
 import { workAsyncStorage } from '../../client/components/work-async-storage.external'
-import { cacheAsyncStorage } from '../../server/app-render/cache-async-storage.external'
+import { workUnitAsyncStorage } from '../../server/app-render/work-unit-async-storage.external'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 
 import { markCurrentScopeAsDynamic } from '../app-render/dynamic-rendering'
@@ -12,7 +12,7 @@ export type AfterCallback<T = unknown> = () => T | Promise<T>
  */
 export function unstable_after<T>(task: AfterTask<T>): void {
   const workStore = workAsyncStorage.getStore()
-  const cacheStore = cacheAsyncStorage.getStore()
+  const workUnitStore = workUnitAsyncStorage.getStore()
 
   if (workStore) {
     const { afterContext } = workStore
@@ -29,7 +29,7 @@ export function unstable_after<T>(task: AfterTask<T>): void {
         `Route ${workStore.route} with \`dynamic = "force-static"\` couldn't be rendered statically because it used \`${callingExpression}\`. See more info here: https://nextjs.org/docs/app/building-your-application/rendering/static-and-dynamic#dynamic-rendering`
       )
     } else {
-      markCurrentScopeAsDynamic(workStore, cacheStore, callingExpression)
+      markCurrentScopeAsDynamic(workStore, workUnitStore, callingExpression)
     }
 
     afterContext.after(task)

packages/next/src/server/app-render/action-handler.ts

@@ -1,6 +1,6 @@
 import type { IncomingHttpHeaders, OutgoingHttpHeaders } from 'http'
 import type { SizeLimit } from '../../types'
-import type { RequestStore } from '../../client/components/request-async-storage.external'
+import type { RequestStore } from '../../server/app-render/work-unit-async-storage.external'
 import type { AppRenderContext, GenerateFlight } from './app-render'
 import type { AppPageModule } from '../../server/route-modules/app-page/module'
 import type { BaseNextRequest, BaseNextResponse } from '../base-http'