[scheduler] Deadline object -> shouldYield (#14025)

* [scheduler] Deadline object -> shouldYield

Instead of using a requestIdleCallback-style deadline object, expose a
method Scheduler.shouldYield that returns true if there's a higher
priority event in the queue.

* Nits

packages/react-art/src/ReactARTHostConfig.js

@@ -8,6 +8,7 @@
 export {
   unstable_now as now,
   unstable_scheduleCallback as scheduleDeferredCallback,
+  unstable_shouldYield as shouldYield,
   unstable_cancelCallback as cancelDeferredCallback,
 } from 'scheduler';
 import Transform from 'art/core/transform';

packages/react-dom/src/client/ReactDOMHostConfig.js

@@ -68,6 +68,7 @@ export type NoTimeout = -1;
 export {
   unstable_now as now,
   unstable_scheduleCallback as scheduleDeferredCallback,
+  unstable_shouldYield as shouldYield,
   unstable_cancelCallback as cancelDeferredCallback,
 } from 'scheduler';
 

packages/react-native-renderer/src/ReactFabricHostConfig.js

@@ -332,6 +332,7 @@ export const scheduleDeferredCallback =
   ReactNativeFrameScheduling.scheduleDeferredCallback;
 export const cancelDeferredCallback =
   ReactNativeFrameScheduling.cancelDeferredCallback;
+export const shouldYield = ReactNativeFrameScheduling.shouldYield;
 
 export const scheduleTimeout = setTimeout;
 export const cancelTimeout = clearTimeout;

packages/react-native-renderer/src/ReactNativeFrameScheduling.js

@@ -7,8 +7,6 @@
  * @flow
  */
 
-import type {Deadline} from 'react-reconciler/src/ReactFiberScheduler';
-
 const hasNativePerformanceNow =
   typeof performance === 'object' && typeof performance.now === 'function';
 
@@ -16,16 +14,9 @@ const now = hasNativePerformanceNow
   ? () => performance.now()
   : () => Date.now();
 
-type Callback = (deadline: Deadline) => void;
-
-let scheduledCallback: Callback | null = null;
+let scheduledCallback: (() => mixed) | null = null;
 let frameDeadline: number = 0;
 
-const frameDeadlineObject: Deadline = {
-  timeRemaining: () => frameDeadline - now(),
-  didTimeout: false,
-};
-
 function setTimeoutCallback() {
   // TODO (bvaughn) Hard-coded 5ms unblocks initial async testing.
   // React API probably changing to boolean rather than time remaining.
@@ -36,7 +27,7 @@ function setTimeoutCallback() {
   const callback = scheduledCallback;
   scheduledCallback = null;
   if (callback !== null) {
-    callback(frameDeadlineObject);
+    callback();
   }
 }
 
@@ -44,7 +35,7 @@ function setTimeoutCallback() {
 // This implementation is only intended for short-term use anyway.
 // We also don't implement cancel functionality b'c Fiber doesn't currently need it.
 function scheduleDeferredCallback(
-  callback: Callback,
+  callback: () => mixed,
   options?: {timeout: number},
 ): number {
   // We assume only one callback is scheduled at a time b'c that's how Fiber works.
@@ -58,4 +49,8 @@ function cancelDeferredCallback(callbackID: number) {
   clearTimeout((callbackID: any)); // Timeouts are always numbers on RN
 }
 
-export {now, scheduleDeferredCallback, cancelDeferredCallback};
+function shouldYield() {
+  return frameDeadline <= now();
+}
+
+export {now, scheduleDeferredCallback, cancelDeferredCallback, shouldYield};

packages/react-native-renderer/src/ReactNativeHostConfig.js

@@ -238,6 +238,7 @@ export const scheduleDeferredCallback =
   ReactNativeFrameScheduling.scheduleDeferredCallback;
 export const cancelDeferredCallback =
   ReactNativeFrameScheduling.cancelDeferredCallback;
+export const shouldYield = ReactNativeFrameScheduling.shouldYield;
 
 export const scheduleTimeout = setTimeout;
 export const cancelTimeout = clearTimeout;

packages/react-noop-renderer/src/createReactNoop.js

@@ -319,7 +319,7 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
           scheduledCallbackTimeout === -1 ||
           scheduledCallbackTimeout > newTimeout
         ) {
-          scheduledCallbackTimeout = newTimeout;
+          scheduledCallbackTimeout = elapsedTimeInMs + newTimeout;
         }
       }
       return 0;
@@ -333,6 +333,8 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       scheduledCallbackTimeout = -1;
     },
 
+    shouldYield,
+
     scheduleTimeout: setTimeout,
     cancelTimeout: clearTimeout,
     noTimeout: -1,
@@ -512,35 +514,44 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
 
   let yieldedValues = null;
 
+  let didYield;
   let unitsRemaining;
 
-  function* flushUnitsOfWork(n: number): Generator<Array<mixed>, void, void> {
-    let didStop = false;
-    while (!didStop && scheduledCallback !== null) {
-      let cb = scheduledCallback;
-      scheduledCallback = null;
-      unitsRemaining = n;
-      cb({
-        timeRemaining() {
-          if (yieldedValues !== null) {
-            return 0;
-          }
-          if (unitsRemaining-- > 0) {
-            return 999;
-          }
-          didStop = true;
-          return 0;
-        },
-        didTimeout:
-          scheduledCallbackTimeout !== -1 &&
-          elapsedTimeInMs > scheduledCallbackTimeout,
-      });
-
-      if (yieldedValues !== null) {
-        const values = yieldedValues;
-        yieldedValues = null;
-        yield values;
+  function shouldYield() {
+    if (
+      scheduledCallbackTimeout === -1 ||
+      elapsedTimeInMs > scheduledCallbackTimeout
+    ) {
+      return false;
+    } else {
+      if (didYield || yieldedValues !== null) {
+        return true;
       }
+      if (unitsRemaining-- > 0) {
+        return false;
+      }
+      didYield = true;
+      return true;
+    }
+  }
+
+  function* flushUnitsOfWork(n: number): Generator<Array<mixed>, void, void> {
+    unitsRemaining = n + 1;
+    didYield = false;
+    try {
+      while (!didYield && scheduledCallback !== null) {
+        let cb = scheduledCallback;
+        scheduledCallback = null;
+        cb();
+        if (yieldedValues !== null) {
+          const values = yieldedValues;
+          yieldedValues = null;
+          yield values;
+        }
+      }
+    } finally {
+      unitsRemaining = -1;
+      didYield = false;
     }
   }
 

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -66,18 +66,16 @@ import getComponentName from 'shared/getComponentName';
 import invariant from 'shared/invariant';
 import warningWithoutStack from 'shared/warningWithoutStack';
 
-import {
-  scheduleTimeout,
-  cancelTimeout,
-  noTimeout,
-} from './ReactFiberHostConfig';
-
 import ReactFiberInstrumentation from './ReactFiberInstrumentation';
 import * as ReactCurrentFiber from './ReactCurrentFiber';
 import {
   now,
   scheduleDeferredCallback,
   cancelDeferredCallback,
+  shouldYield,
+  scheduleTimeout,
+  cancelTimeout,
+  noTimeout,
   prepareForCommit,
   resetAfterCommit,
 } from './ReactFiberHostConfig';
@@ -168,11 +166,6 @@ import {
 } from './ReactFiberCommitWork';
 import {Dispatcher} from './ReactFiberDispatcher';
 
-export type Deadline = {
-  timeRemaining(): number,
-  didTimeout: boolean,
-};
-
 export type Thenable = {
   then(resolve: () => mixed, reject?: () => mixed): mixed,
 };
@@ -1205,18 +1198,14 @@ function workLoop(isYieldy) {
       nextUnitOfWork = performUnitOfWork(nextUnitOfWork);
     }
   } else {
-    // Flush asynchronous work until the deadline runs out of time.
-    while (nextUnitOfWork !== null && !shouldYield()) {
+    // Flush asynchronous work until there's a higher priority event
+    while (nextUnitOfWork !== null && !shouldYieldToRenderer()) {
       nextUnitOfWork = performUnitOfWork(nextUnitOfWork);
     }
   }
 }
 
-function renderRoot(
-  root: FiberRoot,
-  isYieldy: boolean,
-  isExpired: boolean,
-): void {
+function renderRoot(root: FiberRoot, isYieldy: boolean): void {
   invariant(
     !isWorking,
     'renderRoot was called recursively. This error is likely caused ' +
@@ -1445,7 +1434,7 @@ function renderRoot(
       // similar to a suspend, but without a timeout because we're not waiting
       // for a promise to resolve.
       !root.didError &&
-      !isExpired
+      isYieldy
     ) {
       root.didError = true;
       const suspendedExpirationTime = (root.nextExpirationTimeToWorkOn = expirationTime);
@@ -1461,7 +1450,7 @@ function renderRoot(
     }
   }
 
-  if (!isExpired && nextLatestAbsoluteTimeoutMs !== -1) {
+  if (isYieldy && nextLatestAbsoluteTimeoutMs !== -1) {
     // The tree was suspended.
     const suspendedExpirationTime = expirationTime;
     markSuspendedPriorityLevel(root, suspendedExpirationTime);
@@ -1891,10 +1880,8 @@ let isRendering: boolean = false;
 let nextFlushedRoot: FiberRoot | null = null;
 let nextFlushedExpirationTime: ExpirationTime = NoWork;
 let lowestPriorityPendingInteractiveExpirationTime: ExpirationTime = NoWork;
-let deadlineDidExpire: boolean = false;
 let hasUnhandledError: boolean = false;
 let unhandledError: mixed | null = null;
-let deadline: Deadline | null = null;
 
 let isBatchingUpdates: boolean = false;
 let isUnbatchingUpdates: boolean = false;
@@ -1913,8 +1900,6 @@ const NESTED_UPDATE_LIMIT = 50;
 let nestedUpdateCount: number = 0;
 let lastCommittedRootDuringThisBatch: FiberRoot | null = null;
 
-const timeHeuristicForUnitOfWork = 1;
-
 function recomputeCurrentRendererTime() {
   const currentTimeMs = now() - originalStartTimeMs;
   currentRendererTime = msToExpirationTime(currentTimeMs);
@@ -1973,7 +1958,7 @@ function onSuspend(
   msUntilTimeout: number,
 ): void {
   root.expirationTime = rootExpirationTime;
-  if (msUntilTimeout === 0 && !shouldYield()) {
+  if (msUntilTimeout === 0 && !shouldYieldToRenderer()) {
     // Don't wait an additional tick. Commit the tree immediately.
     root.pendingCommitExpirationTime = suspendedExpirationTime;
     root.finishedWork = finishedWork;
@@ -2068,7 +2053,7 @@ function requestWork(root: FiberRoot, expirationTime: ExpirationTime) {
       // flush it now.
       nextFlushedRoot = root;
       nextFlushedExpirationTime = Sync;
-      performWorkOnRoot(root, Sync, true);
+      performWorkOnRoot(root, Sync, false);
     }
     return;
   }
@@ -2176,38 +2161,55 @@ function findHighestPriorityRoot() {
   nextFlushedExpirationTime = highestPriorityWork;
 }
 
-function performAsyncWork(dl) {
-  if (dl.didTimeout) {
-    // The callback timed out. That means at least one update has expired.
-    // Iterate through the root schedule. If they contain expired work, set
-    // the next render expiration time to the current time. This has the effect
-    // of flushing all expired work in a single batch, instead of flushing each
-    // level one at a time.
-    if (firstScheduledRoot !== null) {
-      recomputeCurrentRendererTime();
-      let root: FiberRoot = firstScheduledRoot;
-      do {
-        didExpireAtExpirationTime(root, currentRendererTime);
-        // The root schedule is circular, so this is never null.
-        root = (root.nextScheduledRoot: any);
-      } while (root !== firstScheduledRoot);
+// TODO: This wrapper exists because many of the older tests (the ones that use
+// flushDeferredPri) rely on the number of times `shouldYield` is called. We
+// should get rid of it.
+let didYield: boolean = false;
+function shouldYieldToRenderer() {
+  if (didYield) {
+    return true;
+  }
+  if (shouldYield()) {
+    didYield = true;
+    return true;
+  }
+  return false;
+}
+
+function performAsyncWork() {
+  try {
+    if (!shouldYieldToRenderer()) {
+      // The callback timed out. That means at least one update has expired.
+      // Iterate through the root schedule. If they contain expired work, set
+      // the next render expiration time to the current time. This has the effect
+      // of flushing all expired work in a single batch, instead of flushing each
+      // level one at a time.
+      if (firstScheduledRoot !== null) {
+        recomputeCurrentRendererTime();
+        let root: FiberRoot = firstScheduledRoot;
+        do {
+          didExpireAtExpirationTime(root, currentRendererTime);
+          // The root schedule is circular, so this is never null.
+          root = (root.nextScheduledRoot: any);
+        } while (root !== firstScheduledRoot);
+      }
     }
+    performWork(NoWork, true);
+  } finally {
+    didYield = false;
   }
-  performWork(NoWork, dl);
 }
 
 function performSyncWork() {
-  performWork(Sync, null);
+  performWork(Sync, false);
 }
 
-function performWork(minExpirationTime: ExpirationTime, dl: Deadline | null) {
-  deadline = dl;
-
-  // Keep working on roots until there's no more work, or until we reach
-  // the deadline.
+function performWork(minExpirationTime: ExpirationTime, isYieldy: boolean) {
+  // Keep working on roots until there's no more work, or until there's a higher
+  // priority event.
   findHighestPriorityRoot();
 
-  if (deadline !== null) {
+  if (isYieldy) {
     recomputeCurrentRendererTime();
     currentSchedulerTime = currentRendererTime;
 
@@ -2222,12 +2224,12 @@ function performWork(minExpirationTime: ExpirationTime, dl: Deadline | null) {
       nextFlushedExpirationTime !== NoWork &&
       (minExpirationTime === NoWork ||
         minExpirationTime >= nextFlushedExpirationTime) &&
-      (!deadlineDidExpire || currentRendererTime >= nextFlushedExpirationTime)
+      !(didYield && currentRendererTime < nextFlushedExpirationTime)
     ) {
       performWorkOnRoot(
         nextFlushedRoot,
         nextFlushedExpirationTime,
-        currentRendererTime >= nextFlushedExpirationTime,
+        currentRendererTime < nextFlushedExpirationTime,
       );
       findHighestPriorityRoot();
       recomputeCurrentRendererTime();
@@ -2240,7 +2242,7 @@ function performWork(minExpirationTime: ExpirationTime, dl: Deadline | null) {
       (minExpirationTime === NoWork ||
         minExpirationTime >= nextFlushedExpirationTime)
     ) {
-      performWorkOnRoot(nextFlushedRoot, nextFlushedExpirationTime, true);
+      performWorkOnRoot(nextFlushedRoot, nextFlushedExpirationTime, false);
       findHighestPriorityRoot();
     }
   }
@@ -2249,7 +2251,7 @@ function performWork(minExpirationTime: ExpirationTime, dl: Deadline | null) {
   // or there's no more work left with sufficient priority.
 
   // If we're inside a callback, set this to false since we just completed it.
-  if (deadline !== null) {
+  if (isYieldy) {
     callbackExpirationTime = NoWork;
     callbackID = null;
   }
@@ -2262,9 +2264,6 @@ function performWork(minExpirationTime: ExpirationTime, dl: Deadline | null) {
   }
 
   // Clean-up.
-  deadline = null;
-  deadlineDidExpire = false;
-
   finishRendering();
 }
 
@@ -2279,7 +2278,7 @@ function flushRoot(root: FiberRoot, expirationTime: ExpirationTime) {
   // including the given time.
   nextFlushedRoot = root;
   nextFlushedExpirationTime = expirationTime;
-  performWorkOnRoot(root, expirationTime, true);
+  performWorkOnRoot(root, expirationTime, false);
   // Flush any sync work that was scheduled by lifecycles
   performSyncWork();
 }
@@ -2315,7 +2314,7 @@ function finishRendering() {
 function performWorkOnRoot(
   root: FiberRoot,
   expirationTime: ExpirationTime,
-  isExpired: boolean,
+  isYieldy: boolean,
 ) {
   invariant(
     !isRendering,
@@ -2326,7 +2325,7 @@ function performWorkOnRoot(
   isRendering = true;
 
   // Check if this is async work or sync/expired work.
-  if (deadline === null || isExpired) {
+  if (!isYieldy) {
     // Flush work without yielding.
     // TODO: Non-yieldy work does not necessarily imply expired work. A renderer
     // may want to perform some work without yielding, but also without
@@ -2346,8 +2345,7 @@ function performWorkOnRoot(
         // $FlowFixMe Complains noTimeout is not a TimeoutID, despite the check above
         cancelTimeout(timeoutHandle);
       }
-      const isYieldy = false;
-      renderRoot(root, isYieldy, isExpired);
+      renderRoot(root, isYieldy);
       finishedWork = root.finishedWork;
       if (finishedWork !== null) {
         // We've completed the root. Commit it.
@@ -2370,13 +2368,12 @@ function performWorkOnRoot(
         // $FlowFixMe Complains noTimeout is not a TimeoutID, despite the check above
         cancelTimeout(timeoutHandle);
       }
-      const isYieldy = true;
-      renderRoot(root, isYieldy, isExpired);
+      renderRoot(root, isYieldy);
       finishedWork = root.finishedWork;
       if (finishedWork !== null) {
-        // We've completed the root. Check the deadline one more time
+        // We've completed the root. Check the if we should yield one more time
         // before committing.
-        if (!shouldYield()) {
+        if (!shouldYieldToRenderer()) {
           // Still time left. Commit the root.
           completeRoot(root, finishedWork, expirationTime);
         } else {
@@ -2430,24 +2427,6 @@ function completeRoot(
   commitRoot(root, finishedWork);
 }
 
-// When working on async work, the reconciler asks the renderer if it should
-// yield execution. For DOM, we implement this with requestIdleCallback.
-function shouldYield() {
-  if (deadlineDidExpire) {
-    return true;
-  }
-  if (
-    deadline === null ||
-    deadline.timeRemaining() > timeHeuristicForUnitOfWork
-  ) {
-    // Disregard deadline.didTimeout. Only expired work should be flushed
-    // during a timeout. This path is only hit for non-expired work.
-    return false;
-  }
-  deadlineDidExpire = true;
-  return true;
-}
-
 function onUncaughtError(error: mixed) {
   invariant(
     nextFlushedRoot !== null,
@@ -2524,7 +2503,7 @@ function interactiveUpdates<A, B, R>(fn: (A, B) => R, a: A, b: B): R {
     lowestPriorityPendingInteractiveExpirationTime !== NoWork
   ) {
     // Synchronously flush pending interactive updates.
-    performWork(lowestPriorityPendingInteractiveExpirationTime, null);
+    performWork(lowestPriorityPendingInteractiveExpirationTime, false);
     lowestPriorityPendingInteractiveExpirationTime = NoWork;
   }
   const previousIsBatchingInteractiveUpdates = isBatchingInteractiveUpdates;
@@ -2548,7 +2527,7 @@ function flushInteractiveUpdates() {
     lowestPriorityPendingInteractiveExpirationTime !== NoWork
   ) {
     // Synchronously flush pending interactive updates.
-    performWork(lowestPriorityPendingInteractiveExpirationTime, null);
+    performWork(lowestPriorityPendingInteractiveExpirationTime, false);
     lowestPriorityPendingInteractiveExpirationTime = NoWork;
   }
 }

packages/react-reconciler/src/forks/ReactFiberHostConfig.custom.js

@@ -52,6 +52,7 @@ export const shouldDeprioritizeSubtree =
 export const createTextInstance = $$$hostConfig.createTextInstance;
 export const scheduleDeferredCallback = $$$hostConfig.scheduleDeferredCallback;
 export const cancelDeferredCallback = $$$hostConfig.cancelDeferredCallback;
+export const shouldYield = $$$hostConfig.shouldYield;
 export const scheduleTimeout = $$$hostConfig.setTimeout;
 export const cancelTimeout = $$$hostConfig.clearTimeout;
 export const noTimeout = $$$hostConfig.noTimeout;

packages/react-test-renderer/src/ReactTestHostConfig.js

@@ -202,6 +202,7 @@ export const scheduleDeferredCallback =
   TestRendererScheduling.scheduleDeferredCallback;
 export const cancelDeferredCallback =
   TestRendererScheduling.cancelDeferredCallback;
+export const shouldYield = TestRendererScheduling.shouldYield;
 
 export const scheduleTimeout = setTimeout;
 export const cancelTimeout = clearTimeout;

packages/react-test-renderer/src/ReactTestRendererScheduling.js

@@ -7,15 +7,16 @@
  * @flow
  */
 
-import type {Deadline} from 'react-reconciler/src/ReactFiberScheduler';
-
 // Current virtual time
 export let nowImplementation = () => 0;
-export let scheduledCallback: ((deadline: Deadline) => mixed) | null = null;
+export let scheduledCallback: (() => mixed) | null = null;
 export let yieldedValues: Array<mixed> = [];
 
+let didStop: boolean = false;
+let expectedNumberOfYields: number = -1;
+
 export function scheduleDeferredCallback(
-  callback: (deadline: Deadline) => mixed,
+  callback: () => mixed,
   options?: {timeout: number},
 ): number {
   scheduledCallback = callback;
@@ -31,21 +32,25 @@ export function setNowImplementation(implementation: () => number): void {
   nowImplementation = implementation;
 }
 
+export function shouldYield() {
+  if (
+    expectedNumberOfYields !== -1 &&
+    yieldedValues.length >= expectedNumberOfYields
+  ) {
+    // We yielded at least as many values as expected. Stop rendering.
+    didStop = true;
+    return true;
+  }
+  // Keep rendering.
+  return false;
+}
+
 export function flushAll(): Array<mixed> {
   yieldedValues = [];
   while (scheduledCallback !== null) {
     const cb = scheduledCallback;
     scheduledCallback = null;
-    cb({
-      timeRemaining() {
-        // Keep rendering until there's no more work
-        return 999;
-      },
-      // React's scheduler has its own way of keeping track of expired
-      // work and doesn't read this, so don't bother setting it to the
-      // correct value.
-      didTimeout: false,
-    });
+    cb();
   }
   const values = yieldedValues;
   yieldedValues = [];
@@ -53,30 +58,21 @@ export function flushAll(): Array<mixed> {
 }
 
 export function flushNumberOfYields(count: number): Array<mixed> {
-  let didStop = false;
+  expectedNumberOfYields = count;
+  didStop = false;
   yieldedValues = [];
-  while (scheduledCallback !== null && !didStop) {
-    const cb = scheduledCallback;
-    scheduledCallback = null;
-    cb({
-      timeRemaining() {
-        if (yieldedValues.length >= count) {
-          // We at least as many values as expected. Stop rendering.
-          didStop = true;
-          return 0;
-        }
-        // Keep rendering.
-        return 999;
-      },
-      // React's scheduler has its own way of keeping track of expired
-      // work and doesn't read this, so don't bother setting it to the
-      // correct value.
-      didTimeout: false,
-    });
+  try {
+    while (scheduledCallback !== null && !didStop) {
+      const cb = scheduledCallback;
+      scheduledCallback = null;
+      cb();
+    }
+    return yieldedValues;
+  } finally {
+    expectedNumberOfYields = -1;
+    didStop = false;
+    yieldedValues = [];
   }
-  const values = yieldedValues;
-  yieldedValues = [];
-  return values;
 }
 
 export function yieldValue(value: mixed): void {

packages/react/src/ReactSharedInternals.js

@@ -8,6 +8,7 @@
 import assign from 'object-assign';
 import {
   unstable_cancelCallback,
+  unstable_shouldYield,
   unstable_now,
   unstable_scheduleCallback,
   unstable_runWithPriority,
@@ -43,6 +44,7 @@ if (__UMD__) {
   Object.assign(ReactSharedInternals, {
     Scheduler: {
       unstable_cancelCallback,
+      unstable_shouldYield,
       unstable_now,
       unstable_scheduleCallback,
       unstable_runWithPriority,

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

@@ -39,6 +39,13 @@
     );
   }
 
+  function unstable_shouldYield() {
+    return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_shouldYield.apply(
+      this,
+      arguments
+    );
+  }
+
   function unstable_runWithPriority() {
     return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_runWithPriority.apply(
       this,
@@ -64,6 +71,7 @@
     unstable_now: unstable_now,
     unstable_scheduleCallback: unstable_scheduleCallback,
     unstable_cancelCallback: unstable_cancelCallback,
+    unstable_shouldYield: unstable_shouldYield,
     unstable_runWithPriority: unstable_runWithPriority,
     unstable_wrapCallback: unstable_wrapCallback,
     unstable_getCurrentPriorityLevel: unstable_getCurrentPriorityLevel,

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

@@ -39,6 +39,13 @@
     );
   }
 
+  function unstable_shouldYield() {
+    return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_shouldYield.apply(
+      this,
+      arguments
+    );
+  }
+
   function unstable_runWithPriority() {
     return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_runWithPriority.apply(
       this,
@@ -64,6 +71,7 @@
     unstable_now: unstable_now,
     unstable_scheduleCallback: unstable_scheduleCallback,
     unstable_cancelCallback: unstable_cancelCallback,
+    unstable_shouldYield: unstable_shouldYield,
     unstable_runWithPriority: unstable_runWithPriority,
     unstable_wrapCallback: unstable_wrapCallback,
     unstable_getCurrentPriorityLevel: unstable_getCurrentPriorityLevel,

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

@@ -39,6 +39,13 @@
     );
   }
 
+  function unstable_shouldYield() {
+    return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_shouldYield.apply(
+      this,
+      arguments
+    );
+  }
+
   function unstable_runWithPriority() {
     return global.React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.Scheduler.unstable_runWithPriority.apply(
       this,
@@ -64,6 +71,7 @@
     unstable_now: unstable_now,
     unstable_scheduleCallback: unstable_scheduleCallback,
     unstable_cancelCallback: unstable_cancelCallback,
+    unstable_shouldYield: unstable_shouldYield,
     unstable_runWithPriority: unstable_runWithPriority,
     unstable_wrapCallback: unstable_wrapCallback,
     unstable_getCurrentPriorityLevel: unstable_getCurrentPriorityLevel,

packages/scheduler/src/Scheduler.js

@@ -30,6 +30,7 @@ var IDLE_PRIORITY = maxSigned31BitInt;
 // Callbacks are stored as a circular, doubly linked list.
 var firstCallbackNode = null;
 
+var currentDidTimeout = false;
 var currentPriorityLevel = NormalPriority;
 var currentEventStartTime = -1;
 var currentExpirationTime = -1;
@@ -42,41 +43,6 @@ var isHostCallbackScheduled = false;
 var hasNativePerformanceNow =
   typeof performance === 'object' && typeof performance.now === 'function';
 
-var timeRemaining;
-if (hasNativePerformanceNow) {
-  timeRemaining = function() {
-    if (
-      firstCallbackNode !== null &&
-      firstCallbackNode.expirationTime < currentExpirationTime
-    ) {
-      // A higher priority callback was scheduled. Yield so we can switch to
-      // working on that.
-      return 0;
-    }
-    // We assume that if we have a performance timer that the rAF callback
-    // gets a performance timer value. Not sure if this is always true.
-    var remaining = getFrameDeadline() - performance.now();
-    return remaining > 0 ? remaining : 0;
-  };
-} else {
-  timeRemaining = function() {
-    // Fallback to Date.now()
-    if (
-      firstCallbackNode !== null &&
-      firstCallbackNode.expirationTime < currentExpirationTime
-    ) {
-      return 0;
-    }
-    var remaining = getFrameDeadline() - Date.now();
-    return remaining > 0 ? remaining : 0;
-  };
-}
-
-var deadlineObject = {
-  timeRemaining,
-  didTimeout: false,
-};
-
 function ensureHostCallbackIsScheduled() {
   if (isExecutingCallback) {
     // Don't schedule work yet; wait until the next time we yield.
@@ -121,7 +87,7 @@ function flushFirstCallback() {
   currentExpirationTime = expirationTime;
   var continuationCallback;
   try {
-    continuationCallback = callback(deadlineObject);
+    continuationCallback = callback();
   } finally {
     currentPriorityLevel = previousPriorityLevel;
     currentExpirationTime = previousExpirationTime;
@@ -184,7 +150,6 @@ function flushImmediateWork() {
     firstCallbackNode.priorityLevel === ImmediatePriority
   ) {
     isExecutingCallback = true;
-    deadlineObject.didTimeout = true;
     try {
       do {
         flushFirstCallback();
@@ -207,7 +172,8 @@ function flushImmediateWork() {
 
 function flushWork(didTimeout) {
   isExecutingCallback = true;
-  deadlineObject.didTimeout = didTimeout;
+  const previousDidTimeout = currentDidTimeout;
+  currentDidTimeout = didTimeout;
   try {
     if (didTimeout) {
       // Flush all the expired callbacks without yielding.
@@ -232,14 +198,12 @@ function flushWork(didTimeout) {
       if (firstCallbackNode !== null) {
         do {
           flushFirstCallback();
-        } while (
-          firstCallbackNode !== null &&
-          getFrameDeadline() - getCurrentTime() > 0
-        );
+        } while (firstCallbackNode !== null && !shouldYieldToHost());
       }
     }
   } finally {
     isExecutingCallback = false;
+    currentDidTimeout = previousDidTimeout;
     if (firstCallbackNode !== null) {
       // There's still work remaining. Request another callback.
       ensureHostCallbackIsScheduled();
@@ -399,6 +363,15 @@ function unstable_getCurrentPriorityLevel() {
   return currentPriorityLevel;
 }
 
+function unstable_shouldYield() {
+  return (
+    !currentDidTimeout &&
+    ((firstCallbackNode !== null &&
+      firstCallbackNode.expirationTime < currentExpirationTime) ||
+      shouldYieldToHost())
+  );
+}
+
 // The remaining code is essentially a polyfill for requestIdleCallback. It
 // works by scheduling a requestAnimationFrame, storing the time for the start
 // of the frame, then scheduling a postMessage which gets scheduled after paint.
@@ -466,14 +439,14 @@ if (hasNativePerformanceNow) {
 
 var requestHostCallback;
 var cancelHostCallback;
-var getFrameDeadline;
+var shouldYieldToHost;
 
 if (typeof window !== 'undefined' && window._schedMock) {
   // Dynamic injection, only for testing purposes.
   var impl = window._schedMock;
   requestHostCallback = impl[0];
   cancelHostCallback = impl[1];
-  getFrameDeadline = impl[2];
+  shouldYieldToHost = impl[2];
 } else if (
   // If Scheduler runs in a non-DOM environment, it falls back to a naive
   // implementation using setTimeout.
@@ -509,8 +482,8 @@ if (typeof window !== 'undefined' && window._schedMock) {
   cancelHostCallback = function() {
     _callback = null;
   };
-  getFrameDeadline = function() {
-    return Infinity;
+  shouldYieldToHost = function() {
+    return false;
   };
   getCurrentTime = function() {
     return _currentTime === -1 ? 0 : _currentTime;
@@ -549,8 +522,8 @@ if (typeof window !== 'undefined' && window._schedMock) {
   var previousFrameTime = 33;
   var activeFrameTime = 33;
 
-  getFrameDeadline = function() {
-    return frameDeadline;
+  shouldYieldToHost = function() {
+    return frameDeadline <= getCurrentTime();
   };
 
   // We use the postMessage trick to defer idle work until after the repaint.
@@ -687,5 +660,6 @@ export {
   unstable_cancelCallback,
   unstable_wrapCallback,
   unstable_getCurrentPriorityLevel,
+  unstable_shouldYield,
   getCurrentTime as unstable_now,
 };

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

@@ -17,6 +17,7 @@ let scheduleCallback;
 let cancelCallback;
 let wrapCallback;
 let getCurrentPriorityLevel;
+let shouldYield;
 let flushWork;
 let advanceTime;
 let doWork;
@@ -119,8 +120,8 @@ describe('Scheduler', () => {
       _flushWork = null;
       clearTimeout(timeoutID);
     }
-    function getTimeRemaining() {
-      return endOfFrame;
+    function shouldYieldToHost() {
+      return endOfFrame <= currentTime;
     }
 
     // Override host implementation
@@ -131,7 +132,7 @@ describe('Scheduler', () => {
     window._schedMock = [
       requestHostCallback,
       cancelHostCallback,
-      getTimeRemaining,
+      shouldYieldToHost,
     ];
 
     const Schedule = require('scheduler');
@@ -143,6 +144,7 @@ describe('Scheduler', () => {
     cancelCallback = Schedule.unstable_cancelCallback;
     wrapCallback = Schedule.unstable_wrapCallback;
     getCurrentPriorityLevel = Schedule.unstable_getCurrentPriorityLevel;
+    shouldYield = Schedule.unstable_shouldYield;
   });
 
   it('flushes work incrementally', () => {
@@ -273,14 +275,10 @@ describe('Scheduler', () => {
     scheduleCallback(() => doWork('B', 100));
 
     const tasks = [['C1', 100], ['C2', 100], ['C3', 100]];
-    const C = deadline => {
+    const C = () => {
       while (tasks.length > 0) {
         doWork(...tasks.shift());
-        if (
-          tasks.length > 0 &&
-          !deadline.didTimeout &&
-          deadline.timeRemaining() <= 0
-        ) {
+        if (shouldYield()) {
           yieldValue('Yield!');
           return C;
         }
@@ -299,14 +297,10 @@ describe('Scheduler', () => {
 
   it('continuation callbacks inherit the expiration of the previous callback', () => {
     const tasks = [['A', 125], ['B', 124], ['C', 100], ['D', 100]];
-    const work = deadline => {
+    const work = () => {
       while (tasks.length > 0) {
         doWork(...tasks.shift());
-        if (
-          tasks.length > 0 &&
-          !deadline.didTimeout &&
-          deadline.timeRemaining() <= 0
-        ) {
+        if (shouldYield()) {
           yieldValue('Yield!');
           return work;
         }
@@ -344,14 +338,10 @@ describe('Scheduler', () => {
 
   it('continuations are interrupted by higher priority work', () => {
     const tasks = [['A', 100], ['B', 100], ['C', 100], ['D', 100]];
-    const work = deadline => {
+    const work = () => {
       while (tasks.length > 0) {
         doWork(...tasks.shift());
-        if (
-          tasks.length > 0 &&
-          !deadline.didTimeout &&
-          deadline.timeRemaining() <= 0
-        ) {
+        if (tasks.length > 0 && shouldYield()) {
           yieldValue('Yield!');
           return work;
         }
@@ -372,7 +362,7 @@ describe('Scheduler', () => {
       'inside an executing callback',
     () => {
       const tasks = [['A', 100], ['B', 100], ['C', 100], ['D', 100]];
-      const work = deadline => {
+      const work = () => {
         while (tasks.length > 0) {
           const task = tasks.shift();
           doWork(...task);
@@ -383,11 +373,7 @@ describe('Scheduler', () => {
               scheduleCallback(() => doWork('High pri', 100)),
             );
           }
-          if (
-            tasks.length > 0 &&
-            !deadline.didTimeout &&
-            deadline.timeRemaining() <= 0
-          ) {
+          if (tasks.length > 0 && shouldYield()) {
             yieldValue('Yield!');
             return work;
           }

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

@@ -105,9 +105,6 @@ describe('SchedulerDOM', () => {
       scheduleCallback(cb);
       advanceOneFrame({timeLeftInFrame: 15});
       expect(cb).toHaveBeenCalledTimes(1);
-      // should not have timed out and should include a timeRemaining method
-      expect(cb.mock.calls[0][0].didTimeout).toBe(false);
-      expect(typeof cb.mock.calls[0][0].timeRemaining()).toBe('number');
     });
 
     it('inserts its rAF callback as early into the queue as possible', () => {
@@ -144,16 +141,6 @@ describe('SchedulerDOM', () => {
         // after a delay, calls as many callbacks as it has time for
         advanceOneFrame({timeLeftInFrame: 15});
         expect(callbackLog).toEqual(['A', 'B']);
-        // callbackA should not have timed out and should include a timeRemaining method
-        expect(callbackA.mock.calls[0][0].didTimeout).toBe(false);
-        expect(typeof callbackA.mock.calls[0][0].timeRemaining()).toBe(
-          'number',
-        );
-        // callbackA should not have timed out and should include a timeRemaining method
-        expect(callbackB.mock.calls[0][0].didTimeout).toBe(false);
-        expect(typeof callbackB.mock.calls[0][0].timeRemaining()).toBe(
-          'number',
-        );
       });
 
       it("accepts callbacks betweeen animationFrame and postMessage and doesn't stall", () => {

packages/shared/forks/Scheduler.umd.js

@@ -15,6 +15,12 @@ const {
   unstable_cancelCallback,
   unstable_now,
   unstable_scheduleCallback,
+  unstable_shouldYield,
 } = ReactInternals.Scheduler;
 
-export {unstable_cancelCallback, unstable_now, unstable_scheduleCallback};
+export {
+  unstable_cancelCallback,
+  unstable_now,
+  unstable_scheduleCallback,
+  unstable_shouldYield,
+};