Use setImmediate when available over MessageChannel (#20834)

* Move direct port access into a function

* Fork based on presence of setImmediate

* Copy SchedulerDOM-test into another file

* Change the new test to use shimmed setImmediate

* Clarify comment

* Fix test to work with existing feature detection

* Add flags

* Disable OSS flag and skip tests

* Use VARIANT to reenable tests

* lol

packages/scheduler/src/SchedulerFeatureFlags.js

@@ -9,3 +9,6 @@
 export const enableSchedulerDebugging = false;
 export const enableIsInputPending = false;
 export const enableProfiling = __PROFILE__;
+
+// TODO: enable to fix https://github.com/facebook/react/issues/20756.
+export const enableSetImmediate = __VARIANT__;

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

@@ -0,0 +1,281 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+
+/* eslint-disable no-for-of-loops/no-for-of-loops */
+
+'use strict';
+
+let Scheduler;
+let runtime;
+let performance;
+let cancelCallback;
+let scheduleCallback;
+let NormalPriority;
+
+// The Scheduler implementation uses browser APIs like `MessageChannel` and
+// `setTimeout` to schedule work on the main thread. Most of our tests treat
+// these as implementation details; however, the sequence and timing of these
+// APIs are not precisely specified, and can vary across browsers.
+//
+// To prevent regressions, we need the ability to simulate specific edge cases
+// that we may encounter in various browsers.
+//
+// This test suite mocks all browser methods used in our implementation. It
+// assumes as little as possible about the order and timing of events.
+describe('SchedulerDOMSetImmediate', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    // Un-mock scheduler
+    jest.mock('scheduler', () => require.requireActual('scheduler'));
+
+    runtime = installMockBrowserRuntime();
+    performance = global.performance;
+    Scheduler = require('scheduler');
+    cancelCallback = Scheduler.unstable_cancelCallback;
+    scheduleCallback = Scheduler.unstable_scheduleCallback;
+    NormalPriority = Scheduler.unstable_NormalPriority;
+  });
+
+  afterEach(() => {
+    delete global.performance;
+
+    if (!runtime.isLogEmpty()) {
+      throw Error('Test exited without clearing log.');
+    }
+  });
+
+  function installMockBrowserRuntime() {
+    let timerIDCounter = 0;
+    // let timerIDs = new Map();
+
+    let eventLog = [];
+
+    let currentTime = 0;
+
+    global.performance = {
+      now() {
+        return currentTime;
+      },
+    };
+
+    const window = {};
+    global.window = window;
+
+    // TODO: Scheduler no longer requires these methods to be polyfilled. But
+    // maybe we want to continue warning if they don't exist, to preserve the
+    // option to rely on it in the future?
+    window.requestAnimationFrame = window.cancelAnimationFrame = () => {};
+
+    window.setTimeout = (cb, delay) => {
+      const id = timerIDCounter++;
+      log(`Set Timer`);
+      // TODO
+      return id;
+    };
+    window.clearTimeout = id => {
+      // TODO
+    };
+
+    // Unused: we expect setImmediate to be preferred.
+    global.MessageChannel = function() {
+      return {
+        port1: {},
+        port2: {
+          postMessage() {
+            throw Error('Should be unused');
+          },
+        },
+      };
+    };
+
+    let pendingSetImmediateCallback = null;
+    window.setImmediate = function(cb) {
+      if (pendingSetImmediateCallback) {
+        throw Error('Message event already scheduled');
+      }
+      log('Set Immediate');
+      pendingSetImmediateCallback = cb;
+    };
+
+    function ensureLogIsEmpty() {
+      if (eventLog.length !== 0) {
+        throw Error('Log is not empty. Call assertLog before continuing.');
+      }
+    }
+    function advanceTime(ms) {
+      currentTime += ms;
+    }
+    function fireSetImmediate() {
+      ensureLogIsEmpty();
+      if (!pendingSetImmediateCallback) {
+        throw Error('No setImmediate was scheduled');
+      }
+      const cb = pendingSetImmediateCallback;
+      pendingSetImmediateCallback = null;
+      log('setImmediate Callback');
+      cb();
+    }
+    function log(val) {
+      eventLog.push(val);
+    }
+    function isLogEmpty() {
+      return eventLog.length === 0;
+    }
+    function assertLog(expected) {
+      const actual = eventLog;
+      eventLog = [];
+      expect(actual).toEqual(expected);
+    }
+    return {
+      advanceTime,
+      fireSetImmediate,
+      log,
+      isLogEmpty,
+      assertLog,
+    };
+  }
+
+  // @gate enableSchedulerSetImmediate
+  it('task that finishes before deadline', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('Task');
+    });
+    runtime.assertLog(['Set Immediate']);
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'Task']);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('task with continuation', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('Task');
+      while (!Scheduler.unstable_shouldYield()) {
+        runtime.advanceTime(1);
+      }
+      runtime.log(`Yield at ${performance.now()}ms`);
+      return () => {
+        runtime.log('Continuation');
+      };
+    });
+    runtime.assertLog(['Set Immediate']);
+
+    runtime.fireSetImmediate();
+    runtime.assertLog([
+      'setImmediate Callback',
+      'Task',
+      'Yield at 5ms',
+      'Set Immediate',
+    ]);
+
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'Continuation']);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('multiple tasks', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('A');
+    });
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('B');
+    });
+    runtime.assertLog(['Set Immediate']);
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'A', 'B']);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('multiple tasks with a yield in between', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('A');
+      runtime.advanceTime(4999);
+    });
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('B');
+    });
+    runtime.assertLog(['Set Immediate']);
+    runtime.fireSetImmediate();
+    runtime.assertLog([
+      'setImmediate Callback',
+      'A',
+      // Ran out of time. Post a continuation event.
+      'Set Immediate',
+    ]);
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'B']);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('cancels tasks', () => {
+    const task = scheduleCallback(NormalPriority, () => {
+      runtime.log('Task');
+    });
+    runtime.assertLog(['Set Immediate']);
+    cancelCallback(task);
+    runtime.assertLog([]);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('throws when a task errors then continues in a new event', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('Oops!');
+      throw Error('Oops!');
+    });
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('Yay');
+    });
+    runtime.assertLog(['Set Immediate']);
+
+    expect(() => runtime.fireSetImmediate()).toThrow('Oops!');
+    runtime.assertLog(['setImmediate Callback', 'Oops!', 'Set Immediate']);
+
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'Yay']);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('schedule new task after queue has emptied', () => {
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('A');
+    });
+
+    runtime.assertLog(['Set Immediate']);
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'A']);
+
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('B');
+    });
+    runtime.assertLog(['Set Immediate']);
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'B']);
+  });
+
+  // @gate enableSchedulerSetImmediate
+  it('schedule new task after a cancellation', () => {
+    const handle = scheduleCallback(NormalPriority, () => {
+      runtime.log('A');
+    });
+
+    runtime.assertLog(['Set Immediate']);
+    cancelCallback(handle);
+
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback']);
+
+    scheduleCallback(NormalPriority, () => {
+      runtime.log('B');
+    });
+    runtime.assertLog(['Set Immediate']);
+    runtime.fireSetImmediate();
+    runtime.assertLog(['setImmediate Callback', 'B']);
+  });
+});

packages/scheduler/src/forks/SchedulerDOM.js

@@ -11,6 +11,7 @@
 import {
   enableSchedulerDebugging,
   enableProfiling,
+  enableSetImmediate,
 } from '../SchedulerFeatureFlags';
 
 import {push, pop, peek} from '../SchedulerMinHeap';
@@ -88,6 +89,7 @@ var isHostTimeoutScheduled = false;
 // Capture local references to native APIs, in case a polyfill overrides them.
 const setTimeout = window.setTimeout;
 const clearTimeout = window.clearTimeout;
+const setImmediate = window.setImmediate; // IE and Node.js + jsdom
 
 if (typeof console !== 'undefined') {
   // TODO: Scheduler no longer requires these methods to be polyfilled. But
@@ -533,7 +535,7 @@ const performWorkUntilDeadline = () => {
       if (hasMoreWork) {
         // If there's more work, schedule the next message event at the end
         // of the preceding one.
-        port.postMessage(null);
+        schedulePerformWorkUntilDeadline();
       } else {
         isMessageLoopRunning = false;
         scheduledHostCallback = null;
@@ -547,15 +549,36 @@ const performWorkUntilDeadline = () => {
   needsPaint = false;
 };
 
-const channel = new MessageChannel();
-const port = channel.port2;
-channel.port1.onmessage = performWorkUntilDeadline;
+let schedulePerformWorkUntilDeadline;
+if (enableSetImmediate && typeof setImmediate === 'function') {
+  // Node.js and old IE.
+  // There's a few reasons for why we prefer setImmediate.
+  //
+  // Unlike MessageChannel, it doesn't prevent a Node.js process from exiting.
+  // (Even though this is a DOM fork of the Scheduler, you could get here
+  // with a mix of Node.js 15+, which has a MessageChannel, and jsdom.)
+  // https://github.com/facebook/react/issues/20756
+  //
+  // But also, it runs earlier which is the semantic we want.
+  // If other browsers ever implement it, it's better to use it.
+  // Although both of these would be inferior to native scheduling.
+  schedulePerformWorkUntilDeadline = () => {
+    setImmediate(performWorkUntilDeadline);
+  };
+} else {
+  const channel = new MessageChannel();
+  const port = channel.port2;
+  channel.port1.onmessage = performWorkUntilDeadline;
+  schedulePerformWorkUntilDeadline = () => {
+    port.postMessage(null);
+  };
+}
 
 function requestHostCallback(callback) {
   scheduledHostCallback = callback;
   if (!isMessageLoopRunning) {
     isMessageLoopRunning = true;
-    port.postMessage(null);
+    schedulePerformWorkUntilDeadline();
   }
 }
 

packages/scheduler/src/forks/SchedulerFeatureFlags.www.js

@@ -10,6 +10,7 @@ export const {
   enableIsInputPending,
   enableSchedulerDebugging,
   enableProfiling: enableProfilingFeatureFlag,
+  enableSetImmediate,
 } = require('SchedulerFeatureFlags');
 
 export const enableProfiling = __PROFILE__ && enableProfilingFeatureFlag;

scripts/jest/TestFlags.js

@@ -55,6 +55,7 @@ function getTestFlags() {
   // These are required on demand because some of our tests mutate them. We try
   // not to but there are exceptions.
   const featureFlags = require('shared/ReactFeatureFlags');
+  const schedulerFeatureFlags = require('scheduler/src/SchedulerFeatureFlags');
 
   // TODO: This is a heuristic to detect the release channel by checking a flag
   // that is known to only be enabled in www. What we should do instead is set
@@ -89,6 +90,10 @@ function getTestFlags() {
       // tests, Jest doesn't expose the API correctly. Fix then remove
       // this override.
       enableCache: __EXPERIMENTAL__,
+
+      // This is from SchedulerFeatureFlags. Needed because there's no equivalent
+      // of ReactFeatureFlags-www.dynamic for it. Remove when enableSetImmediate is gone.
+      enableSchedulerSetImmediate: schedulerFeatureFlags.enableSetImmediate,
     },
     {
       get(flags, flagName) {