Process completion callbacks immediately after completing a root

We need to process completion callbacks in two places. The first is
intuitive: right after a root completes. It might seem like that is
sufficient. But if a completion callback is scheduled on an already
completed root, it's possible we won't complete that root again. So
we also need to process completion callbacks whenever we skip over an
already completed root.

Author: acdlite

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -398,33 +398,10 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
       // no longer blocked, return the time at which it completed so that we
       // can commit it.
       if (isRootBlocked(root, completedAt)) {
-        // Process pending completion callbacks so that they are called at
-        // the end of the current batch.
-        const completionCallbacks = root.completionCallbacks;
-        if (completionCallbacks !== null) {
-          processUpdateQueue(
-            completionCallbacks,
-            null,
-            null,
-            null,
-            completedAt,
-          );
-          const callbackList = completionCallbacks.callbackList;
-          if (callbackList !== null) {
-            // Add new callbacks to list of completion callbacks
-            if (rootCompletionCallbackList === null) {
-              rootCompletionCallbackList = callbackList;
-            } else {
-              for (let i = 0; i < callbackList.length; i++) {
-                rootCompletionCallbackList.push(callbackList[i]);
-              }
-            }
-            completionCallbacks.callbackList = null;
-            if (completionCallbacks.first === null) {
-              root.completionCallbacks = null;
-            }
-          }
-        }
+        // We usually process completion callbacks right after a root is
+        // completed. But this root already completed, and it's possible that
+        // we received new completion callbacks since then.
+        processCompletionCallbacks(root, completedAt);
         return Done;
       }
 
@@ -434,6 +411,33 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     return expirationTime;
   }
 
+  function processCompletionCallbacks(
+    root: FiberRoot,
+    completedAt: ExpirationTime,
+  ) {
+    // Process pending completion callbacks so that they are called at
+    // the end of the current batch.
+    const completionCallbacks = root.completionCallbacks;
+    if (completionCallbacks !== null) {
+      processUpdateQueue(completionCallbacks, null, null, null, completedAt);
+      const callbackList = completionCallbacks.callbackList;
+      if (callbackList !== null) {
+        // Add new callbacks to list of completion callbacks
+        if (rootCompletionCallbackList === null) {
+          rootCompletionCallbackList = callbackList;
+        } else {
+          for (let i = 0; i < callbackList.length; i++) {
+            rootCompletionCallbackList.push(callbackList[i]);
+          }
+        }
+        completionCallbacks.callbackList = null;
+        if (completionCallbacks.first === null) {
+          root.completionCallbacks = null;
+        }
+      }
+    }
+  }
+
   function commitAllHostEffects() {
     while (nextEffect !== null) {
       if (__DEV__) {
@@ -819,6 +823,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
           // The root is not blocked, so we can commit it now.
           pendingCommit = workInProgress;
         }
+        processCompletionCallbacks(root, nextRenderExpirationTime);
         return null;
       }
     }
@@ -1601,12 +1606,9 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
                 }
                 break;
               case TaskPriority:
-                invariant(
-                  isBatchingUpdates,
-                  'Task updates can only be scheduled as a nested update or ' +
-                    'inside batchedUpdates. This error is likely caused by a ' +
-                    'bug in React. Please file an issue.',
-                );
+                if (!isPerformingWork && !isBatchingUpdates) {
+                  performWork(TaskPriority, null);
+                }
                 break;
               default:
                 // This update is async. Schedule a callback.

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

@@ -95,7 +95,62 @@ describe('ReactIncrementalRoot', () => {
     expect(root.getChildren()).toEqual([]);
   });
 
-  it(
-    'does not work on on a blocked tree if the expiration time is greater than the blocked update',
-  );
+  it('does not work on on a blocked tree if the expiration time is greater than the blocked update', () => {
+    let ops = [];
+    function Foo(props) {
+      ops.push('Foo: ' + props.children);
+      return <span prop={props.children} />;
+    }
+    const root = ReactNoop.createRoot();
+    root.prerender(<Foo>A</Foo>);
+    ReactNoop.flush();
+
+    expect(ops).toEqual(['Foo: A']);
+    expect(root.getChildren()).toEqual([]);
+
+    // workB has a later expiration time
+    ReactNoop.expire(1000);
+    root.prerender(<Foo>B</Foo>);
+    ReactNoop.flush();
+
+    // Should not have re-rendered the root at the later expiration time
+    expect(ops).toEqual(['Foo: A']);
+    expect(root.getChildren()).toEqual([]);
+  });
+
+  it('commits earlier work without committing later work', () => {
+    const root = ReactNoop.createRoot();
+    const work1 = root.prerender(<span prop="A" />);
+    ReactNoop.flush();
+
+    expect(root.getChildren()).toEqual([]);
+
+    // Second prerender has a later expiration time
+    ReactNoop.expire(1000);
+    root.prerender(<span prop="B" />);
+
+    work1.commit();
+
+    // Should not have re-rendered the root at the later expiration time
+    expect(root.getChildren()).toEqual([span('A')]);
+  });
+
+  it('flushes ealier work if later work is committed', () => {
+    let ops = [];
+    const root = ReactNoop.createRoot();
+    const work1 = root.prerender(<span prop="A" />);
+    // Second prerender has a later expiration time
+    ReactNoop.expire(1000);
+    const work2 = root.prerender(<span prop="B" />);
+
+    work1.then(() => ops.push('complete 1'));
+    work2.then(() => ops.push('complete 2'));
+
+    work2.commit();
+
+    // Because the later prerender was committed, the earlier one should have
+    // committed, too.
+    expect(root.getChildren()).toEqual([span('B')]);
+    expect(ops).toEqual(['complete 1', 'complete 2']);
+  });
 });