Wire up owner stacks in Flight

This exposes it to captureOwnerStack().

In this case we install it permanently as we only allow one RSC renderer
which then supports async contexts.

It also exposes it to component stack addendums that React adds to its own
console.errors. At least for now.

Author: sebmarkbage

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -1441,9 +1441,7 @@ describe('ReactFlight', () => {
         <div>{Array(6).fill(<NoKey />)}</div>,
       );
       ReactNoopFlightClient.read(transport);
-    }).toErrorDev('Each child in a list should have a unique "key" prop.', {
-      withoutStack: gate(flags => flags.enableOwnerStacks),
-    });
+    }).toErrorDev('Each child in a list should have a unique "key" prop.');
   });
 
   it('should warn in DEV a child is missing keys in client component', async () => {
@@ -2728,4 +2726,36 @@ describe('ReactFlight', () => {
 
     expect(ReactNoop).toMatchRenderedOutput(<span>Hello, Seb</span>);
   });
+
+  // @gate __DEV__ && enableOwnerStacks
+  it('can get the component owner stacks during rendering in dev', () => {
+    let stack;
+
+    function Foo() {
+      return ReactServer.createElement(Bar, null);
+    }
+    function Bar() {
+      return ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Baz, null),
+      );
+    }
+
+    function Baz() {
+      stack = ReactServer.captureOwnerStack();
+      return ReactServer.createElement('span', null, 'hi');
+    }
+    ReactNoopFlightServer.render(
+      ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Foo, null),
+      ),
+    );
+
+    expect(normalizeCodeLocInfo(stack)).toBe(
+      '\n    in Bar (at **)' + '\n    in Foo (at **)',
+    );
+  });
 });

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

@@ -39,6 +39,15 @@ let ReactServerDOMServer;
 let ReactServerDOMClient;
 let use;
 
+function normalizeCodeLocInfo(str) {
+  return (
+    str &&
+    str.replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
+      return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
+    })
+  );
+}
+
 describe('ReactFlightDOMEdge', () => {
   beforeEach(() => {
     jest.resetModules();
@@ -883,4 +892,42 @@ describe('ReactFlightDOMEdge', () => {
       );
     }
   });
+
+  // @gate __DEV__ && enableOwnerStacks
+  it('can get the component owner stacks asynchronously', async () => {
+    let stack;
+
+    function Foo() {
+      return ReactServer.createElement(Bar, null);
+    }
+    function Bar() {
+      return ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Baz, null),
+      );
+    }
+
+    const promise = Promise.resolve(0);
+
+    async function Baz() {
+      await promise;
+      stack = ReactServer.captureOwnerStack();
+      return ReactServer.createElement('span', null, 'hi');
+    }
+
+    const stream = ReactServerDOMServer.renderToReadableStream(
+      ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Foo, null),
+      ),
+      webpackMap,
+    );
+    await readResult(stream);
+
+    expect(normalizeCodeLocInfo(stack)).toBe(
+      '\n    in Bar (at **)' + '\n    in Foo (at **)',
+    );
+  });
 });

packages/react-server/src/ReactFlightServer.js

@@ -97,6 +97,8 @@ import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
 
 import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
 
+import {getOwnerStackByComponentInfoInDev} from './flight/ReactFlightComponentStack';
+
 import {
   getIteratorFn,
   REACT_ELEMENT_TYPE,
@@ -317,6 +319,21 @@ if (
   patchConsole(console, 'warn');
 }
 
+function getCurrentStackInDEV(): string {
+  if (__DEV__) {
+    if (enableOwnerStacks) {
+      const owner: null | ReactComponentInfo = resolveOwner();
+      if (owner === null) {
+        return '';
+      }
+      return getOwnerStackByComponentInfoInDev(owner);
+    }
+    // We don't have Parent Stacks in Flight.
+    return '';
+  }
+  return '';
+}
+
 const ObjectPrototype = Object.prototype;
 
 type JSONValue =
@@ -491,6 +508,12 @@ function RequestInstance(
     );
   }
   ReactSharedInternals.A = DefaultAsyncDispatcher;
+  if (__DEV__) {
+    // Unlike Fizz or Fiber, we don't reset this and just keep it on permanently.
+    // This lets it act more like the AsyncDispatcher so that we can get the
+    // stack asynchronously too.
+    ReactSharedInternals.getCurrentStack = getCurrentStackInDEV;
+  }
 
   const abortSet: Set<Task> = new Set();
   const pingedTasks: Array<Task> = [];

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

@@ -0,0 +1,52 @@
+/**
+ * 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 {describeBuiltInComponentFrame} from 'shared/ReactComponentStackFrame';
+
+import {enableOwnerStacks} from 'shared/ReactFeatureFlags';
+
+export function getOwnerStackByComponentInfoInDev(
+  componentInfo: ReactComponentInfo,
+): string {
+  if (!enableOwnerStacks || !__DEV__) {
+    return '';
+  }
+  try {
+    let info = '';
+
+    // The owner stack of the current component will be where it was created, i.e. inside its owner.
+    // There's no actual name of the currently executing component. Instead, that is available
+    // on the regular stack that's currently executing. However, if there is no owner at all, then
+    // there's no stack frame so we add the name of the root component to the stack to know which
+    // component is currently executing.
+    if (!componentInfo.owner && typeof componentInfo.name === 'string') {
+      return describeBuiltInComponentFrame(componentInfo.name);
+    }
+
+    let owner: void | null | ReactComponentInfo = componentInfo;
+
+    while (owner) {
+      if (typeof owner.stack === 'string') {
+        // Server Component
+        const ownerStack: string = owner.stack;
+        owner = owner.owner;
+        if (owner && ownerStack !== '') {
+          info += '\n' + ownerStack;
+        }
+      } else {
+        break;
+      }
+    }
+    return info;
+  } catch (x) {
+    return '\nError generating stack: ' + x.message + '\n' + x.stack;
+  }
+}