From b22cf43d7c5fc77496086a2d01dcc505c3a15bc1 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Wed, 17 Jul 2019 13:14:47 -0700
Subject: [PATCH] [Scheduler] Bugfix: Cancelling a continuation

Cancelling the original task should also cancel its continuation.
---
 packages/scheduler/src/Scheduler.js                | 10 ++--------
 packages/scheduler/src/__tests__/Scheduler-test.js | 13 +++++++++++++
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/packages/scheduler/src/Scheduler.js b/packages/scheduler/src/Scheduler.js
index 251c3ad11ea6d..cdb9930d4c973 100644
--- a/packages/scheduler/src/Scheduler.js
+++ b/packages/scheduler/src/Scheduler.js
@@ -144,14 +144,8 @@ function flushTask(task, currentTime) {
   // with the same priority and expiration as the just-finished callback.
   if (typeof continuationCallback === 'function') {
     var expirationTime = task.expirationTime;
-    var continuationTask = {
-      callback: continuationCallback,
-      priorityLevel: task.priorityLevel,
-      startTime: task.startTime,
-      expirationTime,
-      next: null,
-      previous: null,
-    };
+    var continuationTask = task;
+    continuationTask.callback = continuationCallback;
 
     // Insert the new callback into the list, sorted by its timeout. This is
     // almost the same as the code in `scheduleCallback`, except the callback
diff --git a/packages/scheduler/src/__tests__/Scheduler-test.js b/packages/scheduler/src/__tests__/Scheduler-test.js
index c2b03cedb4449..07377ba0d0e36 100644
--- a/packages/scheduler/src/__tests__/Scheduler-test.js
+++ b/packages/scheduler/src/__tests__/Scheduler-test.js
@@ -275,6 +275,19 @@ describe('Scheduler', () => {
     },
   );
 
+  it('cancelling a continuation', () => {
+    const task = scheduleCallback(NormalPriority, () => {
+      Scheduler.unstable_yieldValue('Yield');
+      return () => {
+        Scheduler.unstable_yieldValue('Continuation');
+      };
+    });
+
+    expect(Scheduler).toFlushAndYieldThrough(['Yield']);
+    cancelCallback(task);
+    expect(Scheduler).toFlushWithoutYielding();
+  });
+
   it('top-level immediate callbacks fire in a subsequent task', () => {
     scheduleCallback(ImmediatePriority, () =>
       Scheduler.unstable_yieldValue('A'),