Use asyncLocalStorage for request information

package.json

@@ -201,7 +201,7 @@
     "taskr": "1.1.0",
     "tree-kill": "1.2.2",
     "tsec": "0.2.1",
-    "turbo": "1.3.2-canary.1",
+    "turbo": "1.5.3",
     "typescript": "4.8.2",
     "wait-port": "0.2.2",
     "webpack": "5.74.0",

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 headers(): Headers {
+  staticGenerationBailout('headers')
+  const requestStore =
+    requestAsyncStorage && 'getStore' in requestAsyncStorage
+      ? requestAsyncStorage.getStore()!
+      : requestAsyncStorage
+
+  return requestStore.headers
 }
 
-export function usePreviewData() {
-  useStaticGenerationBailout('usePreviewData')
-  return useContext(PreviewDataContext)
+export function previewData() {
+  staticGenerationBailout('previewData')
+  const requestStore =
+    requestAsyncStorage && 'getStore' in requestAsyncStorage
+      ? requestAsyncStorage.getStore()!
+      : requestAsyncStorage
+  return requestStore.previewData
 }
 
-export function useCookies() {
-  useStaticGenerationBailout('useCookies')
-  return useContext(CookiesContext)
+export function cookies() {
+  staticGenerationBailout('cookies')
+  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/package.json

@@ -57,7 +57,7 @@
     "release": "taskr release",
     "build": "pnpm release && pnpm types",
     "prepublishOnly": "cd ../../ && turbo run build",
-    "types": "tsc --declaration --emitDeclarationOnly --declarationDir dist",
+    "types": "cat server/app-render.tsx && tsc --declaration --emitDeclarationOnly --declarationDir dist",
     "typescript": "tsec --noEmit",
     "ncc-compiled": "ncc cache clean && taskr ncc"
   },

packages/next/server/app-render.tsx

@@ -28,8 +28,98 @@ import {
 import { FlushEffectsContext } from '../shared/lib/flush-effects'
 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 { NextCookies } from './web/spec-extension/cookies'
+
+const INTERNAL_HEADERS_INSTANCE = Symbol('internal for headers readonly')
+
+function readonlyHeadersError() {
+  return new Error('ReadonlyHeaders cannot be modified')
+}
+class ReadonlyHeaders {
+  [INTERNAL_HEADERS_INSTANCE]: Headers
+
+  entries: Headers['entries']
+  forEach: Headers['forEach']
+  get: Headers['get']
+  has: Headers['has']
+  keys: Headers['keys']
+  values: Headers['values']
+
+  constructor(headers: IncomingHttpHeaders) {
+    // Since `new Headers` uses `this.append()` to fill the headers object ReadonlyHeaders can't extend from Headers directly as it would throw.
+    const headersInstance = new Headers(headers as any)
+    this[INTERNAL_HEADERS_INSTANCE] = headersInstance
+
+    this.entries = headersInstance.entries.bind(headersInstance)
+    this.forEach = headersInstance.forEach.bind(headersInstance)
+    this.get = headersInstance.get.bind(headersInstance)
+    this.has = headersInstance.has.bind(headersInstance)
+    this.keys = headersInstance.keys.bind(headersInstance)
+    this.values = headersInstance.values.bind(headersInstance)
+  }
+  [Symbol.iterator]() {
+    return this[INTERNAL_HEADERS_INSTANCE][Symbol.iterator]()
+  }
+
+  append() {
+    throw readonlyHeadersError()
+  }
+  delete() {
+    throw readonlyHeadersError()
+  }
+  set() {
+    throw readonlyHeadersError()
+  }
+}
+
+const INTERNAL_COOKIES_INSTANCE = Symbol('internal for cookies readonly')
+function readonlyCookiesError() {
+  return new Error('ReadonlyCookies cannot be modified')
+}
+class ReadonlyNextCookies {
+  [INTERNAL_COOKIES_INSTANCE]: NextCookies
+
+  entries: NextCookies['entries']
+  forEach: NextCookies['forEach']
+  get: NextCookies['get']
+  getWithOptions: NextCookies['getWithOptions']
+  has: NextCookies['has']
+  keys: NextCookies['keys']
+  values: NextCookies['values']
+
+  constructor(request: {
+    headers: {
+      get(key: 'cookie'): string | null | undefined
+    }
+  }) {
+    // Since `new Headers` uses `this.append()` to fill the headers object ReadonlyHeaders can't extend from Headers directly as it would throw.
+    // Request overridden to not have to provide a fully request object.
+    const cookiesInstance = new NextCookies(request as Request)
+    this[INTERNAL_COOKIES_INSTANCE] = cookiesInstance
+
+    this.entries = cookiesInstance.entries.bind(cookiesInstance)
+    this.forEach = cookiesInstance.forEach.bind(cookiesInstance)
+    this.get = cookiesInstance.get.bind(cookiesInstance)
+    this.getWithOptions = cookiesInstance.getWithOptions.bind(cookiesInstance)
+    this.has = cookiesInstance.has.bind(cookiesInstance)
+    this.keys = cookiesInstance.keys.bind(cookiesInstance)
+    this.values = cookiesInstance.values.bind(cookiesInstance)
+  }
+  [Symbol.iterator]() {
+    return this[INTERNAL_COOKIES_INSTANCE][Symbol.iterator]()
+  }
+
+  clear() {
+    throw readonlyCookiesError()
+  }
+  delete() {
+    throw readonlyCookiesError()
+  }
+  set() {
+    throw readonlyCookiesError()
+  }
+}
 
 // this needs to be required lazily so that `next-server` can set
 // the env before we require
@@ -541,6 +631,7 @@ export async function renderToHTMLOrFlight(
   patchFetch(ComponentMod)
 
   const staticGenerationAsyncStorage = ComponentMod.staticGenerationAsyncStorage
+  const requestAsyncStorage = ComponentMod.requestAsyncStorage
 
   if (
     !('getStore' in staticGenerationAsyncStorage) &&
@@ -558,9 +649,6 @@ export async function renderToHTMLOrFlight(
         ? staticGenerationAsyncStorage.getStore()
         : staticGenerationAsyncStorage
 
-    const { CONTEXT_NAMES } =
-      ComponentMod.serverHooks as typeof import('../client/components/hooks-server-context')
-
     // don't modify original query object
     query = Object.assign({}, query)
 
@@ -607,29 +695,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 +707,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 +1372,64 @@ 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 ReadonlyHeaders(headersWithoutFlight(req.headers)),
+    cookies: new ReadonlyNextCookies({
+      headers: {
+        get: (key) => {
+          if (key !== 'cookie') {
+            throw new Error('Only cookie header is supported')
+          }
+          return 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())
+  )
 }