Use asyncLocalStorage for request information

packages/next/build/webpack/loaders/next-app-loader.ts

@@ -193,6 +193,7 @@ const nextAppLoader: webpack.LoaderDefinitionFunction<{
     }
 
     export const staticGenerationAsyncStorage = require('next/dist/client/components/static-generation-async-storage.js').staticGenerationAsyncStorage
+    export const requestAsyncStorage = require('next/dist/client/components/request-async-storage.js').requestAsyncStorage
 
     export const serverHooks = require('next/dist/client/components/hooks-server-context.js')
 

packages/next/client/components/hooks-server-context.ts

@@ -13,38 +13,3 @@ export class DynamicServerError extends Error {
     super(`Dynamic server usage: ${type}`)
   }
 }
-
-// Ensure serverContext is not created more than once as React will throw when creating it more than once
-// https://github.com/facebook/react/blob/dd2d6522754f52c70d02c51db25eb7cbd5d1c8eb/packages/react/src/ReactServerContext.js#L101
-const createContext = <T>(name: string, defaultValue: T | null = null) => {
-  // @ts-expect-error __NEXT_DEV_SERVER_CONTEXT__ is a global
-  if (!global.__NEXT_SERVER_CONTEXT__) {
-    // @ts-expect-error __NEXT_SERVER_CONTEXT__ is a global
-    global.__NEXT_SERVER_CONTEXT__ = {}
-  }
-
-  // @ts-expect-error __NEXT_SERVER_CONTEXT__ is a global
-  if (!global.__NEXT_SERVER_CONTEXT__[name]) {
-    // @ts-expect-error __NEXT_SERVER_CONTEXT__ is a global
-    global.__NEXT_SERVER_CONTEXT__[name] = createServerContext(
-      name,
-      defaultValue
-    )
-  }
-
-  // @ts-expect-error __NEXT_SERVER_CONTEXT__ is a global
-  return global.__NEXT_SERVER_CONTEXT__[name]
-}
-
-export const CONTEXT_NAMES = {
-  HeadersContext: 'HeadersContext',
-  PreviewDataContext: 'PreviewDataContext',
-  CookiesContext: 'CookiesContext',
-  FetchRevalidateContext: 'FetchRevalidateContext',
-} as const
-
-export const HeadersContext = createContext(CONTEXT_NAMES.HeadersContext)
-export const PreviewDataContext = createContext(
-  CONTEXT_NAMES.PreviewDataContext
-)
-export const CookiesContext = createContext(CONTEXT_NAMES.CookiesContext)

packages/next/client/components/hooks-server.ts

@@ -1,13 +1,8 @@
-import { useContext } from 'react'
-import {
-  HeadersContext,
-  PreviewDataContext,
-  CookiesContext,
-  DynamicServerError,
-} from './hooks-server-context'
+import { DynamicServerError } from './hooks-server-context'
+import { requestAsyncStorage } from './request-async-storage'
 import { staticGenerationAsyncStorage } from './static-generation-async-storage'
 
-function useStaticGenerationBailout(reason: string) {
+function staticGenerationBailout(reason: string) {
   const staticGenerationStore =
     staticGenerationAsyncStorage && 'getStore' in staticGenerationAsyncStorage
       ? staticGenerationAsyncStorage?.getStore()
@@ -22,17 +17,31 @@ function useStaticGenerationBailout(reason: string) {
   }
 }
 
-export function useHeaders() {
-  useStaticGenerationBailout('useHeaders')
-  return useContext(HeadersContext)
+export function useHeaders(): Headers {
+  staticGenerationBailout('useHeaders')
+  const requestStore =
+    requestAsyncStorage && 'getStore' in requestAsyncStorage
+      ? requestAsyncStorage.getStore()!
+      : requestAsyncStorage
+
+  return requestStore.headers
 }
 
 export function usePreviewData() {
-  useStaticGenerationBailout('usePreviewData')
-  return useContext(PreviewDataContext)
+  staticGenerationBailout('usePreviewData')
+  const requestStore =
+    requestAsyncStorage && 'getStore' in requestAsyncStorage
+      ? requestAsyncStorage.getStore()!
+      : requestAsyncStorage
+  return requestStore.previewData
 }
 
 export function useCookies() {
-  useStaticGenerationBailout('useCookies')
-  return useContext(CookiesContext)
+  staticGenerationBailout('useCookies')
+  const requestStore =
+    requestAsyncStorage && 'getStore' in requestAsyncStorage
+      ? requestAsyncStorage.getStore()!
+      : requestAsyncStorage
+
+  return requestStore.cookies
 }

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

@@ -0,0 +1,15 @@
+import type { AsyncLocalStorage } from 'async_hooks'
+import { Cookies } from '../../server/web/spec-extension/cookies'
+
+export interface RequestStore {
+  headers: Headers
+  cookies: Cookies
+  previewData: any
+}
+
+export let requestAsyncStorage: AsyncLocalStorage<RequestStore> | RequestStore =
+  {} as any
+
+if (process.env.NEXT_RUNTIME !== 'edge' && typeof window === 'undefined') {
+  requestAsyncStorage = new (require('async_hooks').AsyncLocalStorage)()
+}

packages/next/server/app-render.tsx

@@ -30,6 +30,7 @@ import { stripInternalQueries } from './internal-utils'
 import type { ComponentsType } from '../build/webpack/loaders/next-app-loader'
 import type { UnwrapPromise } from '../lib/coalesced-function'
 import { REDIRECT_ERROR_CODE } from '../client/components/redirect'
+import { Cookies } from './web/spec-extension/cookies'
 
 // this needs to be required lazily so that `next-server` can set
 // the env before we require
@@ -541,6 +542,7 @@ export async function renderToHTMLOrFlight(
   patchFetch(ComponentMod)
 
   const staticGenerationAsyncStorage = ComponentMod.staticGenerationAsyncStorage
+  const requestAsyncStorage = ComponentMod.requestAsyncStorage
 
   if (
     !('getStore' in staticGenerationAsyncStorage) &&
@@ -607,29 +609,10 @@ export async function renderToHTMLOrFlight(
       | typeof import('../client/components/hot-reloader.client').default
       | null
 
-    const headers = headersWithoutFlight(req.headers)
-    // TODO-APP: fix type of req
-    // @ts-expect-error
-    const cookies = req.cookies
-
     /**
      * The tree created in next-app-loader that holds component segments and modules
      */
     const loaderTree: LoaderTree = ComponentMod.tree
-
-    const tryGetPreviewData =
-      process.env.NEXT_RUNTIME === 'edge'
-        ? () => false
-        : require('./api-utils/node').tryGetPreviewData
-
-    // Reads of this are cached on the `req` object, so this should resolve
-    // instantly. There's no need to pass this data down from a previous
-    // invoke, where we'd have to consider server & serverless.
-    const previewData = tryGetPreviewData(
-      req,
-      res,
-      (renderOpts as any).previewProps
-    )
     /**
      * Server Context is specifically only available in Server Components.
      * It has to hold values that can't change while rendering from the common layout down.
@@ -638,9 +621,6 @@ export async function renderToHTMLOrFlight(
 
     const serverContexts: Array<[string, any]> = [
       ['WORKAROUND', null], // TODO-APP: First value has a bug currently where the value is not set on the second request: https://github.com/facebook/react/issues/24849
-      [CONTEXT_NAMES.HeadersContext, headers],
-      [CONTEXT_NAMES.CookiesContext, cookies],
-      [CONTEXT_NAMES.PreviewDataContext, previewData],
     ]
 
     type CreateSegmentPath = (child: FlightSegmentPath) => FlightSegmentPath
@@ -1306,18 +1286,55 @@ export async function renderToHTMLOrFlight(
     pathname,
   }
 
-  if ('getStore' in staticGenerationAsyncStorage) {
-    return new Promise<UnwrapPromise<ReturnType<typeof renderToHTMLOrFlight>>>(
-      (resolve, reject) => {
+  const tryGetPreviewData =
+    process.env.NEXT_RUNTIME === 'edge'
+      ? () => false
+      : require('./api-utils/node').tryGetPreviewData
+
+  // Reads of this are cached on the `req` object, so this should resolve
+  // instantly. There's no need to pass this data down from a previous
+  // invoke, where we'd have to consider server & serverless.
+  const previewData = tryGetPreviewData(
+    req,
+    res,
+    (renderOpts as any).previewProps
+  )
+
+  const requestStore = {
+    headers: new Headers(headersWithoutFlight(req.headers)),
+    cookies: new Cookies(req.headers.cookie),
+    previewData,
+  }
+
+  function handleRequestStoreRun<T>(fn: () => T): Promise<T> {
+    if ('getStore' in requestAsyncStorage) {
+      return new Promise((resolve, reject) => {
+        requestAsyncStorage.run(requestStore, () => {
+          return Promise.resolve(fn()).then(resolve).catch(reject)
+        })
+      })
+    } else {
+      Object.assign(requestAsyncStorage, requestStore)
+      return Promise.resolve(fn())
+    }
+  }
+
+  function handleStaticGenerationStoreRun<T>(fn: () => T): Promise<T> {
+    if ('getStore' in staticGenerationAsyncStorage) {
+      return new Promise((resolve, reject) => {
         staticGenerationAsyncStorage.run(initialStaticGenerationStore, () => {
-          return wrappedRender().then(resolve).catch(reject)
+          return Promise.resolve(fn()).then(resolve).catch(reject)
         })
-      }
-    )
-  } else {
-    Object.assign(staticGenerationAsyncStorage, initialStaticGenerationStore)
-    return wrappedRender().finally(() => {
-      staticGenerationAsyncStorage.inUse = false
-    })
+      })
+    } else {
+      Object.assign(staticGenerationAsyncStorage, initialStaticGenerationStore)
+      return Promise.resolve(fn()).finally(() => {
+        staticGenerationAsyncStorage.inUse = false
+      })
+    }
   }
+
+  return handleRequestStoreRun(() =>
+    handleStaticGenerationStoreRun(() => wrappedRender())
+  )
 }

test/e2e/app-dir/app/app/hooks/use-cookies/page.js

@@ -3,8 +3,7 @@ import { useCookies } from 'next/dist/client/components/hooks-server'
 export default function Page() {
   const cookies = useCookies()
 
-  const hasCookie =
-    'use-cookies' in cookies && cookies['use-cookies'] === 'value'
+  const hasCookie = cookies.get('use-cookies') === 'value'
 
   return (
     <>

test/e2e/app-dir/app/app/hooks/use-headers/page.js

@@ -3,8 +3,7 @@ import { useHeaders } from 'next/dist/client/components/hooks-server'
 export default function Page() {
   const headers = useHeaders()
 
-  const hasHeader =
-    'x-use-headers' in headers && headers['x-use-headers'] === 'value'
+  const hasHeader = headers.get('x-use-headers') === 'value'
 
   return (
     <>

test/e2e/app-dir/rsc-basic/app/page.js

@@ -5,7 +5,8 @@ const envVar = process.env.ENV_VAR_TEST
 const headerKey = 'x-next-test-client'
 
 export default function Index(props) {
-  const header = useHeaders()[headerKey]
+  const headers = useHeaders()
+  const header = headers.get(headerKey)
 
   return (
     <div>