[scheduler] Pass didTimeout argument to callbacks (#14931)

As I prepare to refactor the Fiber scheduler, I've noticed some quirks
in our implementation. This PR addressed one of them.

---

There's no reason for a timed out Scheduler callback to check
`shouldYield`, because the value will always be false until the work
has completed. The `didTimeout` argument provides this information to
the callback so it can avoid the redundant checks.

React's existing check for whether a callback has timed out didn't make
any sense, but happened to work anyway. I don't think the wrongness of
the old implementation is observable via React APIs but it's
incoherent regardless.

packages/react-noop-renderer/src/createReactNoop.js

@@ -540,31 +540,37 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
   let unitsRemaining;
 
   function shouldYield() {
+    // Check if we already yielded
+    if (didYield || yieldedValues !== null) {
+      return true;
+    }
+
+    // If there are no remaining units of work, and we haven't timed out, then
+    // we should yield.
     if (
-      scheduledCallbackTimeout === -1 ||
-      elapsedTimeInMs > scheduledCallbackTimeout
+      unitsRemaining-- <= 0 &&
+      (scheduledCallbackTimeout === -1 ||
+        elapsedTimeInMs < scheduledCallbackTimeout)
     ) {
-      return false;
-    } else {
-      if (didYield || yieldedValues !== null) {
-        return true;
-      }
-      if (unitsRemaining-- > 0) {
-        return false;
-      }
       didYield = true;
       return true;
     }
+
+    // Otherwise, keep working.
+    return false;
   }
 
   function* flushUnitsOfWork(n: number): Generator<Array<mixed>, void, void> {
-    unitsRemaining = n + 1;
+    unitsRemaining = n;
     didYield = false;
     try {
       while (!didYield && scheduledCallback !== null) {
         let cb = scheduledCallback;
         scheduledCallback = null;
-        cb();
+        const didTimeout =
+          scheduledCallbackTimeout !== -1 &&
+          scheduledCallbackTimeout < elapsedTimeInMs;
+        cb(didTimeout);
         if (yieldedValues !== null) {
           const values = yieldedValues;
           yieldedValues = null;

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -2219,9 +2219,9 @@ function shouldYieldToRenderer() {
   return false;
 }
 
-function performAsyncWork() {
+function performAsyncWork(didTimeout) {
   try {
-    if (!shouldYieldToRenderer()) {
+    if (didTimeout) {
       // The callback timed out. That means at least one update has expired.
       // Iterate through the root schedule. If they contain expired work, set
       // the next render expiration time to the current time. This has the effect

packages/scheduler/src/Scheduler.js

@@ -94,7 +94,7 @@ function flushFirstCallback() {
   currentExpirationTime = expirationTime;
   var continuationCallback;
   try {
-    continuationCallback = callback();
+    continuationCallback = callback(currentDidTimeout);
   } finally {
     currentPriorityLevel = previousPriorityLevel;
     currentExpirationTime = previousExpirationTime;

packages/scheduler/src/__tests__/Scheduler-test.js

@@ -246,26 +246,48 @@ describe('Scheduler', () => {
   });
 
   it('expires work', () => {
-    scheduleCallback(() => doWork('A', 100));
+    scheduleCallback(didTimeout =>
+      doWork(`A (did timeout: ${didTimeout})`, 100),
+    );
     runWithPriority(UserBlockingPriority, () => {
-      scheduleCallback(() => doWork('B', 100));
+      scheduleCallback(didTimeout =>
+        doWork(`B (did timeout: ${didTimeout})`, 100),
+      );
     });
-    scheduleCallback(() => doWork('C', 100));
     runWithPriority(UserBlockingPriority, () => {
-      scheduleCallback(() => doWork('D', 100));
+      scheduleCallback(didTimeout =>
+        doWork(`C (did timeout: ${didTimeout})`, 100),
+      );
     });
 
     // Advance time, but not by enough to expire any work
     advanceTime(249);
     expect(clearYieldedValues()).toEqual([]);
 
-    // Advance by just a bit more to expire the high pri callbacks
+    // Schedule a few more callbacks
+    scheduleCallback(didTimeout =>
+      doWork(`D (did timeout: ${didTimeout})`, 100),
+    );
+    scheduleCallback(didTimeout =>
+      doWork(`E (did timeout: ${didTimeout})`, 100),
+    );
+
+    // Advance by just a bit more to expire the user blocking callbacks
     advanceTime(1);
-    expect(clearYieldedValues()).toEqual(['B', 'D']);
+    expect(clearYieldedValues()).toEqual([
+      'B (did timeout: true)',
+      'C (did timeout: true)',
+    ]);
 
-    // Expire the rest
-    advanceTime(10000);
-    expect(clearYieldedValues()).toEqual(['A', 'C']);
+    // Expire A
+    advanceTime(4600);
+    expect(clearYieldedValues()).toEqual(['A (did timeout: true)']);
+
+    // Flush the rest without expiring
+    expect(flushWork()).toEqual([
+      'D (did timeout: false)',
+      'E (did timeout: false)',
+    ]);
   });
 
   it('has a default expiration of ~5 seconds', () => {