[Scheduler] Canceling an already running task

Canceling an already running task is currently a no-op. It does not
cause the task to yield execution, and if the task does yield with
a continuation, then the continuation will run.

This change addresses both parts: canceling the current task causes
`shouldYield` to return `true`, and if the canceled task returns a
continuation, the continuation is ignored.

This fixes the regression test introduced in the preceding commit.
There's likely a better way to model a restart of an in-progress
render, but this approach is conceptually similar to how a regular
high priority interruption already works.

Author: acdlite

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -2046,8 +2046,16 @@ function updateDehydratedSuspenseComponent(
         // at even higher pri.
         let attemptHydrationAtExpirationTime = renderExpirationTime + 1;
         suspenseState.retryTime = attemptHydrationAtExpirationTime;
+        // TODO: This happens to abort the current render and switch to a
+        // hydration pass, because the priority of the new task is slightly
+        // higher priority, which causes the work loop to cancel the current
+        // Scheduler task via `Scheduler.cancelCallback`. But we should probably
+        // model this entirely within React instead of relying on Scheduler's
+        // semantics for canceling in-progress tasks. I've chosen this approach
+        // for now since it's a fairly non-invasive change and it conceptually
+        // matches how other types of interuptions (e.g. due to input events)
+        // already work.
         scheduleWork(current, attemptHydrationAtExpirationTime);
-        // TODO: Early abort this render.
       } else {
         // We have already tried to ping at a higher priority than we're rendering with
         // so if we got here, we must have failed to hydrate at those levels. We must

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -515,6 +515,9 @@ function scheduleCallbackForRoot(
     // New callback has higher priority than the existing one.
     const existingCallbackNode = root.callbackNode;
     if (existingCallbackNode !== null) {
+      // If this happens during render, this task represents the currently
+      // running task. Canceling has the effect of interrupting the render and
+      // starting over at the higher priority.
       cancelCallback(existingCallbackNode);
     }
     root.callbackExpirationTime = expirationTime;

packages/scheduler/src/Scheduler.js

@@ -75,6 +75,8 @@ var currentPriorityLevel = NormalPriority;
 // This is set while performing work, to prevent re-entrancy.
 var isPerformingWork = false;
 
+var didCancelCurrentTask = false;
+
 var isHostCallbackScheduled = false;
 var isHostTimeoutScheduled = false;
 
@@ -184,16 +186,20 @@ function workLoop(hasTimeRemaining, initialTime) {
       markTaskRun(currentTask, currentTime);
       const continuationCallback = callback(didUserCallbackTimeout);
       currentTime = getCurrentTime();
-      if (typeof continuationCallback === 'function') {
-        currentTask.callback = continuationCallback;
-        markTaskYield(currentTask, currentTime);
+      if (didCancelCurrentTask) {
+        didCancelCurrentTask = false;
       } else {
-        if (enableProfiling) {
-          markTaskCompleted(currentTask, currentTime);
-          currentTask.isQueued = false;
-        }
-        if (currentTask === peek(taskQueue)) {
-          pop(taskQueue);
+        if (typeof continuationCallback === 'function') {
+          currentTask.callback = continuationCallback;
+          markTaskYield(currentTask, currentTime);
+        } else {
+          if (enableProfiling) {
+            markTaskCompleted(currentTask, currentTime);
+            currentTask.isQueued = false;
+          }
+          if (currentTask === peek(taskQueue)) {
+            pop(taskQueue);
+          }
         }
       }
       advanceTimers(currentTime);
@@ -389,6 +395,9 @@ function unstable_cancelCallback(task) {
   // remove from the queue because you can't remove arbitrary nodes from an
   // array based heap, only the first one.)
   task.callback = null;
+  if (task === currentTask) {
+    didCancelCurrentTask = true;
+  }
 }
 
 function unstable_getCurrentPriorityLevel() {
@@ -406,6 +415,7 @@ function unstable_shouldYield() {
       firstTask.callback !== null &&
       firstTask.startTime <= currentTime &&
       firstTask.expirationTime < currentTask.expirationTime) ||
+    didCancelCurrentTask ||
     shouldYieldToHost()
   );
 }

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

@@ -288,6 +288,26 @@ describe('Scheduler', () => {
     expect(Scheduler).toFlushWithoutYielding();
   });
 
+  it('cancelling the currently running task', () => {
+    const task = scheduleCallback(NormalPriority, () => {
+      Scheduler.unstable_yieldValue('Start');
+      for (let i = 0; !shouldYield(); i++) {
+        if (i === 5) {
+          // Canceling the current task will cause `shouldYield` to return
+          // `true`. Otherwise this would infinite loop.
+          Scheduler.unstable_yieldValue('Cancel');
+          cancelCallback(task);
+        }
+      }
+      Scheduler.unstable_yieldValue('Finish');
+      // The continuation should be ignored, since the task was
+      // already canceled.
+      return () => Scheduler.unstable_yieldValue('Continuation');
+    });
+
+    expect(Scheduler).toFlushAndYield(['Start', 'Cancel', 'Finish']);
+  });
+
   it('top-level immediate callbacks fire in a subsequent task', () => {
     scheduleCallback(ImmediatePriority, () =>
       Scheduler.unstable_yieldValue('A'),