store executed redirectTask in RequestBag and propagate cancellation

Author: gregcotten

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -956,7 +956,7 @@ internal struct RedirectHandler<ResponseType> {
         status: HTTPResponseStatus,
         to redirectURL: URL,
         promise: EventLoopPromise<ResponseType>
-    ) {
+    ) -> HTTPClient.Task<ResponseType>? {
         do {
             var redirectState = self.redirectState
             try redirectState.redirect(to: redirectURL.absoluteString)
@@ -976,13 +976,19 @@ internal struct RedirectHandler<ResponseType> {
                 headers: headers,
                 body: body
             )
-            self.execute(newRequest, redirectState).futureResult.whenComplete { result in
+
+            let newTask = self.execute(newRequest, redirectState)
+
+            newTask.futureResult.whenComplete { result in
                 promise.futureResult.eventLoop.execute {
                     promise.completeWith(result)
                 }
             }
+
+            return newTask
         } catch {
             promise.fail(error)
+            return nil
         }
     }
 }

Sources/AsyncHTTPClient/RequestBag+StateMachine.swift

@@ -594,6 +594,7 @@ extension RequestBag.StateMachine {
     enum FailAction {
         case failTask(Error, HTTPRequestScheduler?, HTTPRequestExecutor?)
         case cancelExecutor(HTTPRequestExecutor)
+        case propagateCancellation
         case none
     }
 
@@ -624,8 +625,11 @@ extension RequestBag.StateMachine {
             self.state = .finished(error: error)
             return .failTask(error, nil, nil)
         case .finished(.none):
-            // An error occurred after the request has finished. Ignore...
-            return .none
+            if (error as? HTTPClientError) == .cancelled {
+                return .propagateCancellation
+            } else {
+                return .none
+            }
         case .deadlineExceededWhileQueued:
             let realError: Error = {
                 if (error as? HTTPClientError) == .cancelled {

Sources/AsyncHTTPClient/RequestBag.swift

@@ -43,6 +43,9 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
     // the consume body part stack depth is synchronized on the task event loop.
     private var consumeBodyPartStackDepth: Int
 
+    // if a redirect occurs, we store the task for it so we can propagate cancellation
+    private var redirectTask: HTTPClient.Task<Delegate.Response>? = nil
+
     // MARK: HTTPClientTask properties
 
     var logger: Logger {
@@ -234,7 +237,7 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             executor.demandResponseBodyStream(self)
 
         case .redirect(let executor, let handler, let head, let newURL):
-            handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
+            self.redirectTask = handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
             executor.cancelRequest(self)
 
         case .forwardResponseHead(let head):
@@ -258,7 +261,7 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             executor.demandResponseBodyStream(self)
 
         case .redirect(let executor, let handler, let head, let newURL):
-            handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
+            self.redirectTask = handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
             executor.cancelRequest(self)
 
         case .forwardResponsePart(let part):
@@ -294,7 +297,7 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             }
 
         case .redirect(let handler, let head, let newURL):
-            handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
+            self.redirectTask = handler.redirect(status: head.status, to: newURL, promise: self.task.promise)
         }
     }
 
@@ -368,6 +371,8 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             self.failTask0(error)
         case .cancelExecutor(let executor):
             executor.cancelRequest(self)
+        case .propagateCancellation:
+            self.redirectTask?.cancel()
         case .none:
             break
         }