Task work inside batched updates is always sync, even for initial mounts

Behavior now matches Stack. It's unfortunate that this prevents us
from unifying SynchronousPriority and TaskPriority.

scripts/fiber/tests-passing.txt

@@ -1309,6 +1309,7 @@ src/renderers/dom/shared/__tests__/ReactMount-test.js
 * should warn if render removes React-rendered children
 * should warn if the unmounted node was rendered by another copy of React
 * passes the correct callback context
+* initial mount is sync inside batchedUpdates, but task work is deferred until the end of the batch
 
 src/renderers/dom/shared/__tests__/ReactMountDestruction-test.js
 * should destroy a react root upon request
@@ -1832,6 +1833,7 @@ src/renderers/shared/fiber/__tests__/ReactIncrementalPerf-test.js
 * deduplicates lifecycle names during commit to reduce overhead
 * supports coroutines
 * supports portals
+* does not schedule an extra callback if setState is called during a synchronous commit phase
 
 src/renderers/shared/fiber/__tests__/ReactIncrementalReflection-test.js
 * handles isMounted even when the initial render is deferred

src/renderers/dom/shared/__tests__/ReactMount-test.js

@@ -310,4 +310,39 @@ describe('ReactMount', () => {
 
     expect(calls).toBe(5);
   });
+
+  it('initial mount is sync inside batchedUpdates, but task work is deferred until the end of the batch', () => {
+    var container1 = document.createElement('div');
+    var container2 = document.createElement('div');
+
+    class Foo extends React.Component {
+      state = {active: false};
+      componentDidMount() {
+        this.setState({active: true});
+      }
+      render() {
+        return (
+          <div>{this.props.children + (this.state.active ? '!' : '')}</div>
+        );
+      }
+    }
+
+    ReactDOM.render(<div>1</div>, container1);
+
+    ReactDOM.unstable_batchedUpdates(() => {
+      // Update. Does not flush yet.
+      ReactDOM.render(<div>2</div>, container1);
+      expect(container1.textContent).toEqual('1');
+
+      // Initial mount on another root. Should flush immediately.
+      ReactDOM.render(<Foo>a</Foo>, container2);
+      // The update did not flush yet.
+      expect(container1.textContent).toEqual('1');
+      // The initial mount flushed, but not the update scheduled in cDU.
+      expect(container2.textContent).toEqual('a');
+    });
+    // All updates have flushed.
+    expect(container1.textContent).toEqual('2');
+    expect(container2.textContent).toEqual('a!');
+  });
 });

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -195,6 +195,10 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
   // Keeps track of whether we should should batch sync updates.
   let isBatchingUpdates: boolean = false;
 
+  // This is needed for the weird case where the initial mount is synchronous
+  // even inside batchedUpdates :(
+  let isUnbatchingUpdates: boolean = false;
+
   // The next work in progress fiber that we're currently working on.
   let nextUnitOfWork: Fiber | null = null;
   let nextPriorityLevel: PriorityLevel = NoWork;
@@ -759,7 +763,10 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     }
   }
 
-  function workLoop(deadline: Deadline | null) {
+  function workLoop(
+    minPriorityLevel: PriorityLevel,
+    deadline: Deadline | null,
+  ) {
     // Clear any errors.
     clearErrors();
 
@@ -771,6 +778,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     while (
       nextUnitOfWork !== null &&
       nextPriorityLevel !== NoWork &&
+      nextPriorityLevel <= minPriorityLevel &&
       nextPriorityLevel <= TaskPriority
     ) {
       nextUnitOfWork = performUnitOfWork(nextUnitOfWork);
@@ -788,6 +796,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
         nextUnitOfWork !== null &&
         !deadlineHasExpired &&
         nextPriorityLevel !== NoWork &&
+        nextPriorityLevel <= minPriorityLevel &&
         nextPriorityLevel >= HighPriority
       ) {
         if (deadline.timeRemaining() > timeHeuristicForUnitOfWork) {
@@ -815,7 +824,14 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     }
   }
 
-  function performWork(deadline: Deadline | null) {
+  function performDeferredWork(deadline: Deadline) {
+    performWork(OffscreenPriority, deadline);
+  }
+
+  function performWork(
+    minPriorityLevel: PriorityLevel,
+    deadline: Deadline | null,
+  ) {
     if (__DEV__) {
       startWorkLoopTimer();
     }
@@ -849,10 +865,16 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
       priorityContextBeforeReconciliation = priorityContext;
       let error = null;
       if (__DEV__) {
-        error = invokeGuardedCallback(null, workLoop, null, deadline);
+        error = invokeGuardedCallback(
+          null,
+          workLoop,
+          null,
+          minPriorityLevel,
+          deadline,
+        );
       } else {
         try {
-          workLoop(deadline);
+          workLoop(minPriorityLevel, deadline);
         } catch (e) {
           error = e;
         }
@@ -903,7 +925,10 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
           case TaskPriority:
             // We have remaining synchronous or task work. Keep performing it,
             // regardless of whether we're inside a callback.
-            continue;
+            if (nextPriorityLevel <= minPriorityLevel) {
+              continue;
+            }
+            break;
           case HighPriority:
           case LowPriority:
           case OffscreenPriority:
@@ -914,7 +939,10 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
               hasRemainingAsyncWork = true;
             } else {
               // We are inside a callback.
-              if (!deadlineHasExpired) {
+              if (
+                !deadlineHasExpired &&
+                nextPriorityLevel <= minPriorityLevel
+              ) {
                 // We still have time. Keep working.
                 continue;
               }
@@ -939,7 +967,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     }
     // If there's remaining async work, make sure we schedule another callback.
     if (hasRemainingAsyncWork && !isCallbackScheduled) {
-      scheduleDeferredCallback(performWork);
+      scheduleDeferredCallback(performDeferredWork);
       isCallbackScheduled = true;
     }
 
@@ -1264,11 +1292,34 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
         if (node.tag === HostRoot) {
           const root: FiberRoot = (node.stateNode: any);
           scheduleRoot(root, priorityLevel);
-          if (priorityLevel === SynchronousPriority) {
-            performWork(null);
-          } else if (priorityLevel !== NoWork && !isCallbackScheduled) {
-            scheduleDeferredCallback(performWork);
-            isCallbackScheduled = true;
+          if (!isPerformingWork) {
+            switch (priorityLevel) {
+              case SynchronousPriority:
+                // Perform this update now.
+                if (isUnbatchingUpdates) {
+                  // We're inside unbatchedUpdates, which is inside either
+                  // batchedUpdates or a lifecycle. We should only flush
+                  // synchronous work, not task work.
+                  performWork(SynchronousPriority, null);
+                } else {
+                  // Flush both synchronous and task work.
+                  performWork(TaskPriority, null);
+                }
+                break;
+              case TaskPriority:
+                invariant(
+                  isBatchingUpdates,
+                  'Task updates can only be scheduled as a nested update or ' +
+                    'inside batchedUpdates.',
+                );
+                break;
+              default:
+                // Schedule a callback to perform the work later.
+                if (!isCallbackScheduled) {
+                  scheduleDeferredCallback(performDeferredWork);
+                  isCallbackScheduled = true;
+                }
+            }
           }
         } else {
           if (__DEV__) {
@@ -1335,18 +1386,22 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
       // If we're not already inside a batch, we need to flush any task work
       // that was created by the user-provided function.
       if (!isPerformingWork && !isBatchingUpdates) {
-        performWork(null);
+        performWork(TaskPriority, null);
       }
     }
   }
 
   function unbatchedUpdates<A>(fn: () => A): A {
+    const previousIsUnbatchingUpdates = isUnbatchingUpdates;
     const previousIsBatchingUpdates = isBatchingUpdates;
+    // This is only true if we're nested inside batchedUpdates.
+    isUnbatchingUpdates = isBatchingUpdates;
     isBatchingUpdates = false;
     try {
       return fn();
     } finally {
       isBatchingUpdates = previousIsBatchingUpdates;
+      isUnbatchingUpdates = previousIsUnbatchingUpdates;
     }
   }
 

src/renderers/shared/fiber/__tests__/ReactIncrementalPerf-test.js

@@ -487,4 +487,20 @@ describe('ReactDebugFiberPerf', () => {
     ReactNoop.flush();
     expect(getFlameChart()).toMatchSnapshot();
   });
+
+  it('does not schedule an extra callback if setState is called during a synchronous commit phase', () => {
+    class Component extends React.Component {
+      state = {step: 1};
+      componentDidMount() {
+        this.setState({step: 2});
+      }
+      render() {
+        return <span prop={this.state.step} />;
+      }
+    }
+    ReactNoop.syncUpdates(() => {
+      ReactNoop.render(<Component />);
+    });
+    expect(getFlameChart()).toMatchSnapshot();
+  });
 });

src/renderers/shared/fiber/__tests__/__snapshots__/ReactIncrementalPerf-test.js.snap

@@ -67,6 +67,20 @@ exports[`ReactDebugFiberPerf deduplicates lifecycle names during commit to reduc
 "
 `;
 
+exports[`ReactDebugFiberPerf does not schedule an extra callback if setState is called during a synchronous commit phase 1`] = `
+"⛔ (React Tree Reconciliation) Warning: There were cascading updates
+  ⚛ Component [mount]
+  ⛔ (Committing Changes) Warning: Lifecycle hook scheduled a cascading update
+    ⚛ (Committing Host Effects: 1 Total)
+    ⚛ (Calling Lifecycle Methods: 1 Total)
+      ⛔ Component.componentDidMount Warning: Scheduled a cascading update
+  ⚛ Component [update]
+  ⛔ (Committing Changes) Warning: Caused by a cascading update in earlier commit
+    ⚛ (Committing Host Effects: 1 Total)
+    ⚛ (Calling Lifecycle Methods: 1 Total)
+"
+`;
+
 exports[`ReactDebugFiberPerf does not treat setState from cWM or cWRP as cascading 1`] = `
 "// Should not print a warning
 ⚛ (React Tree Reconciliation)