diff --git a/packages/scheduler/src/__tests__/SchedulerPostTask-test.js b/packages/scheduler/src/__tests__/SchedulerPostTask-test.js index 03952c50b9159..d76ca570c2f31 100644 --- a/packages/scheduler/src/__tests__/SchedulerPostTask-test.js +++ b/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 ? '' : 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']); + }); + }); }); diff --git a/packages/scheduler/src/forks/SchedulerPostTask.js b/packages/scheduler/src/forks/SchedulerPostTask.js index 1e37882394651..5af147117ef56 100644 --- a/packages/scheduler/src/forks/SchedulerPostTask.js +++ b/packages/scheduler/src/forks/SchedulerPostTask.js @@ -138,18 +138,25 @@ function runTask( // 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