From 2d68bd09603fc54d927240003810db9aa6f2e3e5 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Mon, 19 Aug 2019 21:20:21 +0100
Subject: [PATCH] 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
---
 .../SchedulerBrowser-test.internal.js         | 73 ++++++++++++++++++-
 .../src/forks/SchedulerHostConfig.default.js  |  2 +
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/packages/scheduler/src/__tests__/SchedulerBrowser-test.internal.js b/packages/scheduler/src/__tests__/SchedulerBrowser-test.internal.js
index dba2c3b5da56..657e59762063 100644
--- a/packages/scheduler/src/__tests__/SchedulerBrowser-test.internal.js
+++ b/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']);
+    });
   });
 });
diff --git a/packages/scheduler/src/forks/SchedulerHostConfig.default.js b/packages/scheduler/src/forks/SchedulerHostConfig.default.js
index 8e6afbcc8ad1..c11fca6cd029 100644
--- a/packages/scheduler/src/forks/SchedulerHostConfig.default.js
+++ 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.