Track owner on componentStorage

Author: sebmarkbage

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';
@@ -91,7 +93,7 @@ import {
 } from './ReactFlightHooks';
 import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
 
-import {currentOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
+import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
 
 import {
   getIteratorFn,
@@ -160,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]
@@ -822,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
@@ -851,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,7 +15,7 @@ import {resolveRequest, getCache} from '../ReactFlightServer';
 
 import {disableStringRefs} from 'shared/ReactFeatureFlags';
 
-import {currentOwner} from './ReactFlightCurrentOwner';
+import {resolveOwner} from './ReactFlightCurrentOwner';
 
 function resolveCache(): Map<Function, mixed> {
   const request = resolveRequest();
@@ -39,9 +39,7 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({
 }: any);
 
 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 => {

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

@@ -9,8 +9,22 @@
 
 import type {ReactComponentInfo} from 'shared/ReactTypes';
 
-export let currentOwner: ReactComponentInfo | null = null;
+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;
+}