Wire AsyncLocalStorage within a cached context (#70573)

This ensures we exit the RequestStorage/PrerenderStorage context and
anything else but keeps the StaticGenerationStore (which should no
longer have page/request specific information). I also create new
CacheStore which will be used for life/tags but for now is just used to
provide a better error when cookies/headers/draftMode are accessed
within the cache scope.

I tried to make `React.cache` AsyncLocalStorage scope be shared between
the invocation of the cached function and anything it renders later.
This is also necessary in case the cached function calls Float methods
like `preload()`. However, I hit issues with preserving debug info,
having the right owner and that rendering a Promise in a Server
Component turns it into a `React.lazy` for now which doesn't preserve
the type we want. The proper solution is for React to expose a larger
scope for `React.cache` and float.

errors/next-request-in-use-cache.md

@@ -0,0 +1,51 @@
+---
+title: Cannot access `cookies()` or `headers()` in `"use cache"`
+---
+
+#### Why This Error Occurred
+
+A function is trying to read from the current incoming request inside the scope of a function annotated with `"use cache"`. This is not supported because it would make the cache invalidated by every request which is probably not what you intended.
+
+#### Possible Ways to Fix It
+
+Instead of calling this inside the `"use cache"` function, move it outside the function and pass the value in as an argument. The specific value will now be part of the cache key through its arguments.
+
+Before:
+
+```jsx filename="app/page.js"
+import { cookies } from 'next/headers'
+
+async function getExampleData() {
+  "use cache"
+- const isLoggedIn = (await cookies()).has('token')
+  ...
+}
+
+export default async function Page() {
+  const data = await getExampleData()
+  return ...
+}
+```
+
+After:
+
+```jsx filename="app/page.js"
+import { cookies } from 'next/headers'
+
+async function getExampleData(isLoggedIn) {
+  "use cache"
+  ...
+}
+
+export default async function Page() {
++ const isLoggedIn = (await cookies()).has('token')
+  const data = await getExampleData(isLoggedIn)
+  return ...
+}
+```
+
+### Useful Links
+
+- [`headers()` function](https://nextjs.org/docs/app/api-reference/functions/headers)
+- [`cookies()` function](https://nextjs.org/docs/app/api-reference/functions/cookies)
+- [`draftMode()` function](https://nextjs.org/docs/app/api-reference/functions/draft-mode)

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

@@ -10,6 +10,8 @@ import type { DeepReadonly } from '../../shared/lib/deep-readonly'
 import type { AfterContext } from '../../server/after/after-context'
 import type { ServerComponentsHmrCache } from '../../server/response-cache'
 
+import { cacheAsyncStorage } from '../../server/app-render/cache-async-storage.external'
+
 export interface RequestStore {
   /**
    * The URL of the request. This only specifies the pathname and the search
@@ -48,6 +50,11 @@ export { requestAsyncStorage }
 export function getExpectedRequestStore(callingExpression: string) {
   const store = requestAsyncStorage.getStore()
   if (store) return store
+  if (cacheAsyncStorage.getStore()) {
+    throw new Error(
+      `\`${callingExpression}\` cannot be called inside "use cache". Call it outside and pass an argument instead. Read more: https://nextjs.org/docs/messages/next-request-in-use-cache`
+    )
+  }
   throw new Error(
     `\`${callingExpression}\` was called outside a request scope. Read more: https://nextjs.org/docs/messages/next-dynamic-api-wrong-context`
   )

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

@@ -155,6 +155,8 @@ import {
 import { CacheSignal } from './cache-signal'
 import { getTracedMetadata } from '../lib/trace/utils'
 
+import './clean-async-snapshot.external'
+
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
   segment: string

packages/next/src/server/app-render/cache-async-storage-instance.ts

@@ -0,0 +1,4 @@
+import type { CacheAsyncStorage } from './cache-async-storage.external'
+import { createAsyncLocalStorage } from '../../client/components/async-local-storage'
+
+export const cacheAsyncStorage: CacheAsyncStorage = createAsyncLocalStorage()

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

@@ -0,0 +1,15 @@
+import type { AsyncLocalStorage } from 'async_hooks'
+
+// Share the instance module in the next-shared layer
+import { cacheAsyncStorage } from './cache-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
+
+/**
+ * The Cache store is for tracking information inside a "use cache" or unstable_cache context.
+ * Inside this context we should never expose any request or page specific information.
+ */
+export type CacheStore = {
+  // TODO: Inside this scope we'll track tags and life times of this scope.
+}
+
+export type CacheAsyncStorage = AsyncLocalStorage<CacheStore>
+export { cacheAsyncStorage }

packages/next/src/server/app-render/clean-async-snapshot-instance.ts

@@ -0,0 +1,6 @@
+import { createSnapshot } from '../../client/components/async-local-storage'
+
+export const runInCleanSnapshot: <R, TArgs extends any[]>(
+  fn: (...args: TArgs) => R,
+  ...args: TArgs
+) => R = createSnapshot()

packages/next/src/server/app-render/clean-async-snapshot.external.ts

@@ -0,0 +1,4 @@
+// Share the instance module in the next-shared layer
+import { runInCleanSnapshot } from './clean-async-snapshot-instance' with { 'turbopack-transition': 'next-shared' }
+
+export { runInCleanSnapshot }

packages/next/src/server/use-cache/use-cache-wrapper.ts

@@ -1,5 +1,4 @@
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
-import { createSnapshot } from '../../client/components/async-local-storage'
 /* eslint-disable import/no-extraneous-dependencies */
 import {
   renderToReadableStream,
@@ -15,6 +14,9 @@ import {
 
 import type { StaticGenerationStore } from '../../client/components/static-generation-async-storage.external'
 import { staticGenerationAsyncStorage } from '../../client/components/static-generation-async-storage.external'
+import type { CacheStore } from '../app-render/cache-async-storage.external'
+import { cacheAsyncStorage } from '../app-render/cache-async-storage.external'
+import { runInCleanSnapshot } from '../app-render/clean-async-snapshot.external'
 
 import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
 
@@ -73,10 +75,80 @@ cacheHandlerMap.set('default', {
   },
 })
 
-// TODO: Consider moving this another module that is guaranteed to be required in a safe scope.
-const runInCleanSnapshot = createSnapshot()
+function generateCacheEntry(
+  staticGenerationStore: StaticGenerationStore,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  cacheHandler: CacheHandler,
+  serializedCacheKey: string | ArrayBuffer,
+  encodedArguments: FormData | string,
+  fn: any
+): Promise<any> {
+  // We need to run this inside a clean AsyncLocalStorage snapshot so that the cache
+  // generation cannot read anything from the context we're currently executing which
+  // might include request specific things like cookies() inside a React.cache().
+  // Note: It is important that we await at least once before this because it lets us
+  // pop out of any stack specific contexts as well - aka "Sync" Local Storage.
+  return runInCleanSnapshot(
+    generateCacheEntryWithRestoredStaticGenerationStore,
+    staticGenerationStore,
+    clientReferenceManifest,
+    cacheHandler,
+    serializedCacheKey,
+    encodedArguments,
+    fn
+  )
+}
+
+function generateCacheEntryWithRestoredStaticGenerationStore(
+  staticGenerationStore: StaticGenerationStore,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  cacheHandler: CacheHandler,
+  serializedCacheKey: string | ArrayBuffer,
+  encodedArguments: FormData | string,
+  fn: any
+) {
+  // Since we cleared the AsyncLocalStorage we need to restore the staticGenerationStore.
+  // Note: We explicitly don't restore the RequestStore nor the PrerenderStore.
+  // We don't want any request specific information leaking an we don't want to create a
+  // bloated fake request mock for every cache call. So any feature that currently lives
+  // in RequestStore but should be available to Caches need to move to StaticGenerationStore.
+  // PrerenderStore is not needed inside the cache scope because the outer most one will
+  // be the one to report its result to the outer Prerender.
+  return staticGenerationAsyncStorage.run(
+    staticGenerationStore,
+    generateCacheEntryWithCacheContext,
+    staticGenerationStore,
+    clientReferenceManifest,
+    cacheHandler,
+    serializedCacheKey,
+    encodedArguments,
+    fn
+  )
+}
+
+function generateCacheEntryWithCacheContext(
+  staticGenerationStore: StaticGenerationStore,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  cacheHandler: CacheHandler,
+  serializedCacheKey: string | ArrayBuffer,
+  encodedArguments: FormData | string,
+  fn: any
+) {
+  // Initialize the Store for this Cache entry.
+  const cacheStore: CacheStore = {}
+  return cacheAsyncStorage.run(
+    cacheStore,
+    generateCacheEntryImpl,
+    staticGenerationStore,
+    clientReferenceManifest,
+    cacheHandler,
+    serializedCacheKey,
+    encodedArguments,
+    fn
+  )
+}
 
-async function generateCacheEntry(
+async function generateCacheEntryImpl(
   staticGenerationStore: StaticGenerationStore,
   clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
   cacheHandler: CacheHandler,
@@ -230,8 +302,7 @@ export function cache(kind: string, id: string, fn: any) {
         const clientReferenceManifestSingleton =
           getClientReferenceManifestSingleton()
 
-        stream = await runInCleanSnapshot(
-          generateCacheEntry,
+        stream = await generateCacheEntry(
           staticGenerationStore,
           clientReferenceManifestSingleton,
           cacheHandler,
@@ -246,8 +317,7 @@ export function cache(kind: string, id: string, fn: any) {
           // then we should warm up the cache with a fresh revalidated entry.
           const clientReferenceManifestSingleton =
             getClientReferenceManifestSingleton()
-          const ignoredStream = await runInCleanSnapshot(
-            generateCacheEntry,
+          const ignoredStream = await generateCacheEntry(
             staticGenerationStore,
             clientReferenceManifestSingleton,
             cacheHandler,

test/e2e/app-dir/use-cache/app/errors/error-boundary.tsx

@@ -0,0 +1,19 @@
+'use client'
+
+import { Component } from 'react'
+
+export default class ErrorBoundary extends Component<
+  { id: string; children: React.ReactNode },
+  { message: null | string }
+> {
+  state = { message: null }
+  static getDerivedStateFromError(error: any) {
+    return { message: error.message }
+  }
+  render() {
+    if (this.state.message !== null) {
+      return <p id={this.props.id}>{this.state.message}</p>
+    }
+    return this.props.children
+  }
+}

test/e2e/app-dir/use-cache/app/errors/page.tsx

@@ -0,0 +1,37 @@
+import ErrorBoundary from './error-boundary'
+
+import { cookies } from 'next/headers'
+
+import { currentUser, currentReferer, isEditing } from './util'
+
+async function Cookie() {
+  'use cache'
+  return <div>User: {currentUser()}</div>
+}
+
+async function Header() {
+  'use cache'
+  return <div>Referer: {currentReferer()}</div>
+}
+
+async function DraftMode() {
+  'use cache'
+  return <div>Editing: {isEditing()}</div>
+}
+
+export default async function Page() {
+  await cookies()
+  return (
+    <div>
+      <ErrorBoundary id="cookies">
+        <Cookie />
+      </ErrorBoundary>
+      <ErrorBoundary id="headers">
+        <Header />
+      </ErrorBoundary>
+      <ErrorBoundary id="draft-mode">
+        <DraftMode />
+      </ErrorBoundary>
+    </div>
+  )
+}

test/e2e/app-dir/use-cache/app/errors/util.ts

@@ -0,0 +1,13 @@
+import { cookies, headers, draftMode } from 'next/headers'
+
+export async function currentUser() {
+  return (await cookies()).get('user')?.value
+}
+
+export async function currentReferer() {
+  return (await headers()).get('referer')
+}
+
+export async function isEditing() {
+  return (await draftMode()).isEnabled
+}

test/e2e/app-dir/use-cache/app/layout.tsx

@@ -1,7 +1,11 @@
+import { Suspense } from 'react'
+
 export default function Root({ children }: { children: React.ReactNode }) {
   return (
     <html>
-      <body>{children}</body>
+      <body>
+        <Suspense fallback={null}>{children}</Suspense>
+      </body>
     </html>
   )
 }

test/e2e/app-dir/use-cache/app/react-cache/page.tsx

@@ -0,0 +1,24 @@
+import { cache } from 'react'
+
+const number = cache(() => {
+  return Math.random()
+})
+
+function Component() {
+  // Read number again in a component. This should be deduped.
+  return <p id="b">{number()}</p>
+}
+
+async function getCachedComponent() {
+  'use cache'
+  return (
+    <div>
+      <p id="a">{number()}</p>
+      <Component />
+    </div>
+  )
+}
+
+export default async function Page() {
+  return <div>{getCachedComponent()}</div>
+}