integrate scheduler.yield in SchedulerPostTask

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

@@ -94,6 +94,22 @@ describe('SchedulerPostTask', () => {
       });
     };
 
+    scheduler.yield = function ({priority, signal}) {
+      const id = idCounter++;
+      log(`Yield ${id} [${priority === undefined ? '<default>' : priority}]`);
+      const controller = signal._controller;
+      let callback;
+
+      return {
+        then(cb) {
+          callback = cb;
+          return new Promise((resolve, reject) => {
+            taskQueue.set(controller, {id, callback, resolve, reject});
+          });
+        },
+      };
+    };
+
     global.TaskController = class TaskController {
       constructor() {
         this.signal = {_controller: this};
@@ -178,7 +194,7 @@ describe('SchedulerPostTask', () => {
       'Task 0 Fired',
       'A',
       'Yield at 5ms',
-      'Post Task 1 [user-visible]',
+      'Yield 1 [user-visible]',
     ]);
 
     runtime.flushTasks();
@@ -321,7 +337,7 @@ describe('SchedulerPostTask', () => {
 
       // The continuation should be scheduled in a separate macrotask even
       // though there's time remaining.
-      'Post Task 1 [user-visible]',
+      'Yield 1 [user-visible]',
     ]);
 
     // No time has elapsed
@@ -330,4 +346,67 @@ describe('SchedulerPostTask', () => {
     runtime.flushTasks();
     runtime.assertLog(['Task 1 Fired', 'Continuation Task']);
   });
+
+  describe('falls back to postTask for scheduling continuations when scheduler.yield is not available', () => {
+    beforeEach(() => {
+      delete global.scheduler.yield;
+    });
+
+    it('task with continuation', () => {
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('A');
+        while (!Scheduler.unstable_shouldYield()) {
+          runtime.advanceTime(1);
+        }
+        runtime.log(`Yield at ${performance.now()}ms`);
+        return () => {
+          runtime.log('Continuation');
+        };
+      });
+      runtime.assertLog(['Post Task 0 [user-visible]']);
+
+      runtime.flushTasks();
+      runtime.assertLog([
+        'Task 0 Fired',
+        'A',
+        'Yield at 5ms',
+        'Post Task 1 [user-visible]',
+      ]);
+
+      runtime.flushTasks();
+      runtime.assertLog(['Task 1 Fired', 'Continuation']);
+    });
+
+    it('yielding continues in a new task regardless of how much time is remaining', () => {
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('Original Task');
+        runtime.log('shouldYield: ' + shouldYield());
+        runtime.log('Return a continuation');
+        return () => {
+          runtime.log('Continuation Task');
+        };
+      });
+      runtime.assertLog(['Post Task 0 [user-visible]']);
+
+      runtime.flushTasks();
+      runtime.assertLog([
+        'Task 0 Fired',
+        'Original Task',
+        // Immediately before returning a continuation, `shouldYield` returns
+        // false, which means there must be time remaining in the frame.
+        'shouldYield: false',
+        'Return a continuation',
+
+        // The continuation should be scheduled in a separate macrotask even
+        // though there's time remaining.
+        'Post Task 1 [user-visible]',
+      ]);
+
+      // No time has elapsed
+      expect(performance.now()).toBe(0);
+
+      runtime.flushTasks();
+      runtime.assertLog(['Task 1 Fired', 'Continuation Task']);
+    });
+  });
 });

packages/scheduler/src/forks/SchedulerPostTask.js

@@ -138,18 +138,25 @@ function runTask<T>(
       // Update the original callback node's controller, since even though we're
       // posting a new task, conceptually it's the same one.
       node._controller = continuationController;
-      scheduler
-        .postTask(
-          runTask.bind(
-            null,
-            priorityLevel,
-            postTaskPriority,
-            node,
-            continuation,
-          ),
-          continuationOptions,
-        )
-        .catch(handleAbortError);
+
+      const nextTask = runTask.bind(
+        null,
+        priorityLevel,
+        postTaskPriority,
+        node,
+        continuation,
+      );
+
+      if (scheduler.yield !== undefined) {
+        scheduler
+          .yield(continuationOptions)
+          .then(nextTask)
+          .catch(handleAbortError);
+      } else {
+        scheduler
+          .postTask(nextTask, continuationOptions)
+          .catch(handleAbortError);
+      }
     }
   } catch (error) {
     // We're inside a `postTask` promise. If we don't handle this error, then it