setState() in componentDidMount() should flush synchronously even with createBatch() (#12466)

* Add a failing test for setState in cDM during batch.commit()

* Copy pasta

* Flush all follow-up Sync work on the committed batch

* Nit: Use performSyncWork

Call performSyncWork right after flushing the batch. Does effectively
the same thing by reusing the existing function.

Also added some comments.

* Delete accidentally duplicated test

Author: gaearon

packages/react-dom/src/__tests__/ReactDOMRoot-test.internal.js

@@ -185,6 +185,46 @@ describe('ReactDOMRoot', () => {
     expect(container.textContent).toEqual('Hi');
   });
 
+  it('applies setState in componentDidMount synchronously in a batch', done => {
+    class App extends React.Component {
+      state = {mounted: false};
+      componentDidMount() {
+        this.setState({
+          mounted: true,
+        });
+      }
+      render() {
+        return this.state.mounted ? 'Hi' : 'Bye';
+      }
+    }
+
+    const root = ReactDOM.createRoot(container);
+    const batch = root.createBatch();
+    batch.render(
+      <AsyncMode>
+        <App />
+      </AsyncMode>,
+    );
+
+    flush();
+
+    // Hasn't updated yet
+    expect(container.textContent).toEqual('');
+
+    let ops = [];
+    batch.then(() => {
+      // Still hasn't updated
+      ops.push(container.textContent);
+
+      // Should synchronously commit
+      batch.commit();
+      ops.push(container.textContent);
+
+      expect(ops).toEqual(['', 'Hi']);
+      done();
+    });
+  });
+
   it('does not restart a completed batch when committing if there were no intervening updates', () => {
     let ops = [];
     function Foo(props) {

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -1169,7 +1169,15 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
             interruptedBy = fiber;
             resetStack();
           }
-          if (nextRoot !== root || !isWorking) {
+          if (
+            // If we're in the render phase, we don't need to schedule this root
+            // for an update, because we'll do it before we exit...
+            !isWorking ||
+            isCommitting ||
+            // ...unless this is a different root than the one we're rendering.
+            nextRoot !== root
+          ) {
+            // Add this root to the root schedule.
             requestWork(root, expirationTime);
           }
           if (nestedUpdateCount > NESTED_UPDATE_LIMIT) {
@@ -1500,6 +1508,8 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
     nextFlushedRoot = root;
     nextFlushedExpirationTime = expirationTime;
     performWorkOnRoot(root, expirationTime, false);
+    // Flush any sync work that was scheduled by lifecycles
+    performSyncWork();
     finishRendering();
   }