Fix queueing updates in cWM/cWRP when batching (#8398)

I tried to add a temporary check for the correct state in https://gist.github.com/spicyj/338147e989215b6eeaf48a6ce2d68d93 and run our test suites over it, but none of our existing test cases would have triggered that invariant. The new one I added does.

scripts/fiber/tests-failing.txt

@@ -113,7 +113,6 @@ src/renderers/shared/stack/reconciler/__tests__/ReactStatelessComponent-test.js
 
 src/renderers/shared/stack/reconciler/__tests__/ReactUpdates-test.js
 * should queue mount-ready handlers across different roots
-* should queue updates from during mount
 * marks top-level updates
 * throws in setState if the update callback is not a function
 * throws in replaceState if the update callback is not a function

scripts/fiber/tests-passing.txt

@@ -1294,6 +1294,7 @@ src/renderers/shared/stack/reconciler/__tests__/ReactCompositeComponent-test.js
 * unmasked context propagates through updates
 * should trigger componentWillReceiveProps for context changes
 * only renders once if updated in componentWillReceiveProps
+* only renders once if updated in componentWillReceiveProps when batching
 * should allow access to findDOMNode in componentWillUnmount
 * context should be passed down from the parent
 * should replace state
@@ -1449,6 +1450,7 @@ src/renderers/shared/stack/reconciler/__tests__/ReactUpdates-test.js
 * should flush updates in the correct order
 * should flush updates in the correct order across roots
 * should queue nested updates
+* should queue updates from during mount
 * calls componentWillReceiveProps setState callback properly
 * does not call render after a component as been deleted
 * does not update one component twice in a batch (#2410)

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -784,10 +784,13 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
     try {
       return fn();
     } finally {
-      shouldBatchUpdates = prev;
-      // If we've exited the batch, perform any scheduled task work
-      if (!shouldBatchUpdates) {
-        performTaskWork();
+      // If we're exiting the batch, perform any scheduled task work
+      try {
+        if (!prev) {
+          performTaskWork();
+        }
+      } finally {
+        shouldBatchUpdates = prev;
       }
     }
   }

src/renderers/shared/stack/reconciler/__tests__/ReactCompositeComponent-test.js

@@ -1045,7 +1045,9 @@ describe('ReactCompositeComponent', () => {
 
       componentWillReceiveProps(props) {
         expect(props.update).toBe(1);
+        expect(renders).toBe(1);
         this.setState({updated: true});
+        expect(renders).toBe(1);
       }
 
       render() {
@@ -1063,6 +1065,36 @@ describe('ReactCompositeComponent', () => {
     expect(instance.state.updated).toBe(true);
   });
 
+  it('only renders once if updated in componentWillReceiveProps when batching', () => {
+    var renders = 0;
+
+    class Component extends React.Component {
+      state = {updated: false};
+
+      componentWillReceiveProps(props) {
+        expect(props.update).toBe(1);
+        expect(renders).toBe(1);
+        this.setState({updated: true});
+        expect(renders).toBe(1);
+      }
+
+      render() {
+        renders++;
+        return <div />;
+      }
+    }
+
+    var container = document.createElement('div');
+    var instance = ReactDOM.render(<Component update={0} />, container);
+    expect(renders).toBe(1);
+    expect(instance.state.updated).toBe(false);
+    ReactDOM.unstable_batchedUpdates(() => {
+      ReactDOM.render(<Component update={1} />, container);
+    });
+    expect(renders).toBe(2);
+    expect(instance.state.updated).toBe(true);
+  });
+
   it('should update refs if shouldComponentUpdate gives false', () => {
     class Static extends React.Component {
       shouldComponentUpdate() {