Fix context stack misalignment caused by error replay

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');
+  });
+});