Fix message loop behavior when host callback is cancelled (#16407)

* Add a regression test for cancelCallback with message loop

* If there's nothing scheduled, we're not running

* Add more tests from #16271

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

@@ -15,6 +15,7 @@
 let Scheduler;
 let runtime;
 let performance;
+let cancelCallback;
 let scheduleCallback;
 let NormalPriority;
 
@@ -52,6 +53,7 @@ describe('SchedulerBrowser', () => {
       performance = window.performance;
       require('scheduler/src/SchedulerFeatureFlags').enableMessageLoopImplementation = enableMessageLoopImplementation;
       Scheduler = require('scheduler');
+      cancelCallback = Scheduler.unstable_cancelCallback;
       scheduleCallback = Scheduler.unstable_scheduleCallback;
       NormalPriority = Scheduler.unstable_NormalPriority;
     });
@@ -360,6 +362,15 @@ describe('SchedulerBrowser', () => {
     const enableMessageLoopImplementation = true;
     beforeAndAfterHooks(enableMessageLoopImplementation);
 
+    it('task that finishes before deadline', () => {
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('Task');
+      });
+      runtime.assertLog(['Post Message']);
+      runtime.fireMessageEvent();
+      runtime.assertLog(['Message Event', 'Task']);
+    });
+
     it('task with continutation', () => {
       scheduleCallback(NormalPriority, () => {
         runtime.log('Task');
@@ -385,7 +396,48 @@ describe('SchedulerBrowser', () => {
       runtime.assertLog(['Message Event', 'Continuation']);
     });
 
-    it('task that throws', () => {
+    it('multiple tasks', () => {
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('A');
+      });
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('B');
+      });
+      runtime.assertLog(['Post Message']);
+      runtime.fireMessageEvent();
+      runtime.assertLog(['Message Event', 'A', 'B']);
+    });
+
+    it('multiple tasks with a yield in between', () => {
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('A');
+        runtime.advanceTime(4999);
+      });
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('B');
+      });
+      runtime.assertLog(['Post Message']);
+      runtime.fireMessageEvent();
+      runtime.assertLog([
+        'Message Event',
+        'A',
+        // Ran out of time. Post a continuation event.
+        'Post Message',
+      ]);
+      runtime.fireMessageEvent();
+      runtime.assertLog(['Message Event', 'B']);
+    });
+
+    it('cancels tasks', () => {
+      const task = scheduleCallback(NormalPriority, () => {
+        runtime.log('Task');
+      });
+      runtime.assertLog(['Post Message']);
+      cancelCallback(task);
+      runtime.assertLog([]);
+    });
+
+    it('throws when a task errors then continues in a new event', () => {
       scheduleCallback(NormalPriority, () => {
         runtime.log('Oops!');
         throw Error('Oops!');
@@ -418,5 +470,24 @@ describe('SchedulerBrowser', () => {
       runtime.fireMessageEvent();
       runtime.assertLog(['Message Event', 'B']);
     });
+
+    it('schedule new task after a cancellation', () => {
+      let handle = scheduleCallback(NormalPriority, () => {
+        runtime.log('A');
+      });
+
+      runtime.assertLog(['Post Message']);
+      cancelCallback(handle);
+
+      runtime.fireMessageEvent();
+      runtime.assertLog(['Message Event']);
+
+      scheduleCallback(NormalPriority, () => {
+        runtime.log('B');
+      });
+      runtime.assertLog(['Post Message']);
+      runtime.fireMessageEvent();
+      runtime.assertLog(['Message Event', 'B']);
+    });
   });
 });

packages/scheduler/src/forks/SchedulerHostConfig.default.js

@@ -235,6 +235,8 @@ if (
           port.postMessage(null);
           throw error;
         }
+      } else {
+        isMessageLoopRunning = false;
       }
       // Yielding to the browser will give it a chance to paint, so we can
       // reset this.