Capture the source and not just the stack on first seen error (#31367)

Otherwise we can't capture the owner stack at the right location when
there's a rethrow.

packages/react-noop-renderer/src/createReactNoop.js

@@ -81,6 +81,8 @@ type TextInstance = {
 type HostContext = Object;
 type CreateRootOptions = {
   unstable_transitionCallbacks?: TransitionTracingCallbacks,
+  onUncaughtError?: (error: mixed, errorInfo: {componentStack: string}) => void,
+  onCaughtError?: (error: mixed, errorInfo: {componentStack: string}) => void,
   ...
 };
 
@@ -1069,8 +1071,12 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
         null,
         false,
         '',
-        NoopRenderer.defaultOnUncaughtError,
-        NoopRenderer.defaultOnCaughtError,
+        options && options.onUncaughtError
+          ? options.onUncaughtError
+          : NoopRenderer.defaultOnUncaughtError,
+        options && options.onCaughtError
+          ? options.onCaughtError
+          : NoopRenderer.defaultOnCaughtError,
         onRecoverableError,
         options && options.unstable_transitionCallbacks
           ? options.unstable_transitionCallbacks

packages/react-reconciler/src/ReactCapturedValue.js

@@ -11,7 +11,7 @@ import type {Fiber} from './ReactInternalTypes';
 
 import {getStackByFiberInDevAndProd} from './ReactFiberComponentStack';
 
-const CapturedStacks: WeakMap<any, string> = new WeakMap();
+const CapturedStacks: WeakMap<any, CapturedValue<any>> = new WeakMap();
 
 export type CapturedValue<+T> = {
   +value: T,
@@ -25,36 +25,38 @@ export function createCapturedValueAtFiber<T>(
 ): CapturedValue<T> {
   // If the value is an error, call this function immediately after it is thrown
   // so the stack is accurate.
-  let stack;
   if (typeof value === 'object' && value !== null) {
-    const capturedStack = CapturedStacks.get(value);
-    if (typeof capturedStack === 'string') {
-      stack = capturedStack;
-    } else {
-      stack = getStackByFiberInDevAndProd(source);
-      CapturedStacks.set(value, stack);
+    const existing = CapturedStacks.get(value);
+    if (existing !== undefined) {
+      return existing;
     }
+    const captured = {
+      value,
+      source,
+      stack: getStackByFiberInDevAndProd(source),
+    };
+    CapturedStacks.set(value, captured);
+    return captured;
   } else {
-    stack = getStackByFiberInDevAndProd(source);
+    return {
+      value,
+      source,
+      stack: getStackByFiberInDevAndProd(source),
+    };
   }
-
-  return {
-    value,
-    source,
-    stack,
-  };
 }
 
 export function createCapturedValueFromError(
   value: Error,
   stack: null | string,
 ): CapturedValue<Error> {
-  if (typeof stack === 'string') {
-    CapturedStacks.set(value, stack);
-  }
-  return {
+  const captured = {
     value,
     source: null,
     stack: stack,
   };
+  if (typeof stack === 'string') {
+    CapturedStacks.set(value, captured);
+  }
+  return captured;
 }

packages/react-reconciler/src/__tests__/ReactErrorStacks-test.js

@@ -100,4 +100,77 @@ describe('ReactFragment', () => {
       ]),
     ]);
   });
+
+  it('retains owner stacks when rethrowing an error', async () => {
+    function Foo() {
+      return (
+        <RethrowingBoundary>
+          <Bar />
+        </RethrowingBoundary>
+      );
+    }
+    function Bar() {
+      return <SomethingThatErrors />;
+    }
+    function SomethingThatErrors() {
+      throw new Error('uh oh');
+    }
+
+    class RethrowingBoundary extends React.Component {
+      static getDerivedStateFromError(error) {
+        throw error;
+      }
+
+      render() {
+        return this.props.children;
+      }
+    }
+
+    const errors = [];
+    class CatchingBoundary extends React.Component {
+      constructor() {
+        super();
+        this.state = {};
+      }
+      static getDerivedStateFromError(error) {
+        return {errored: true};
+      }
+      render() {
+        if (this.state.errored) {
+          return null;
+        }
+        return this.props.children;
+      }
+    }
+
+    ReactNoop.createRoot({
+      onCaughtError(error, errorInfo) {
+        errors.push(
+          error.message,
+          normalizeCodeLocInfo(errorInfo.componentStack),
+          React.captureOwnerStack
+            ? normalizeCodeLocInfo(React.captureOwnerStack())
+            : null,
+        );
+      },
+    }).render(
+      <CatchingBoundary>
+        <Foo />
+      </CatchingBoundary>,
+    );
+    await waitForAll([]);
+    expect(errors).toEqual([
+      'uh oh',
+      componentStack([
+        'SomethingThatErrors',
+        'Bar',
+        'RethrowingBoundary',
+        'Foo',
+        'CatchingBoundary',
+      ]),
+      gate(flags => flags.enableOwnerStacks) && __DEV__
+        ? componentStack(['Bar', 'Foo'])
+        : null,
+    ]);
+  });
 });