[Scheduler] shouldYield should always return false from inside an expired task

packages/react-dom/src/events/__tests__/ChangeEventPlugin-test.internal.js

@@ -779,6 +779,7 @@ describe('ChangeEventPlugin', () => {
 
           // 3s should be enough to expire the updates
           Scheduler.unstable_advanceTime(3000);
+          expect(Scheduler).toFlushExpired([]);
           expect(container.textContent).toEqual('hovered');
         });
       },

packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.internal.js

@@ -540,7 +540,7 @@ describe('ReactIncrementalUpdates', () => {
 
       // All the updates should render and commit in a single batch.
       Scheduler.unstable_advanceTime(10000);
-      expect(Scheduler).toHaveYielded(['Render: goodbye']);
+      expect(Scheduler).toFlushExpired(['Render: goodbye']);
       // Passive effect
       expect(Scheduler).toFlushAndYield(['Commit: goodbye']);
     });
@@ -645,7 +645,7 @@ describe('ReactIncrementalUpdates', () => {
 
       // All the updates should render and commit in a single batch.
       Scheduler.unstable_advanceTime(10000);
-      expect(Scheduler).toHaveYielded([
+      expect(Scheduler).toFlushExpired([
         'Render: goodbye',
         'Commit: goodbye',
         'Render: goodbye',

packages/react-reconciler/src/__tests__/ReactSuspenseWithNoopRenderer-test.internal.js

@@ -848,7 +848,7 @@ describe('ReactSuspenseWithNoopRenderer', () => {
       </Suspense>,
     );
     Scheduler.unstable_advanceTime(10000);
-    expect(Scheduler).toHaveYielded([
+    expect(Scheduler).toFlushExpired([
       'Suspend! [A]',
       'Suspend! [B]',
       'Loading...',
@@ -987,7 +987,7 @@ describe('ReactSuspenseWithNoopRenderer', () => {
     Scheduler.unstable_advanceTime(10000);
     jest.advanceTimersByTime(10000);
 
-    expect(Scheduler).toHaveYielded([
+    expect(Scheduler).toFlushExpired([
       'Suspend! [goodbye]',
       'Loading...',
       'Commit: goodbye',

packages/scheduler/src/Scheduler.js

@@ -396,18 +396,58 @@ function unstable_getCurrentPriorityLevel() {
 }
 
 function unstable_shouldYield() {
+  // Note: In general, we should bias toward returning `false` instead of
+  // `true`, since in the extreme cases, the consequences of yielding too much
+  // (indefinite starvation) are worse than the consequences of blocking the
+  // main thread (temporary unresponsiveness).
+
+  if (currentTask === null) {
+    // There's no currently running task.
+    // TODO: Is this desirable behavior? Maybe we should pretend unscheduled
+    // work has default priority.
+    return false;
+  }
+
+  // Check if the current task expired.
   const currentTime = getCurrentTime();
+  const currentExpiration = currentTask.expirationTime;
+  if (currentExpiration <= currentTime) {
+    // Never yield from an expired task, even if there is a pending main thread
+    // task or a higher priority task in the queue.
+    return false;
+  }
+
+  // Check if there's a pending main thread task.
+  if (shouldYieldToHost()) {
+    return true;
+  }
+
+  // Check a higher priority task was scheduled. But first, check if any timers
+  // have elapsed since we started working on this task. If so, this transfers
+  // them to the task queue.
   advanceTimers(currentTime);
-  const firstTask = peek(taskQueue);
-  return (
-    (firstTask !== currentTask &&
-      currentTask !== null &&
-      firstTask !== null &&
-      firstTask.callback !== null &&
-      firstTask.startTime <= currentTime &&
-      firstTask.expirationTime < currentTask.expirationTime) ||
-    shouldYieldToHost()
-  );
+
+  // Now compare the priority of the current task to the highest priority task.
+  // This equivalent to checking if the current task is the head of the queue.
+  let firstTask = peek(taskQueue);
+  while (firstTask !== currentTask && firstTask !== null) {
+    if (firstTask.callback !== null) {
+      // There's a higher priority task. Yield.
+      return true;
+    } else {
+      // There's a higher priority task, but it was canceled. This happens
+      // because tasks are not always removed from the queue immediately when
+      // they are canceled, since removal of an arbitrary node from a binary
+      // heap is O(log n). Instead, we null out the `callback` field and remove
+      // it the next time we traverse the queue.
+      //
+      // Pop the canceled task from the queue and proceed to the next task.
+      pop(taskQueue);
+      firstTask = peek(taskQueue);
+    }
+  }
+  // There's no higher priority task.
+  return false;
 }
 
 const unstable_requestPaint = requestPaint;

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

@@ -117,14 +117,14 @@ describe('Scheduler', () => {
 
     // Advance by just a bit more to expire the user blocking callbacks
     Scheduler.unstable_advanceTime(1);
-    expect(Scheduler).toHaveYielded([
+    expect(Scheduler).toFlushExpired([
       'B (did timeout: true)',
       'C (did timeout: true)',
     ]);
 
     // Expire A
     Scheduler.unstable_advanceTime(4600);
-    expect(Scheduler).toHaveYielded(['A (did timeout: true)']);
+    expect(Scheduler).toFlushExpired(['A (did timeout: true)']);
 
     // Flush the rest without expiring
     expect(Scheduler).toFlushAndYield([
@@ -140,7 +140,7 @@ describe('Scheduler', () => {
     expect(Scheduler).toHaveYielded([]);
 
     Scheduler.unstable_advanceTime(1);
-    expect(Scheduler).toHaveYielded(['A']);
+    expect(Scheduler).toFlushExpired(['A']);
   });
 
   it('continues working on same task after yielding', () => {
@@ -191,6 +191,180 @@ describe('Scheduler', () => {
     expect(Scheduler).toFlushAndYield(['C2', 'C3', 'D', 'E']);
   });
 
+  it("shouldYield returns false if task has expired, even if there's a pending higher priority task", () => {
+    let step = 0;
+    function workLoop() {
+      while (step < 4) {
+        step += 1;
+        Scheduler.unstable_yieldValue('Step ' + step);
+        if (step === 2) {
+          // This schedules a high pri task
+          scheduleCallback(Scheduler.unstable_ImmediatePriority, () => {
+            Scheduler.unstable_yieldValue('High-pri task');
+          });
+          // This expires the current task
+          Scheduler.unstable_yieldValue('Expire!');
+          Scheduler.unstable_advanceTime(10000);
+        }
+        if (shouldYield()) {
+          // We should never hit this path, because should have already expired.
+          Scheduler.unstable_yieldValue('Yield!');
+          return workLoop;
+        }
+      }
+    }
+
+    scheduleCallback(NormalPriority, workLoop);
+    expect(Scheduler).toFlushAndYield([
+      'Step 1',
+      'Step 2',
+      'Expire!',
+      'Step 3',
+      'Step 4',
+      'High-pri task',
+    ]);
+  });
+
+  it(
+    "shouldYield returns false if task has expired, even if there's a " +
+      'pending main thread task (e.g. requestPaint)',
+    () => {
+      let step = 0;
+      function workLoop() {
+        while (step < 4) {
+          step += 1;
+          Scheduler.unstable_yieldValue('Step ' + step);
+          if (step === 2) {
+            // This asks Scheduler to yield to the main thread
+            Scheduler.unstable_requestPaint();
+            // This expires the current task
+            Scheduler.unstable_yieldValue('Expire!');
+            Scheduler.unstable_advanceTime(10000);
+          }
+          if (shouldYield()) {
+            // We should never hit this path, because should have already expired.
+            Scheduler.unstable_yieldValue('Yield!');
+            return workLoop;
+          }
+        }
+      }
+
+      scheduleCallback(NormalPriority, workLoop);
+      expect(Scheduler).toFlushUntilNextPaint([
+        'Step 1',
+        'Step 2',
+        'Expire!',
+        'Step 3',
+        'Step 4',
+      ]);
+    },
+  );
+
+  it(
+    'shouldYield returns false if there are higher priority tasks that were ' +
+      'already canceled.',
+    () => {
+      let step = 0;
+      function workLoop() {
+        while (step < 4) {
+          step += 1;
+          Scheduler.unstable_yieldValue('Step ' + step);
+          if (step === 2) {
+            // This asks Scheduler to yield to the main thread
+            const task1 = scheduleCallback(
+              Scheduler.unstable_ImmediatePriority,
+              () => {
+                Scheduler.unstable_yieldValue('High-pri task 1');
+              },
+            );
+            const task2 = scheduleCallback(
+              Scheduler.unstable_ImmediatePriority,
+              () => {
+                Scheduler.unstable_yieldValue('High-pri task 2');
+              },
+            );
+            cancelCallback(task1);
+            cancelCallback(task2);
+          }
+          if (shouldYield()) {
+            // We should never hit this path, because should have already expired.
+            Scheduler.unstable_yieldValue('Yield!');
+            return workLoop;
+          }
+        }
+      }
+
+      scheduleCallback(NormalPriority, workLoop);
+      expect(Scheduler).toFlushAndYield([
+        'Step 1',
+        'Step 2',
+        'Step 3',
+        'Step 4',
+      ]);
+    },
+  );
+
+  it(
+    'shouldYield returns true if there are higher priority tasks, some of ' +
+      'which were canceled, but not all of them',
+    () => {
+      // Tests an implementation-specific edge case. Canceled are not always
+      // removed from the queue when they are canceled, since removal of an
+      // arbitrary node from a binary heap is O(log n).
+      let step = 0;
+      function workLoop() {
+        while (step < 4) {
+          step += 1;
+          Scheduler.unstable_yieldValue('Step ' + step);
+          if (step === 2) {
+            // This asks Scheduler to yield to the main thread
+            const task1 = scheduleCallback(
+              Scheduler.unstable_ImmediatePriority,
+              () => {
+                Scheduler.unstable_yieldValue('High-pri task 1');
+              },
+            );
+            scheduleCallback(Scheduler.unstable_ImmediatePriority, () => {
+              Scheduler.unstable_yieldValue('High-pri task 2');
+            });
+            // Cancel the first task, but not the second one.
+            cancelCallback(task1);
+          }
+          if (shouldYield()) {
+            // We should never hit this path, because should have already expired.
+            Scheduler.unstable_yieldValue('Yield!');
+            return workLoop;
+          }
+        }
+      }
+
+      scheduleCallback(NormalPriority, workLoop);
+      expect(Scheduler).toFlushUntilNextPaint([
+        'Step 1',
+        'Step 2',
+        'Yield!',
+        'High-pri task 2',
+        'Step 3',
+        'Step 4',
+      ]);
+    },
+  );
+
+  it('shouldYield returns false if there are no pending tasks', () => {
+    expect(Scheduler.unstable_shouldYield()).toBe(false);
+  });
+
+  // TODO: Is this desirable behavior? Maybe we should pretend unscheduled
+  // work has default priority.
+  it(
+    'shouldYield returns false if called outside of a Scheduler task, even ' +
+      'if there are pending tasks',
+    () => {
+      scheduleCallback(ImmediatePriority, () => {});
+      expect(Scheduler.unstable_shouldYield()).toBe(false);
+    },
+  );
+
   it('continuation callbacks inherit the expiration of the previous callback', () => {
     const tasks = [
       ['A', 125],
@@ -217,7 +391,7 @@ describe('Scheduler', () => {
 
     // Advance time by just a bit more. This should expire all the remaining work.
     Scheduler.unstable_advanceTime(1);
-    expect(Scheduler).toHaveYielded(['C', 'D']);
+    expect(Scheduler).toFlushExpired(['C', 'D']);
   });
 
   it('continuations are interrupted by higher priority work', () => {
@@ -705,7 +879,7 @@ describe('Scheduler', () => {
 
       // Now it expires
       Scheduler.unstable_advanceTime(1);
-      expect(Scheduler).toHaveYielded(['A']);
+      expect(Scheduler).toFlushExpired(['A']);
     });
 
     it('cancels a delayed task', () => {

packages/scheduler/src/forks/SchedulerHostConfig.mock.js

@@ -201,13 +201,10 @@ export function unstable_yieldValue(value: mixed): void {
 
 export function unstable_advanceTime(ms: number) {
   currentTime += ms;
-  if (!isFlushing) {
-    if (scheduledTimeout !== null && timeoutTime <= currentTime) {
-      scheduledTimeout(currentTime);
-      timeoutTime = -1;
-      scheduledTimeout = null;
-    }
-    unstable_flushExpired();
+  if (scheduledTimeout !== null && timeoutTime <= currentTime) {
+    scheduledTimeout(currentTime);
+    timeoutTime = -1;
+    scheduledTimeout = null;
   }
 }