From 3af91eb8cee7cbfb482e4bbee91b81822e264916 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Wed, 19 Jun 2019 15:57:33 -0700
Subject: [PATCH] [Scheduler] Use continuation pattern for posting host
 callback (#15910)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* [scheduler] Internal rename: Callback -> Task

Rename Callback type to Task. Does not affect the public API, only
internal names, though eventually we'll probably want to align with the
WICG Main-thread Scheduling proposal
(https://github.com/WICG/main-thread-scheduling).

* [scheduler] flushFirstTask() -> flushTask(task)

Pass task as an argument to `flushTask` instead of using a module-
level variable.

* [scheduler] Add startTime field

This does not change any semantics, but in the future `startTime` may
represent a future time, to support delayed tasks.

* [Scheduler] Use continuation pattern for host cb

As I prepare to implement integrated timers, I noticed some
peculiarities in the Scheduler implementation that could afford to be
cleaned up.

This is a refactor and shouldn't affect any observable behavior; mostly
it removes some concepts that existed in earlier iterations of Scheduler
and are no longer needed.

The main change is to how the DOM implementation schedules an additional
callback before yielding to the main thread. It used to follow the same
code path for scheduling task; now it has its own branch directly
inside the message event handler. The special case for error handling
— where we call `postMessage` immediately without waiting for rAF —
has similarly been localized inside the catch block of the message
event handler.
---
 packages/scheduler/src/Scheduler.js           | 322 ++++++++----------
 .../scheduler/src/__tests__/Scheduler-test.js |   4 +-
 .../src/forks/SchedulerHostConfig.default.js  | 100 +++---
 .../src/forks/SchedulerHostConfig.mock.js     |  97 +++---
 4 files changed, 222 insertions(+), 301 deletions(-)

diff --git a/packages/scheduler/src/Scheduler.js b/packages/scheduler/src/Scheduler.js
index e130014927..73ed0af352 100644
--- a/packages/scheduler/src/Scheduler.js
+++ b/packages/scheduler/src/Scheduler.js
@@ -11,7 +11,6 @@
 import {enableSchedulerDebugging} from './SchedulerFeatureFlags';
 import {
   requestHostCallback,
-  cancelHostCallback,
   shouldYieldToHost,
   getCurrentTime,
   forceFrameRate,
@@ -38,130 +37,105 @@ var LOW_PRIORITY_TIMEOUT = 10000;
 // Never times out
 var IDLE_PRIORITY = maxSigned31BitInt;
 
-// Callbacks are stored as a circular, doubly linked list.
-var firstCallbackNode = null;
+// Tasks are stored as a circular, doubly linked list.
+var firstTask = null;
 
-var currentHostCallbackDidTimeout = false;
 // Pausing the scheduler is useful for debugging.
 var isSchedulerPaused = false;
 
+var currentTask = null;
 var currentPriorityLevel = NormalPriority;
-var currentEventStartTime = -1;
-var currentExpirationTime = -1;
 
 // This is set while performing work, to prevent re-entrancy.
 var isPerformingWork = false;
 
 var isHostCallbackScheduled = false;
 
-function scheduleHostCallbackIfNeeded() {
-  if (isPerformingWork) {
-    // Don't schedule work yet; wait until the next time we yield.
-    return;
-  }
-  if (firstCallbackNode !== null) {
-    // Schedule the host callback using the earliest expiration in the list.
-    var expirationTime = firstCallbackNode.expirationTime;
-    if (isHostCallbackScheduled) {
-      // Cancel the existing host callback.
-      cancelHostCallback();
-    } else {
-      isHostCallbackScheduled = true;
-    }
-    requestHostCallback(flushWork, expirationTime);
-  }
-}
-
-function flushFirstCallback() {
-  const currentlyFlushingCallback = firstCallbackNode;
-
-  // Remove the node from the list before calling the callback. That way the
+function flushTask(task, currentTime) {
+  // Remove the task from the list before calling the callback. That way the
   // list is in a consistent state even if the callback throws.
-  var next = firstCallbackNode.next;
-  if (firstCallbackNode === next) {
-    // This is the last callback in the list.
-    firstCallbackNode = null;
-    next = null;
+  const next = task.next;
+  if (next === task) {
+    // This is the only scheduled task. Clear the list.
+    firstTask = null;
   } else {
-    var lastCallbackNode = firstCallbackNode.previous;
-    firstCallbackNode = lastCallbackNode.next = next;
-    next.previous = lastCallbackNode;
+    // Remove the task from its position in the list.
+    if (task === firstTask) {
+      firstTask = next;
+    }
+    const previous = task.previous;
+    previous.next = next;
+    next.previous = previous;
   }
+  task.next = task.previous = null;
 
-  currentlyFlushingCallback.next = currentlyFlushingCallback.previous = null;
-
-  // Now it's safe to call the callback.
-  var callback = currentlyFlushingCallback.callback;
-  var expirationTime = currentlyFlushingCallback.expirationTime;
-  var priorityLevel = currentlyFlushingCallback.priorityLevel;
+  // Now it's safe to execute the task.
+  var callback = task.callback;
   var previousPriorityLevel = currentPriorityLevel;
-  var previousExpirationTime = currentExpirationTime;
-  currentPriorityLevel = priorityLevel;
-  currentExpirationTime = expirationTime;
+  var previousTask = currentTask;
+  currentPriorityLevel = task.priorityLevel;
+  currentTask = task;
   var continuationCallback;
   try {
-    const didUserCallbackTimeout =
-      currentHostCallbackDidTimeout ||
-      // Immediate priority callbacks are always called as if they timed out
-      priorityLevel === ImmediatePriority;
+    var didUserCallbackTimeout = task.expirationTime <= currentTime;
     continuationCallback = callback(didUserCallbackTimeout);
   } catch (error) {
     throw error;
   } finally {
     currentPriorityLevel = previousPriorityLevel;
-    currentExpirationTime = previousExpirationTime;
+    currentTask = previousTask;
   }
 
   // A callback may return a continuation. The continuation should be scheduled
   // with the same priority and expiration as the just-finished callback.
   if (typeof continuationCallback === 'function') {
-    var continuationNode: CallbackNode = {
+    var expirationTime = task.expirationTime;
+    var continuationTask = {
       callback: continuationCallback,
-      priorityLevel,
+      priorityLevel: task.priorityLevel,
+      startTime: task.startTime,
       expirationTime,
       next: null,
       previous: null,
     };
 
-    // Insert the new callback into the list, sorted by its expiration. This is
+    // Insert the new callback into the list, sorted by its timeout. This is
     // almost the same as the code in `scheduleCallback`, except the callback
-    // is inserted into the list *before* callbacks of equal expiration instead
+    // is inserted into the list *before* callbacks of equal timeout instead
     // of after.
-    if (firstCallbackNode === null) {
+    if (firstTask === null) {
       // This is the first callback in the list.
-      firstCallbackNode = continuationNode.next = continuationNode.previous = continuationNode;
+      firstTask = continuationTask.next = continuationTask.previous = continuationTask;
     } else {
       var nextAfterContinuation = null;
-      var node = firstCallbackNode;
+      var t = firstTask;
       do {
-        if (node.expirationTime >= expirationTime) {
-          // This callback expires at or after the continuation. We will insert
-          // the continuation *before* this callback.
-          nextAfterContinuation = node;
+        if (expirationTime <= t.expirationTime) {
+          // This task times out at or after the continuation. We will insert
+          // the continuation *before* this task.
+          nextAfterContinuation = t;
           break;
         }
-        node = node.next;
-      } while (node !== firstCallbackNode);
-
+        t = t.next;
+      } while (t !== firstTask);
       if (nextAfterContinuation === null) {
-        // No equal or lower priority callback was found, which means the new
-        // callback is the lowest priority callback in the list.
-        nextAfterContinuation = firstCallbackNode;
-      } else if (nextAfterContinuation === firstCallbackNode) {
-        // The new callback is the highest priority callback in the list.
-        firstCallbackNode = continuationNode;
-        scheduleHostCallbackIfNeeded();
+        // No equal or lower priority task was found, which means the new task
+        // is the lowest priority task in the list.
+        nextAfterContinuation = firstTask;
+      } else if (nextAfterContinuation === firstTask) {
+        // The new task is the highest priority task in the list.
+        firstTask = continuationTask;
       }
 
-      var previous = nextAfterContinuation.previous;
-      previous.next = nextAfterContinuation.previous = continuationNode;
-      continuationNode.next = nextAfterContinuation;
-      continuationNode.previous = previous;
+      const previous = nextAfterContinuation.previous;
+      previous.next = nextAfterContinuation.previous = continuationTask;
+      continuationTask.next = nextAfterContinuation;
+      continuationTask.previous = previous;
     }
   }
 }
 
-function flushWork(didUserCallbackTimeout) {
+function flushWork(hasTimeRemaining, initialTime) {
   // Exit right away if we're currently paused
   if (enableSchedulerDebugging && isSchedulerPaused) {
     return;
@@ -171,48 +145,38 @@ function flushWork(didUserCallbackTimeout) {
   isHostCallbackScheduled = false;
 
   isPerformingWork = true;
-  const previousDidTimeout = currentHostCallbackDidTimeout;
-  currentHostCallbackDidTimeout = didUserCallbackTimeout;
   try {
-    if (didUserCallbackTimeout) {
+    if (!hasTimeRemaining) {
       // Flush all the expired callbacks without yielding.
+      // TODO: Split flushWork into two separate functions instead of using
+      // a boolean argument?
+      let currentTime = initialTime;
       while (
-        firstCallbackNode !== null &&
+        firstTask !== null &&
+        firstTask.expirationTime <= currentTime &&
         !(enableSchedulerDebugging && isSchedulerPaused)
       ) {
-        // TODO Wrap in feature flag
-        // Read the current time. Flush all the callbacks that expire at or
-        // earlier than that time. Then read the current time again and repeat.
-        // This optimizes for as few performance.now calls as possible.
-        var currentTime = getCurrentTime();
-        if (firstCallbackNode.expirationTime <= currentTime) {
-          do {
-            flushFirstCallback();
-          } while (
-            firstCallbackNode !== null &&
-            firstCallbackNode.expirationTime <= currentTime &&
-            !(enableSchedulerDebugging && isSchedulerPaused)
-          );
-          continue;
-        }
-        break;
+        flushTask(firstTask, currentTime);
+        currentTime = getCurrentTime();
       }
     } else {
       // Keep flushing callbacks until we run out of time in the frame.
-      if (firstCallbackNode !== null) {
+      let currentTime = initialTime;
+      if (firstTask !== null) {
         do {
-          if (enableSchedulerDebugging && isSchedulerPaused) {
-            break;
-          }
-          flushFirstCallback();
-        } while (firstCallbackNode !== null && !shouldYieldToHost());
+          flushTask(firstTask, currentTime);
+          currentTime = getCurrentTime();
+        } while (
+          firstTask !== null &&
+          !shouldYieldToHost() &&
+          !(enableSchedulerDebugging && isSchedulerPaused)
+        );
       }
     }
+    // Return whether there's additional work
+    return firstTask !== null;
   } finally {
     isPerformingWork = false;
-    currentHostCallbackDidTimeout = previousDidTimeout;
-    // There's still work remaining. Request another callback.
-    scheduleHostCallbackIfNeeded();
   }
 }
 
@@ -229,24 +193,17 @@ function unstable_runWithPriority(priorityLevel, eventHandler) {
   }
 
   var previousPriorityLevel = currentPriorityLevel;
-  var previousEventStartTime = currentEventStartTime;
   currentPriorityLevel = priorityLevel;
-  currentEventStartTime = getCurrentTime();
 
   try {
     return eventHandler();
-  } catch (error) {
-    // There's still work remaining. Request another callback.
-    scheduleHostCallbackIfNeeded();
-    throw error;
   } finally {
     currentPriorityLevel = previousPriorityLevel;
-    currentEventStartTime = previousEventStartTime;
   }
 }
 
 function unstable_next(eventHandler) {
-  let priorityLevel;
+  var priorityLevel;
   switch (currentPriorityLevel) {
     case ImmediatePriority:
     case UserBlockingPriority:
@@ -261,19 +218,12 @@ function unstable_next(eventHandler) {
   }
 
   var previousPriorityLevel = currentPriorityLevel;
-  var previousEventStartTime = currentEventStartTime;
   currentPriorityLevel = priorityLevel;
-  currentEventStartTime = getCurrentTime();
 
   try {
     return eventHandler();
-  } catch (error) {
-    // There's still work remaining. Request another callback.
-    scheduleHostCallbackIfNeeded();
-    throw error;
   } finally {
     currentPriorityLevel = previousPriorityLevel;
-    currentEventStartTime = previousEventStartTime;
   }
 }
 
@@ -282,103 +232,98 @@ function unstable_wrapCallback(callback) {
   return function() {
     // This is a fork of runWithPriority, inlined for performance.
     var previousPriorityLevel = currentPriorityLevel;
-    var previousEventStartTime = currentEventStartTime;
     currentPriorityLevel = parentPriorityLevel;
-    currentEventStartTime = getCurrentTime();
 
     try {
       return callback.apply(this, arguments);
-    } catch (error) {
-      // There's still work remaining. Request another callback.
-      scheduleHostCallbackIfNeeded();
-      throw error;
     } finally {
       currentPriorityLevel = previousPriorityLevel;
-      currentEventStartTime = previousEventStartTime;
     }
   };
 }
 
-function unstable_scheduleCallback(
-  priorityLevel,
-  callback,
-  deprecated_options,
-) {
-  var startTime =
-    currentEventStartTime !== -1 ? currentEventStartTime : getCurrentTime();
+function unstable_scheduleCallback(priorityLevel, callback, options) {
+  var startTime = getCurrentTime();
 
-  var expirationTime;
+  var timeout;
   if (
-    typeof deprecated_options === 'object' &&
-    deprecated_options !== null &&
-    typeof deprecated_options.timeout === 'number'
+    typeof options === 'object' &&
+    options !== null &&
+    typeof options.timeout === 'number'
   ) {
-    // FIXME: Remove this branch once we lift expiration times out of React.
-    expirationTime = startTime + deprecated_options.timeout;
+    timeout = options.timeout;
   } else {
     switch (priorityLevel) {
       case ImmediatePriority:
-        expirationTime = startTime + IMMEDIATE_PRIORITY_TIMEOUT;
+        timeout = IMMEDIATE_PRIORITY_TIMEOUT;
         break;
       case UserBlockingPriority:
-        expirationTime = startTime + USER_BLOCKING_PRIORITY;
+        timeout = USER_BLOCKING_PRIORITY;
         break;
       case IdlePriority:
-        expirationTime = startTime + IDLE_PRIORITY;
+        timeout = IDLE_PRIORITY;
         break;
       case LowPriority:
-        expirationTime = startTime + LOW_PRIORITY_TIMEOUT;
+        timeout = LOW_PRIORITY_TIMEOUT;
         break;
       case NormalPriority:
       default:
-        expirationTime = startTime + NORMAL_PRIORITY_TIMEOUT;
+        timeout = NORMAL_PRIORITY_TIMEOUT;
     }
   }
 
-  var newNode = {
+  var expirationTime = startTime + timeout;
+
+  var newTask = {
     callback,
-    priorityLevel: priorityLevel,
+    priorityLevel,
+    startTime,
     expirationTime,
     next: null,
     previous: null,
   };
 
-  // Insert the new callback into the list, ordered first by expiration, then
-  // by insertion. So the new callback is inserted after any other callback
-  // with equal expiration.
-  if (firstCallbackNode === null) {
-    // This is the first callback in the list.
-    firstCallbackNode = newNode.next = newNode.previous = newNode;
-    scheduleHostCallbackIfNeeded();
+  // Insert the new task into the list, ordered first by its timeout, then by
+  // insertion. So the new task is inserted after any other task the
+  // same timeout
+  if (firstTask === null) {
+    // This is the first task in the list.
+    firstTask = newTask.next = newTask.previous = newTask;
   } else {
     var next = null;
-    var node = firstCallbackNode;
+    var task = firstTask;
     do {
-      if (node.expirationTime > expirationTime) {
-        // The new callback expires before this one.
-        next = node;
+      if (expirationTime < task.expirationTime) {
+        // The new task times out before this one.
+        next = task;
         break;
       }
-      node = node.next;
-    } while (node !== firstCallbackNode);
+      task = task.next;
+    } while (task !== firstTask);
 
     if (next === null) {
-      // No callback with a later expiration was found, which means the new
-      // callback has the latest expiration in the list.
-      next = firstCallbackNode;
-    } else if (next === firstCallbackNode) {
-      // The new callback has the earliest expiration in the entire list.
-      firstCallbackNode = newNode;
-      scheduleHostCallbackIfNeeded();
+      // No task with a later timeout was found, which means the new task has
+      // the latest timeout in the list.
+      next = firstTask;
+    } else if (next === firstTask) {
+      // The new task has the earliest expiration in the entire list.
+      firstTask = newTask;
     }
 
     var previous = next.previous;
-    previous.next = next.previous = newNode;
-    newNode.next = next;
-    newNode.previous = previous;
+    previous.next = next.previous = newTask;
+    newTask.next = next;
+    newTask.previous = previous;
   }
 
-  return newNode;
+  // Schedule a host callback, if needed. If we're already performing work, wait
+  // until the next time we yield.
+  if (!isHostCallbackScheduled && !isPerformingWork) {
+    isHostCallbackScheduled = true;
+    requestHostCallback(flushWork);
+  }
+
+  return newTask;
 }
 
 function unstable_pauseExecution() {
@@ -387,36 +332,37 @@ function unstable_pauseExecution() {
 
 function unstable_continueExecution() {
   isSchedulerPaused = false;
-  if (firstCallbackNode !== null) {
-    scheduleHostCallbackIfNeeded();
+  if (!isHostCallbackScheduled && !isPerformingWork) {
+    isHostCallbackScheduled = true;
+    requestHostCallback(flushWork);
   }
 }
 
 function unstable_getFirstCallbackNode() {
-  return firstCallbackNode;
+  return firstTask;
 }
 
-function unstable_cancelCallback(callbackNode) {
-  var next = callbackNode.next;
+function unstable_cancelCallback(task) {
+  var next = task.next;
   if (next === null) {
     // Already cancelled.
     return;
   }
 
-  if (next === callbackNode) {
-    // This is the only scheduled callback. Clear the list.
-    firstCallbackNode = null;
+  if (next === task) {
+    // This is the only scheduled task. Clear the list.
+    firstTask = null;
   } else {
-    // Remove the callback from its position in the list.
-    if (callbackNode === firstCallbackNode) {
-      firstCallbackNode = next;
+    // Remove the task from its position in the list.
+    if (task === firstTask) {
+      firstTask = next;
     }
-    var previous = callbackNode.previous;
+    var previous = task.previous;
     previous.next = next;
     next.previous = previous;
   }
 
-  callbackNode.next = callbackNode.previous = null;
+  task.next = task.previous = null;
 }
 
 function unstable_getCurrentPriorityLevel() {
@@ -425,10 +371,10 @@ function unstable_getCurrentPriorityLevel() {
 
 function unstable_shouldYield() {
   return (
-    !currentHostCallbackDidTimeout &&
-    ((firstCallbackNode !== null &&
-      firstCallbackNode.expirationTime < currentExpirationTime) ||
-      shouldYieldToHost())
+    (currentTask !== null &&
+      firstTask !== null &&
+      firstTask.expirationTime < currentTask.expirationTime) ||
+    shouldYieldToHost()
   );
 }
 
diff --git a/packages/scheduler/src/__tests__/Scheduler-test.js b/packages/scheduler/src/__tests__/Scheduler-test.js
index 20cd67fa27..fee2e97c45 100644
--- a/packages/scheduler/src/__tests__/Scheduler-test.js
+++ b/packages/scheduler/src/__tests__/Scheduler-test.js
@@ -227,7 +227,7 @@ describe('Scheduler', () => {
   });
 
   it(
-    'continutations are interrupted by higher priority work scheduled ' +
+    'continuations are interrupted by higher priority work scheduled ' +
       'inside an executing callback',
     () => {
       const tasks = [['A', 100], ['B', 100], ['C', 100], ['D', 100]];
@@ -237,7 +237,7 @@ describe('Scheduler', () => {
           const [label, ms] = task;
           Scheduler.advanceTime(ms);
           Scheduler.yieldValue(label);
-          if (task[0] === 'B') {
+          if (label === 'B') {
             // Schedule high pri work from inside another callback
             Scheduler.yieldValue('Schedule high pri');
             scheduleCallback(UserBlockingPriority, () => {
diff --git a/packages/scheduler/src/forks/SchedulerHostConfig.default.js b/packages/scheduler/src/forks/SchedulerHostConfig.default.js
index ec084f1dd1..1078425caa 100644
--- a/packages/scheduler/src/forks/SchedulerHostConfig.default.js
+++ b/packages/scheduler/src/forks/SchedulerHostConfig.default.js
@@ -88,22 +88,26 @@ if (
   // If this accidentally gets imported in a non-browser environment, e.g. JavaScriptCore,
   // fallback to a naive implementation.
   let _callback = null;
-  const _flushCallback = function(didTimeout) {
+  const _flushCallback = function() {
     if (_callback !== null) {
       try {
-        _callback(didTimeout);
-      } finally {
+        const currentTime = getCurrentTime();
+        const hasRemainingTime = true;
+        _callback(hasRemainingTime, currentTime);
         _callback = null;
+      } catch (e) {
+        setTimeout(_flushCallback, 0);
+        throw e;
       }
     }
   };
-  requestHostCallback = function(cb, ms) {
+  requestHostCallback = function(cb) {
     if (_callback !== null) {
       // Protect against re-entrancy.
       setTimeout(requestHostCallback, 0, cb);
     } else {
       _callback = cb;
-      setTimeout(_flushCallback, 0, false);
+      setTimeout(_flushCallback, 0);
     }
   };
   cancelHostCallback = function() {
@@ -134,12 +138,9 @@ if (
 
   let scheduledHostCallback = null;
   let isMessageEventScheduled = false;
-  let timeoutTime = -1;
 
   let isAnimationFrameScheduled = false;
 
-  let isFlushingHostCallback = false;
-
   let frameDeadline = 0;
   // We start out assuming that we run at 30fps but then the heuristic tracking
   // will adjust this value to a faster fps if we get more frequent animation
@@ -175,42 +176,30 @@ if (
   const port = channel.port2;
   channel.port1.onmessage = function(event) {
     isMessageEventScheduled = false;
-
-    const prevScheduledCallback = scheduledHostCallback;
-    const prevTimeoutTime = timeoutTime;
-    scheduledHostCallback = null;
-    timeoutTime = -1;
-
-    const currentTime = getCurrentTime();
-
-    let didTimeout = false;
-    if (frameDeadline - currentTime <= 0) {
-      // There's no time left in this idle period. Check if the callback has
-      // a timeout and whether it's been exceeded.
-      if (prevTimeoutTime !== -1 && prevTimeoutTime <= currentTime) {
-        // Exceeded the timeout. Invoke the callback even though there's no
-        // time left.
-        didTimeout = true;
-      } else {
-        // No timeout.
-        if (!isAnimationFrameScheduled) {
-          // Schedule another animation callback so we retry later.
-          isAnimationFrameScheduled = true;
-          requestAnimationFrameWithTimeout(animationTick);
-        }
-        // Exit without invoking the callback.
-        scheduledHostCallback = prevScheduledCallback;
-        timeoutTime = prevTimeoutTime;
-        return;
-      }
-    }
-
-    if (prevScheduledCallback !== null) {
-      isFlushingHostCallback = true;
+    if (scheduledHostCallback !== null) {
+      const currentTime = getCurrentTime();
+      const hasTimeRemaining = frameDeadline - currentTime > 0;
       try {
-        prevScheduledCallback(didTimeout);
-      } finally {
-        isFlushingHostCallback = false;
+        const hasMoreWork = scheduledHostCallback(
+          hasTimeRemaining,
+          currentTime,
+        );
+        if (hasMoreWork) {
+          // Ensure the next frame is scheduled.
+          if (!isAnimationFrameScheduled) {
+            isAnimationFrameScheduled = true;
+            requestAnimationFrameWithTimeout(animationTick);
+          }
+        } else {
+          scheduledHostCallback = null;
+        }
+      } catch (error) {
+        // If a scheduler task throws, exit the current browser task so the
+        // error can be observed, and post a new task as soon as possible
+        // so we can continue where we left off.
+        isMessageEventScheduled = true;
+        port.postMessage(undefined);
+        throw error;
       }
     }
   };
@@ -262,25 +251,22 @@ if (
     }
   };
 
-  requestHostCallback = function(callback, absoluteTimeout) {
-    scheduledHostCallback = callback;
-    timeoutTime = absoluteTimeout;
-    if (isFlushingHostCallback || absoluteTimeout < 0) {
-      // Don't wait for the next frame. Continue working ASAP, in a new event.
-      port.postMessage(undefined);
-    } else if (!isAnimationFrameScheduled) {
-      // If rAF didn't already schedule one, we need to schedule a frame.
-      // TODO: If this rAF doesn't materialize because the browser throttles, we
-      // might want to still have setTimeout trigger rIC as a backup to ensure
-      // that we keep performing work.
-      isAnimationFrameScheduled = true;
-      requestAnimationFrameWithTimeout(animationTick);
+  requestHostCallback = function(callback) {
+    if (scheduledHostCallback === null) {
+      scheduledHostCallback = callback;
+      if (!isAnimationFrameScheduled) {
+        // If rAF didn't already schedule one, we need to schedule a frame.
+        // TODO: If this rAF doesn't materialize because the browser throttles,
+        // we might want to still have setTimeout trigger rIC as a backup to
+        // ensure that we keep performing work.
+        isAnimationFrameScheduled = true;
+        requestAnimationFrameWithTimeout(animationTick);
+      }
     }
   };
 
   cancelHostCallback = function() {
     scheduledHostCallback = null;
     isMessageEventScheduled = false;
-    timeoutTime = -1;
   };
 }
diff --git a/packages/scheduler/src/forks/SchedulerHostConfig.mock.js b/packages/scheduler/src/forks/SchedulerHostConfig.mock.js
index 4ecd682ff8..51ad74b967 100644
--- a/packages/scheduler/src/forks/SchedulerHostConfig.mock.js
+++ b/packages/scheduler/src/forks/SchedulerHostConfig.mock.js
@@ -8,33 +8,25 @@
  */
 
 let currentTime: number = 0;
-let scheduledCallback: (boolean => void) | null = null;
-let scheduledCallbackExpiration: number = -1;
+let scheduledCallback: ((boolean, number) => void) | null = null;
 let yieldedValues: Array<mixed> | null = null;
 let expectedNumberOfYields: number = -1;
 let didStop: boolean = false;
 let isFlushing: boolean = false;
 
-export function requestHostCallback(
-  callback: boolean => void,
-  expiration: number,
-) {
+export function requestHostCallback(callback: boolean => void) {
   scheduledCallback = callback;
-  scheduledCallbackExpiration = expiration;
 }
 
 export function cancelHostCallback(): void {
   scheduledCallback = null;
-  scheduledCallbackExpiration = -1;
 }
 
 export function shouldYieldToHost(): boolean {
   if (
-    (expectedNumberOfYields !== -1 &&
-      yieldedValues !== null &&
-      yieldedValues.length >= expectedNumberOfYields) ||
-    (scheduledCallbackExpiration !== -1 &&
-      scheduledCallbackExpiration <= currentTime)
+    expectedNumberOfYields !== -1 &&
+    yieldedValues !== null &&
+    yieldedValues.length >= expectedNumberOfYields
   ) {
     // We yielded at least as many values as expected. Stop flushing.
     didStop = true;
@@ -57,7 +49,6 @@ export function reset() {
   }
   currentTime = 0;
   scheduledCallback = null;
-  scheduledCallbackExpiration = -1;
   yieldedValues = null;
   expectedNumberOfYields = -1;
   didStop = false;
@@ -69,21 +60,23 @@ export function unstable_flushNumberOfYields(count: number): void {
   if (isFlushing) {
     throw new Error('Already flushing work.');
   }
-  expectedNumberOfYields = count;
-  isFlushing = true;
-  try {
-    while (scheduledCallback !== null && !didStop) {
-      const cb = scheduledCallback;
-      scheduledCallback = null;
-      const didTimeout =
-        scheduledCallbackExpiration !== -1 &&
-        scheduledCallbackExpiration <= currentTime;
-      cb(didTimeout);
+  if (scheduledCallback !== null) {
+    const cb = scheduledCallback;
+    expectedNumberOfYields = count;
+    isFlushing = true;
+    try {
+      let hasMoreWork = true;
+      do {
+        hasMoreWork = cb(true, currentTime);
+      } while (hasMoreWork && !didStop);
+      if (!hasMoreWork) {
+        scheduledCallback = null;
+      }
+    } finally {
+      expectedNumberOfYields = -1;
+      didStop = false;
+      isFlushing = false;
     }
-  } finally {
-    expectedNumberOfYields = -1;
-    didStop = false;
-    isFlushing = false;
   }
 }
 
@@ -92,11 +85,12 @@ export function unstable_flushExpired() {
     throw new Error('Already flushing work.');
   }
   if (scheduledCallback !== null) {
-    const cb = scheduledCallback;
-    scheduledCallback = null;
     isFlushing = true;
     try {
-      cb(true);
+      const hasMoreWork = scheduledCallback(false, currentTime);
+      if (!hasMoreWork) {
+        scheduledCallback = null;
+      }
     } finally {
       isFlushing = false;
     }
@@ -104,27 +98,27 @@ export function unstable_flushExpired() {
 }
 
 export function unstable_flushWithoutYielding(): boolean {
+  // Returns false if no work was flushed.
   if (isFlushing) {
     throw new Error('Already flushing work.');
   }
-  isFlushing = true;
-  try {
-    if (scheduledCallback === null) {
-      return false;
+  if (scheduledCallback !== null) {
+    const cb = scheduledCallback;
+    isFlushing = true;
+    try {
+      let hasMoreWork = true;
+      do {
+        hasMoreWork = cb(true, currentTime);
+      } while (hasMoreWork);
+      if (!hasMoreWork) {
+        scheduledCallback = null;
+      }
+      return true;
+    } finally {
+      isFlushing = false;
     }
-    while (scheduledCallback !== null) {
-      const cb = scheduledCallback;
-      scheduledCallback = null;
-      const didTimeout =
-        scheduledCallbackExpiration !== -1 &&
-        scheduledCallbackExpiration <= currentTime;
-      cb(didTimeout);
-    }
-    return true;
-  } finally {
-    expectedNumberOfYields = -1;
-    didStop = false;
-    isFlushing = false;
+  } else {
+    return false;
   }
 }
 
@@ -164,12 +158,7 @@ export function yieldValue(value: mixed): void {
 
 export function advanceTime(ms: number) {
   currentTime += ms;
-  // If the host callback timed out, flush the expired work.
-  if (
-    !isFlushing &&
-    scheduledCallbackExpiration !== -1 &&
-    scheduledCallbackExpiration <= currentTime
-  ) {
+  if (!isFlushing) {
     unstable_flushExpired();
   }
 }