Add API to force Scheduler to yield for macrotask

packages/scheduler/npm/umd/scheduler.development.js

@@ -54,6 +54,13 @@
     );
   }
 
+  function unstable_requestYield() {
+    return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_requestYield.apply(
+      this,
+      arguments
+    );
+  }
+
   function unstable_runWithPriority() {
     return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_runWithPriority.apply(
       this,
@@ -116,6 +123,7 @@
     unstable_cancelCallback: unstable_cancelCallback,
     unstable_shouldYield: unstable_shouldYield,
     unstable_requestPaint: unstable_requestPaint,
+    unstable_requestYield: unstable_requestYield,
     unstable_runWithPriority: unstable_runWithPriority,
     unstable_next: unstable_next,
     unstable_wrapCallback: unstable_wrapCallback,

packages/scheduler/npm/umd/scheduler.production.min.js

@@ -54,6 +54,13 @@
     );
   }
 
+  function unstable_requestYield() {
+    return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_requestYield.apply(
+      this,
+      arguments
+    );
+  }
+
   function unstable_runWithPriority() {
     return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_runWithPriority.apply(
       this,
@@ -110,6 +117,7 @@
     unstable_cancelCallback: unstable_cancelCallback,
     unstable_shouldYield: unstable_shouldYield,
     unstable_requestPaint: unstable_requestPaint,
+    unstable_requestYield: unstable_requestYield,
     unstable_runWithPriority: unstable_runWithPriority,
     unstable_next: unstable_next,
     unstable_wrapCallback: unstable_wrapCallback,

packages/scheduler/npm/umd/scheduler.profiling.min.js

@@ -54,6 +54,13 @@
     );
   }
 
+  function unstable_requestYield() {
+    return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_requestYield.apply(
+      this,
+      arguments
+    );
+  }
+
   function unstable_runWithPriority() {
     return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_runWithPriority.apply(
       this,
@@ -110,6 +117,7 @@
     unstable_cancelCallback: unstable_cancelCallback,
     unstable_shouldYield: unstable_shouldYield,
     unstable_requestPaint: unstable_requestPaint,
+    unstable_requestYield: unstable_requestYield,
     unstable_runWithPriority: unstable_runWithPriority,
     unstable_next: unstable_next,
     unstable_wrapCallback: unstable_wrapCallback,

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

@@ -18,6 +18,8 @@ let performance;
 let cancelCallback;
 let scheduleCallback;
 let requestPaint;
+let requestYield;
+let shouldYield;
 let NormalPriority;
 
 // The Scheduler implementation uses browser APIs like `MessageChannel` and
@@ -42,6 +44,8 @@ describe('SchedulerBrowser', () => {
     scheduleCallback = Scheduler.unstable_scheduleCallback;
     NormalPriority = Scheduler.unstable_NormalPriority;
     requestPaint = Scheduler.unstable_requestPaint;
+    requestYield = Scheduler.unstable_requestYield;
+    shouldYield = Scheduler.unstable_shouldYield;
   });
 
   afterEach(() => {
@@ -475,4 +479,49 @@ describe('SchedulerBrowser', () => {
       'Yield at 5ms',
     ]);
   });
+
+  it('requestYield forces a yield immediately', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('Original Task');
+      runtime.log('shouldYield: ' + shouldYield());
+      runtime.log('requestYield');
+      requestYield();
+      runtime.log('shouldYield: ' + shouldYield());
+      return () => {
+        runtime.log('Continuation Task');
+        runtime.log('shouldYield: ' + shouldYield());
+        runtime.log('Advance time past frame deadline');
+        runtime.advanceTime(10000);
+        runtime.log('shouldYield: ' + shouldYield());
+      };
+    });
+    runtime.assertLog(['Post Message']);
+
+    runtime.fireMessageEvent();
+    runtime.assertLog([
+      'Message Event',
+      'Original Task',
+      'shouldYield: false',
+      'requestYield',
+      // Immediately after calling requestYield, shouldYield starts
+      // returning true, even though no time has elapsed in the frame
+      'shouldYield: true',
+
+      // The continuation should be scheduled in a separate macrotask.
+      'Post Message',
+    ]);
+
+    // No time has elapsed
+    expect(performance.now()).toBe(0);
+
+    // Subsequent tasks work as normal
+    runtime.fireMessageEvent();
+    runtime.assertLog([
+      'Message Event',
+      'Continuation Task',
+      'shouldYield: false',
+      'Advance time past frame deadline',
+      'shouldYield: true',
+    ]);
+  });
 });

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

@@ -725,5 +725,39 @@ describe('Scheduler', () => {
       scheduleCallback(ImmediatePriority, 42);
       expect(Scheduler).toFlushWithoutYielding();
     });
+
+    it('requestYield forces a yield immediately', () => {
+      scheduleCallback(NormalPriority, () => {
+        Scheduler.unstable_yieldValue('Original Task');
+        Scheduler.unstable_yieldValue(
+          'shouldYield: ' + Scheduler.unstable_shouldYield(),
+        );
+        Scheduler.unstable_yieldValue('requestYield');
+        Scheduler.unstable_requestYield();
+        Scheduler.unstable_yieldValue(
+          'shouldYield: ' + Scheduler.unstable_shouldYield(),
+        );
+        return () => {
+          Scheduler.unstable_yieldValue('Continuation Task');
+          Scheduler.unstable_yieldValue(
+            'shouldYield: ' + Scheduler.unstable_shouldYield(),
+          );
+          Scheduler.unstable_yieldValue('Advance time past frame deadline');
+          Scheduler.unstable_yieldValue(
+            'shouldYield: ' + Scheduler.unstable_shouldYield(),
+          );
+        };
+      });
+
+      // The continuation should be scheduled in a separate macrotask.
+      expect(Scheduler).toFlushUntilNextPaint([
+        'Original Task',
+        'shouldYield: false',
+        'requestYield',
+        // Immediately after calling requestYield, shouldYield starts
+        // returning true
+        'shouldYield: true',
+      ]);
+    });
   });
 });

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

@@ -22,6 +22,8 @@ let NormalPriority;
 let UserBlockingPriority;
 let LowPriority;
 let IdlePriority;
+let shouldYield;
+let requestYield;
 
 // The Scheduler postTask implementation uses a new postTask browser API to
 // schedule work on the main thread. This test suite mocks all browser methods
@@ -44,6 +46,8 @@ describe('SchedulerPostTask', () => {
     NormalPriority = Scheduler.unstable_NormalPriority;
     LowPriority = Scheduler.unstable_LowPriority;
     IdlePriority = Scheduler.unstable_IdlePriority;
+    shouldYield = Scheduler.unstable_shouldYield;
+    requestYield = Scheduler.unstable_requestYield;
   });
 
   afterEach(() => {
@@ -296,4 +300,49 @@ describe('SchedulerPostTask', () => {
       'E',
     ]);
   });
+
+  it('requestYield forces a yield immediately', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('Original Task');
+      runtime.log('shouldYield: ' + shouldYield());
+      runtime.log('requestYield');
+      requestYield();
+      runtime.log('shouldYield: ' + shouldYield());
+      return () => {
+        runtime.log('Continuation Task');
+        runtime.log('shouldYield: ' + shouldYield());
+        runtime.log('Advance time past frame deadline');
+        runtime.advanceTime(10000);
+        runtime.log('shouldYield: ' + shouldYield());
+      };
+    });
+    runtime.assertLog(['Post Task 0 [user-visible]']);
+
+    runtime.flushTasks();
+    runtime.assertLog([
+      'Task 0 Fired',
+      'Original Task',
+      'shouldYield: false',
+      'requestYield',
+      // Immediately after calling requestYield, shouldYield starts
+      // returning true, even though no time has elapsed in the frame
+      'shouldYield: true',
+
+      // The continuation should be scheduled in a separate macrotask.
+      'Post Task 1 [user-visible]',
+    ]);
+
+    // No time has elapsed
+    expect(performance.now()).toBe(0);
+
+    // Subsequent tasks work as normal
+    runtime.flushTasks();
+    runtime.assertLog([
+      'Task 1 Fired',
+      'Continuation Task',
+      'shouldYield: false',
+      'Advance time past frame deadline',
+      'shouldYield: true',
+    ]);
+  });
 });

packages/scheduler/src/forks/Scheduler.js

@@ -495,6 +495,11 @@ function requestPaint() {
   // Since we yield every frame regardless, `requestPaint` has no effect.
 }
 
+function requestYield() {
+  // Force a yield at the next opportunity.
+  startTime = -99999;
+}
+
 function forceFrameRate(fps) {
   if (fps < 0 || fps > 125) {
     // Using console['error'] to evade Babel and ESLint
@@ -598,8 +603,6 @@ function cancelHostTimeout() {
   taskTimeoutID = -1;
 }
 
-const unstable_requestPaint = requestPaint;
-
 export {
   ImmediatePriority as unstable_ImmediatePriority,
   UserBlockingPriority as unstable_UserBlockingPriority,
@@ -613,7 +616,8 @@ export {
   unstable_wrapCallback,
   unstable_getCurrentPriorityLevel,
   shouldYieldToHost as unstable_shouldYield,
-  unstable_requestPaint,
+  requestPaint as unstable_requestPaint,
+  requestYield as unstable_requestYield,
   unstable_continueExecution,
   unstable_pauseExecution,
   unstable_getFirstCallbackNode,

packages/scheduler/src/forks/SchedulerMock.js

@@ -608,6 +608,11 @@ function requestPaint() {
   needsPaint = true;
 }
 
+function requestYield() {
+  // Force a yield at the next opportunity.
+  shouldYieldForPaint = needsPaint = true;
+}
+
 export {
   ImmediatePriority as unstable_ImmediatePriority,
   UserBlockingPriority as unstable_UserBlockingPriority,
@@ -622,6 +627,7 @@ export {
   unstable_getCurrentPriorityLevel,
   shouldYieldToHost as unstable_shouldYield,
   requestPaint as unstable_requestPaint,
+  requestYield as unstable_requestYield,
   unstable_continueExecution,
   unstable_pauseExecution,
   unstable_getFirstCallbackNode,

packages/scheduler/src/forks/SchedulerPostTask.js

@@ -67,6 +67,11 @@ export function unstable_requestPaint() {
   // Since we yield every frame regardless, `requestPaint` has no effect.
 }
 
+export function unstable_requestYield() {
+  // Force a yield at the next opportunity.
+  deadline = -99999;
+}
+
 type SchedulerCallback<T> = (
   didTimeout_DEPRECATED: boolean,
 ) =>