Fix context stack misalignment caused by error replay (#12508)

* Add regression tests for error boundary replay bugs

* Ensure the context stack is aligned if renderer throws

* Always throw when replaying a failed unit of work

Replaying a failed unit of work should always throw, because the render
phase is meant to be idempotent, If it doesn't throw, rethrow the
original error, so React's internal stack is not misaligned.

* Reset originalReplayError after replaying

* Typo fix

packages/react-reconciler/src/ReactFiberHostContext.js

@@ -60,12 +60,19 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
     // Push current root instance onto the stack;
     // This allows us to reset root when portals are popped.
     push(rootInstanceStackCursor, nextRootInstance, fiber);
-
-    const nextRootContext = getRootHostContext(nextRootInstance);
-
     // Track the context and the Fiber that provided it.
     // This enables us to pop only Fibers that provide unique contexts.
     push(contextFiberStackCursor, fiber, fiber);
+
+    // Finally, we need to push the host context to the stack.
+    // However, we can't just call getRootHostContext() and push it because
+    // we'd have a different number of entries on the stack depending on
+    // whether getRootHostContext() throws somewhere in renderer code or not.
+    // So we push an empty value first. This lets us safely unwind on errors.
+    push(contextStackCursor, NO_CONTEXT, fiber);
+    const nextRootContext = getRootHostContext(nextRootInstance);
+    // Now that we know this function doesn't throw, replace it.
+    pop(contextStackCursor, fiber);
     push(contextStackCursor, nextRootContext, fiber);
   }
 

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -263,10 +263,18 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
 
   let stashedWorkInProgressProperties;
   let replayUnitOfWork;
+  let isReplayingFailedUnitOfWork;
+  let originalReplayError;
   if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
     stashedWorkInProgressProperties = null;
-    replayUnitOfWork = (failedUnitOfWork: Fiber, isAsync: boolean) => {
-      // Retore the original state of the work-in-progress
+    isReplayingFailedUnitOfWork = false;
+    originalReplayError = null;
+    replayUnitOfWork = (
+      failedUnitOfWork: Fiber,
+      error: mixed,
+      isAsync: boolean,
+    ) => {
+      // Restore the original state of the work-in-progress
       assignFiberPropertiesInDEV(
         failedUnitOfWork,
         stashedWorkInProgressProperties,
@@ -290,12 +298,17 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
           break;
       }
       // Replay the begin phase.
+      isReplayingFailedUnitOfWork = true;
+      originalReplayError = error;
       invokeGuardedCallback(null, workLoop, null, isAsync);
+      isReplayingFailedUnitOfWork = false;
+      originalReplayError = null;
       if (hasCaughtError()) {
         clearCaughtError();
       } else {
-        // This should be unreachable because the render phase is
-        // idempotent
+        // If the begin phase did not fail the second time, set this pointer
+        // back to the original value.
+        nextUnitOfWork = failedUnitOfWork;
       }
     };
   }
@@ -855,9 +868,15 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
       );
     }
     let next = beginWork(current, workInProgress, nextRenderExpirationTime);
-
     if (__DEV__) {
       ReactDebugCurrentFiber.resetCurrentFiber();
+      if (isReplayingFailedUnitOfWork) {
+        // Currently replaying a failed unit of work. This should be unreachable,
+        // because the render phase is meant to be idempotent, and it should
+        // have thrown again. Since it didn't, rethrow the original error, so
+        // React's internal stack is not misaligned.
+        throw originalReplayError;
+      }
     }
     if (__DEV__ && ReactFiberInstrumentation.debugTool) {
       ReactFiberInstrumentation.debugTool.onBeginWork(workInProgress);
@@ -935,7 +954,7 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
 
         if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
           const failedUnitOfWork = nextUnitOfWork;
-          replayUnitOfWork(failedUnitOfWork, isAsync);
+          replayUnitOfWork(failedUnitOfWork, thrownValue, isAsync);
         }
 
         const sourceFiber: Fiber = nextUnitOfWork;

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

@@ -0,0 +1,59 @@
+/**
+ * Copyright (c) 2013-present, Facebook, Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+
+'use strict';
+
+let React;
+let ReactNoop;
+
+describe('ReactIncrementalErrorReplay', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+  });
+
+  function div(...children) {
+    children = children.map(c => (typeof c === 'string' ? {text: c} : c));
+    return {type: 'div', children, prop: undefined};
+  }
+
+  function span(prop) {
+    return {type: 'span', children: [], prop};
+  }
+
+  it('should fail gracefully on error in the host environment', () => {
+    ReactNoop.simulateErrorInHostConfig(() => {
+      ReactNoop.render(<span />);
+      expect(() => ReactNoop.flush()).toThrow('Error in host config.');
+    });
+  });
+
+  it('should fail gracefully on error that does not reproduce on replay', () => {
+    let didInit = false;
+
+    function badLazyInit() {
+      const needsInit = !didInit;
+      didInit = true;
+      if (needsInit) {
+        throw new Error('Hi');
+      }
+    }
+
+    class App extends React.Component {
+      render() {
+        badLazyInit();
+        return <div />;
+      }
+    }
+    ReactNoop.render(<App />);
+    expect(() => ReactNoop.flush()).toThrow('Hi');
+  });
+});