[Flight] Implement captureStackTrace and owner stacks on the Server (#30197)

Wire up owner stacks in Flight to the shared internals. This exposes it
to `captureOwnerStack()`.

In this case we install it permanently as we only allow one RSC renderer
which then supports async contexts. Same thing we do for owner.

This also ends up adding it to errors logged by React through
`consoleWithStackDev`. The plan is to eventually remove that but this is
inline with what we do in Fizz and Fiber already.

However, at the same time we've instrumented the console so we need to
strip them back out before sending to the client. This lets the client
control whether to add the stack back in or allowing
`console.createTask` to control it.

This is another reason we shouldn't append them from React but for now
we hack it by removing them after the fact.

Author: sebmarkbage

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

@@ -81,6 +81,7 @@ let ErrorBoundary;
 let NoErrorExpected;
 let Scheduler;
 let assertLog;
+let assertConsoleErrorDev;
 
 describe('ReactFlight', () => {
   beforeEach(() => {
@@ -102,6 +103,7 @@ describe('ReactFlight', () => {
     Scheduler = require('scheduler');
     const InternalTestUtils = require('internal-test-utils');
     assertLog = InternalTestUtils.assertLog;
+    assertConsoleErrorDev = InternalTestUtils.assertConsoleErrorDev;
 
     ErrorBoundary = class extends React.Component {
       state = {hasError: false, error: null};
@@ -1441,9 +1443,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 +2728,76 @@ 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 **)',
+    );
+  });
+
+  // @gate (enableOwnerStacks && enableServerComponentLogs) || !__DEV__
+  it('should not include component stacks in replayed logs (unless DevTools add them)', () => {
+    function Foo() {
+      return 'hi';
+    }
+
+    function Bar() {
+      const array = [];
+      // Trigger key warning
+      array.push(ReactServer.createElement(Foo));
+      return ReactServer.createElement('div', null, array);
+    }
+
+    function App() {
+      return ReactServer.createElement(Bar);
+    }
+
+    const transport = ReactNoopFlightServer.render(
+      ReactServer.createElement(App),
+    );
+    assertConsoleErrorDev([
+      'Each child in a list should have a unique "key" prop.' +
+        ' See https://react.dev/link/warning-keys for more information.\n' +
+        '    in Bar (at **)\n' +
+        '    in App (at **)',
+    ]);
+
+    // Replay logs on the client
+    ReactNoopFlightClient.read(transport);
+    assertConsoleErrorDev(
+      [
+        'Each child in a list should have a unique "key" prop.' +
+          ' See https://react.dev/link/warning-keys for more information.',
+      ],
+      // We should not have a stack in the replay because that should be added either by console.createTask
+      // or React DevTools on the client. Neither of which we do here.
+      {withoutStack: true},
+    );
+  });
 });

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,10 @@ import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
 
 import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
 
+import {getOwnerStackByComponentInfoInDev} from './flight/ReactFlightComponentStack';
+
+import {isWritingAppendedStack} from 'shared/consoleWithStackDev';
+
 import {
   getIteratorFn,
   REACT_ELEMENT_TYPE,
@@ -263,8 +267,9 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
       'name',
     );
     const wrapperMethod = function (this: typeof console) {
+      let args = arguments;
       const request = resolveRequest();
-      if (methodName === 'assert' && arguments[0]) {
+      if (methodName === 'assert' && args[0]) {
         // assert doesn't emit anything unless first argument is falsy so we can skip it.
       } else if (request !== null) {
         // Extract the stack. Not all console logs print the full stack but they have at
@@ -276,7 +281,22 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
         // refer to previous logs in debug info to associate them with a component.
         const id = request.nextChunkId++;
         const owner: null | ReactComponentInfo = resolveOwner();
-        emitConsoleChunk(request, id, methodName, owner, stack, arguments);
+        if (
+          isWritingAppendedStack &&
+          (methodName === 'error' || methodName === 'warn') &&
+          args.length > 1 &&
+          typeof args[0] === 'string' &&
+          args[0].endsWith('%s')
+        ) {
+          // This looks like we've appended the component stack to the error from our own logs.
+          // We don't want those added to the replayed logs since those have the opportunity to add
+          // their own stacks or use console.createTask on the client as needed.
+          // TODO: Remove this special case once we remove consoleWithStackDev.
+          // $FlowFixMe[method-unbinding]
+          args = Array.prototype.slice.call(args, 0, args.length - 1);
+          args[0] = args[0].slice(0, args[0].length - 2);
+        }
+        emitConsoleChunk(request, id, methodName, owner, stack, args);
       }
       // $FlowFixMe[prop-missing]
       return originalMethod.apply(this, arguments);
@@ -317,6 +337,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 +526,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;
+  }
+}

packages/shared/consoleWithStackDev.js

@@ -40,6 +40,8 @@ export function error(format, ...args) {
 // eslint-disable-next-line react-internal/no-production-logging
 const supportsCreateTask = __DEV__ && enableOwnerStacks && !!console.createTask;
 
+export let isWritingAppendedStack = false;
+
 function printWarning(level, format, args, currentStack) {
   // When changing this logic, you might want to also
   // update consoleWithStackDev.www.js as well.
@@ -50,6 +52,7 @@ function printWarning(level, format, args, currentStack) {
       // can be lost while DevTools isn't open but we can't detect this.
       const stack = ReactSharedInternals.getCurrentStack(currentStack);
       if (stack !== '') {
+        isWritingAppendedStack = true;
         format += '%s';
         args = args.concat([stack]);
       }
@@ -60,5 +63,6 @@ function printWarning(level, format, args, currentStack) {
     // breaks IE9: https://github.com/facebook/react/issues/13610
     // eslint-disable-next-line react-internal/no-production-logging
     Function.prototype.apply.call(console[level], console, args);
+    isWritingAppendedStack = false;
   }
 }