Add API to force Scheduler to yield for macrotask (#25044)

acdlite · web-flow · commit ca990e9a7535 · 2022-08-04T21:20:15.000-04:00
We need a way to yield control of the main thread and schedule a
continuation in a separate macrotask. (This is related to some Suspense
optimizations we have planned.)

Our solution needs account for how Scheduler is implemented. Scheduler
tasks are not 1:1 with real browser macrotasks — many Scheduler "tasks"
can be executed within a single browser task. If a Scheduler task yields
control and posts a continuation, but there's still time left in the
frame, Scheduler will execute the continuation immediately
(synchronously) without yielding control back to the main thread. That's
not what we want — we want to schedule a new macrotask regardless of
where we are in the browser's render cycle.

There are several ways we could approach this. What I ended up doing was
adding a new Scheduler method `unstable_requestYield`. (It's similar to
the existing `unstable_requestPaint` that we use to yield at the end of
the frame.)

It works by setting the internal start time of the current work loop to
a large negative number, so that when the `shouldYield` call computes
how much time has elapsed, it's guaranteed to exceed the deadline. The
advantage of doing it this way is that there are no additional checks in
the normal hot path of the work loop.

The existing layering between Scheduler and React DOM is not ideal. None
of the APIs are public, so despite the fact that Scheduler is a separate
package, I consider that a private implementation detail, and think of
them as part of the same unit.

So for now, though, I think it makes sense to implement this macrotask
logic directly inside of Scheduler instead of layering it on top.

The rough eventual plan for Scheduler is turn it into a `postTask`
prollyfill. Because `postTask` does not yet have an equivalent for
`shouldYield`, we would split that out into its own layer, perhaps
directly inside the reconciler. In that world, the macrotask logic I've
added in this commit would likely live in that same layer. When the
native `postTask` is available, we may not even need any additional
logic because it uses actual browser tasks.
diff --git a/packages/scheduler/npm/umd/scheduler.development.js b/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,
diff --git a/packages/scheduler/npm/umd/scheduler.production.min.js b/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,
diff --git a/packages/scheduler/npm/umd/scheduler.profiling.min.js b/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,
diff --git a/packages/scheduler/src/__tests__/Scheduler-test.js b/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',
+    ]);
+  });
 });
diff --git a/packages/scheduler/src/__tests__/SchedulerMock-test.js b/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',
+      ]);
+    });
   });
 });
diff --git a/packages/scheduler/src/__tests__/SchedulerPostTask-test.js b/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',
+    ]);
+  });
 });
diff --git a/packages/scheduler/src/forks/Scheduler.js b/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,
diff --git a/packages/scheduler/src/forks/SchedulerMock.js b/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,
diff --git a/packages/scheduler/src/forks/SchedulerPostTask.js b/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,
 ) =>
