From 9a7ae34232b5de8d4b2fa507af0c30a23f5d2ea7 Mon Sep 17 00:00:00 2001
From: Nikita Lutsenko <nlutsenko@fb.com>
Date: Wed, 24 Jun 2015 19:08:17 -0700
Subject: [PATCH] Fixed never completed task if continuation returns a task and
 cancellation was requested.

---
 Bolts/Common/BFTask.m  |  4 ++--
 BoltsTests/TaskTests.m | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/Bolts/Common/BFTask.m b/Bolts/Common/BFTask.m
index 58f89cff7..a708ef09b 100644
--- a/Bolts/Common/BFTask.m
+++ b/Bolts/Common/BFTask.m
@@ -340,7 +340,7 @@ - (instancetype)continueWithExecutor:(BFExecutor *)executor
             if ([result isKindOfClass:[BFTask class]]) {
 
                 id (^setupWithTask) (BFTask *) = ^id(BFTask *task) {
-                    if (task.cancelled) {
+                    if (cancellationToken.cancellationRequested || task.cancelled) {
                         [tcs cancel];
                     } else if (task.exception) {
                         tcs.exception = task.exception;
@@ -357,7 +357,7 @@ - (instancetype)continueWithExecutor:(BFExecutor *)executor
                 if (resultTask.completed) {
                     setupWithTask(resultTask);
                 } else {
-                    [resultTask continueWithBlock:setupWithTask cancellationToken:cancellationToken];
+                    [resultTask continueWithBlock:setupWithTask];
                 }
 
             } else {
diff --git a/BoltsTests/TaskTests.m b/BoltsTests/TaskTests.m
index 4bea30610..41c67bdfd 100644
--- a/BoltsTests/TaskTests.m
+++ b/BoltsTests/TaskTests.m
@@ -692,4 +692,19 @@ - (void)testDescription {
     XCTAssertTrue([expected isEqualToString:description]);
 }
 
+- (void)testReturnTaskFromContinuationWithCancellation {
+    BFCancellationTokenSource *cts = [BFCancellationTokenSource cancellationTokenSource];
+
+    XCTestExpectation *expectation = [self expectationWithDescription:@"task"];
+    [[[BFTask taskWithDelay:1] continueWithBlock:^id(BFTask *task) {
+        [cts cancel];
+        return [BFTask taskWithDelay:10];
+    } cancellationToken:cts.token] continueWithBlock:^id(BFTask *task) {
+        XCTAssertTrue(task.cancelled);
+        [expectation fulfill];
+        return nil;
+    }];
+    [self waitForExpectationsWithTimeout:10.0 handler:nil];
+}
+
 @end