[Flight] Track Owner on AsyncLocalStorage When Available

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -21,6 +21,8 @@ if (typeof Blob === 'undefined') {
 if (typeof File === 'undefined') {
   global.File = require('buffer').File;
 }
+// Patch for Edge environments for global scope
+global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
 
 // Don't wait before processing work on the server.
 // TODO: we can replace this with FlightServer.act().
@@ -32,6 +34,7 @@ let webpackMap;
 let webpackModules;
 let webpackModuleLoading;
 let React;
+let ReactServer;
 let ReactDOMServer;
 let ReactServerDOMServer;
 let ReactServerDOMClient;
@@ -55,6 +58,7 @@ describe('ReactFlightDOMEdge', () => {
     webpackModules = WebpackMock.webpackModules;
     webpackModuleLoading = WebpackMock.moduleLoading;
 
+    ReactServer = require('react');
     ReactServerDOMServer = require('react-server-dom-webpack/server');
 
     jest.resetModules();
@@ -692,4 +696,71 @@ describe('ReactFlightDOMEdge', () => {
       ),
     );
   });
+
+  it('supports async server component debug info as the element owner in DEV', async () => {
+    function Container({children}) {
+      return children;
+    }
+
+    const promise = Promise.resolve(true);
+    async function Greeting({firstName}) {
+      // We can't use JSX here because it'll use the Client React.
+      const child = ReactServer.createElement(
+        'span',
+        null,
+        'Hello, ' + firstName,
+      );
+      // Yield the synchronous pass
+      await promise;
+      // We should still be able to track owner using AsyncLocalStorage.
+      return ReactServer.createElement(Container, null, child);
+    }
+
+    const model = {
+      greeting: ReactServer.createElement(Greeting, {firstName: 'Seb'}),
+    };
+
+    const stream = ReactServerDOMServer.renderToReadableStream(
+      model,
+      webpackMap,
+    );
+
+    const rootModel = await ReactServerDOMClient.createFromReadableStream(
+      stream,
+      {
+        ssrManifest: {
+          moduleMap: null,
+          moduleLoading: null,
+        },
+      },
+    );
+
+    const ssrStream = await ReactDOMServer.renderToReadableStream(
+      rootModel.greeting,
+    );
+    const result = await readResult(ssrStream);
+    expect(result).toEqual('<span>Hello, Seb</span>');
+
+    // Resolve the React Lazy wrapper which must have resolved by now.
+    const lazyWrapper = rootModel.greeting;
+    const greeting = lazyWrapper._init(lazyWrapper._payload);
+
+    // We've rendered down to the span.
+    expect(greeting.type).toBe('span');
+    if (__DEV__) {
+      const greetInfo = {name: 'Greeting', env: 'Server', owner: null};
+      expect(lazyWrapper._debugInfo).toEqual([
+        greetInfo,
+        {name: 'Container', env: 'Server', owner: greetInfo},
+      ]);
+      // The owner that created the span was the outer server component.
+      // We expect the debug info to be referentially equal to the owner.
+      expect(greeting._owner).toBe(lazyWrapper._debugInfo[0]);
+    } else {
+      expect(lazyWrapper._debugInfo).toBe(undefined);
+      expect(greeting._owner).toBe(
+        gate(flags => flags.disableStringRefs) ? undefined : null,
+      );
+    }
+  });
 });

packages/react-server/src/ReactFlightServer.js

@@ -73,6 +73,8 @@ import {
   isServerReference,
   supportsRequestStorage,
   requestStorage,
+  supportsComponentStorage,
+  componentStorage,
   createHints,
   initAsyncDebugInfo,
 } from './ReactFlightServerConfig';
@@ -89,11 +91,9 @@ import {
   getThenableStateAfterSuspending,
   resetHooksForRequest,
 } from './ReactFlightHooks';
-import {
-  DefaultAsyncDispatcher,
-  currentOwner,
-  setCurrentOwner,
-} from './flight/ReactFlightAsyncDispatcher';
+import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
+
+import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
 
 import {
   getIteratorFn,
@@ -162,7 +162,7 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
         // We don't currently use this id for anything but we emit it so that we can later
         // refer to previous logs in debug info to associate them with a component.
         const id = request.nextChunkId++;
-        const owner: null | ReactComponentInfo = currentOwner;
+        const owner: null | ReactComponentInfo = resolveOwner();
         emitConsoleChunk(request, id, methodName, owner, stack, arguments);
       }
       // $FlowFixMe[prop-missing]
@@ -824,7 +824,11 @@ function renderFunctionComponent<Props>(
   const prevThenableState = task.thenableState;
   task.thenableState = null;
 
-  let componentDebugInfo: null | ReactComponentInfo = null;
+  // The secondArg is always undefined in Server Components since refs error early.
+  const secondArg = undefined;
+  let result;
+
+  let componentDebugInfo: ReactComponentInfo;
   if (__DEV__) {
     if (debugID === null) {
       // We don't have a chunk to assign debug info. We need to outline this
@@ -853,20 +857,25 @@ function renderFunctionComponent<Props>(
       outlineModel(request, componentDebugInfo);
       emitDebugChunk(request, componentDebugID, componentDebugInfo);
     }
-  }
-
-  prepareToUseHooksForComponent(prevThenableState, componentDebugInfo);
-  // The secondArg is always undefined in Server Components since refs error early.
-  const secondArg = undefined;
-  let result;
-  if (__DEV__) {
+    prepareToUseHooksForComponent(prevThenableState, componentDebugInfo);
     setCurrentOwner(componentDebugInfo);
     try {
-      result = Component(props, secondArg);
+      if (supportsComponentStorage) {
+        // Run the component in an Async Context that tracks the current owner.
+        result = componentStorage.run(
+          componentDebugInfo,
+          Component,
+          props,
+          secondArg,
+        );
+      } else {
+        result = Component(props, secondArg);
+      }
     } finally {
       setCurrentOwner(null);
     }
   } else {
+    prepareToUseHooksForComponent(prevThenableState, null);
     result = Component(props, secondArg);
   }
   if (typeof result === 'object' && result !== null) {

packages/react-server/src/flight/ReactFlightAsyncDispatcher.js

@@ -15,6 +15,8 @@ import {resolveRequest, getCache} from '../ReactFlightServer';
 
 import {disableStringRefs} from 'shared/ReactFeatureFlags';
 
+import {resolveOwner} from './ReactFlightCurrentOwner';
+
 function resolveCache(): Map<Function, mixed> {
   const request = resolveRequest();
   if (request) {
@@ -36,19 +38,11 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({
   },
 }: any);
 
-export let currentOwner: ReactComponentInfo | null = null;
-
 if (__DEV__) {
-  DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => {
-    return currentOwner;
-  };
+  DefaultAsyncDispatcher.getOwner = resolveOwner;
 } else if (!disableStringRefs) {
   // Server Components never use string refs but the JSX runtime looks for it.
   DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => {
     return null;
   };
 }
-
-export function setCurrentOwner(componentInfo: null | ReactComponentInfo) {
-  currentOwner = componentInfo;
-}

packages/react-server/src/flight/ReactFlightCurrentOwner.js

@@ -0,0 +1,30 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {ReactComponentInfo} from 'shared/ReactTypes';
+
+import {
+  supportsComponentStorage,
+  componentStorage,
+} from '../ReactFlightServerConfig';
+
+let currentOwner: ReactComponentInfo | null = null;
+
+export function setCurrentOwner(componentInfo: null | ReactComponentInfo) {
+  currentOwner = componentInfo;
+}
+
+export function resolveOwner(): null | ReactComponentInfo {
+  if (currentOwner) return currentOwner;
+  if (supportsComponentStorage) {
+    const owner = componentStorage.getStore();
+    if (owner) return owner;
+  }
+  return null;
+}

packages/react-server/src/forks/ReactFlightServerConfig.custom.js

@@ -8,6 +8,7 @@
  */
 
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from '../ReactFlightServerConfigBundlerCustom';
 
@@ -23,6 +24,10 @@ export const isPrimaryRenderer = false;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
 
+export const supportsComponentStorage = false;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  (null: any);
+
 export function createHints(): any {
   return null;
 }

packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-esm.js

@@ -6,15 +6,18 @@
  *
  * @flow
  */
-import {AsyncLocalStorage} from 'async_hooks';
 
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from 'react-server-dom-esm/src/ReactFlightServerConfigESMBundler';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
 
-export const supportsRequestStorage = true;
-export const requestStorage: AsyncLocalStorage<Request | void> =
-  new AsyncLocalStorage();
+export const supportsRequestStorage = false;
+export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
+
+export const supportsComponentStorage = false;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  (null: any);
 
 export * from '../ReactFlightServerConfigDebugNoop';

packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-turbopack.js

@@ -8,11 +8,16 @@
  */
 
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
 
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
 
+export const supportsComponentStorage = false;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  (null: any);
+
 export * from '../ReactFlightServerConfigDebugNoop';

packages/react-server/src/forks/ReactFlightServerConfig.dom-browser.js

@@ -8,11 +8,16 @@
  */
 
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
 
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
 
+export const supportsComponentStorage = false;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  (null: any);
+
 export * from '../ReactFlightServerConfigDebugNoop';

packages/react-server/src/forks/ReactFlightServerConfig.dom-bun.js

@@ -8,11 +8,16 @@
  */
 
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from '../ReactFlightServerConfigBundlerCustom';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
 
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
 
+export const supportsComponentStorage = false;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  (null: any);
+
 export * from '../ReactFlightServerConfigDebugNoop';

packages/react-server/src/forks/ReactFlightServerConfig.dom-edge-turbopack.js

@@ -7,6 +7,7 @@
  * @flow
  */
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -16,6 +17,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
 export const requestStorage: AsyncLocalStorage<Request | void> =
   supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
 
+export const supportsComponentStorage: boolean =
+  __DEV__ && supportsRequestStorage;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
+
 // We use the Node version but get access to async_hooks from a global.
 import type {HookCallbacks, AsyncHook} from 'async_hooks';
 export const createAsyncHook: HookCallbacks => AsyncHook =

packages/react-server/src/forks/ReactFlightServerConfig.dom-edge.js

@@ -6,7 +6,9 @@
  *
  * @flow
  */
+
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -16,6 +18,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
 export const requestStorage: AsyncLocalStorage<Request | void> =
   supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
 
+export const supportsComponentStorage: boolean =
+  __DEV__ && supportsRequestStorage;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
+
 // We use the Node version but get access to async_hooks from a global.
 import type {HookCallbacks, AsyncHook} from 'async_hooks';
 export const createAsyncHook: HookCallbacks => AsyncHook =

packages/react-server/src/forks/ReactFlightServerConfig.dom-legacy.js

@@ -8,11 +8,16 @@
  */
 
 import type {Request} from 'react-server/src/ReactFlightServer';
+import type {ReactComponentInfo} from 'shared/ReactTypes';
 
 export * from '../ReactFlightServerConfigBundlerCustom';
 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
 
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
 
+export const supportsComponentStorage = false;
+export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
+  (null: any);
+
 export * from '../ReactFlightServerConfigDebugNoop';