Schedule deadline timeout

Author: dnadoba

Sources/AsyncHTTPClient/AsyncAwait/HTTPClient+execute.swift

@@ -94,13 +94,20 @@ extension HTTPClient {
         let cancelHandler = TransactionCancelHandler()
 
         return try await withTaskCancellationHandler(operation: { () async throws -> HTTPClientResponse in
-            try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<HTTPClientResponse, Swift.Error>) -> Void in
+            let eventLoop = self.eventLoopGroup.any()
+            let deadlineTask = eventLoop.scheduleTask(deadline: deadline) {
+                cancelHandler.cancel(reason: .deadlineExceeded)
+            }
+            defer {
+                deadlineTask.cancel()
+            }
+            return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<HTTPClientResponse, Swift.Error>) -> Void in
                 let transaction = Transaction(
                     request: request,
                     requestOptions: .init(idleReadTimeout: nil),
                     logger: logger,
                     connectionDeadline: deadline,
-                    preferredEventLoop: self.eventLoopGroup.next(),
+                    preferredEventLoop: eventLoop,
                     responseContinuation: continuation
                 )
 
@@ -109,7 +116,7 @@ extension HTTPClient {
                 self.poolManager.executeRequest(transaction)
             }
         }, onCancel: {
-            cancelHandler.cancel()
+            cancelHandler.cancel(reason: .taskCanceled)
         })
     }
 }
@@ -119,22 +126,38 @@ extension HTTPClient {
 /// in the `body` closure and cancelation from the `onCancel` closure  of `withTaskCancellationHandler`.
 @available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *)
 private actor TransactionCancelHandler {
+    enum CancelReason {
+        /// swift concurrency task was canceled
+        case taskCanceled
+        /// deadline timeout
+        case deadlineExceeded
+    }
+
     private enum State {
         case initialised
         case register(Transaction)
-        case cancelled
+        case cancelled(CancelReason)
     }
 
     private var state: State = .initialised
 
     init() {}
 
+    private func cancelTransaction(_ transaction: Transaction, for reason: CancelReason) {
+        switch reason {
+        case .taskCanceled:
+            transaction.cancel()
+        case .deadlineExceeded:
+            transaction.deadlineExceeded()
+        }
+    }
+
     private func _registerTransaction(_ transaction: Transaction) {
         switch self.state {
         case .initialised:
             self.state = .register(transaction)
-        case .cancelled:
-            transaction.cancel()
+        case .cancelled(let reason):
+            self.cancelTransaction(transaction, for: reason)
         case .register:
             preconditionFailure("transaction already set")
         }
@@ -146,21 +169,21 @@ private actor TransactionCancelHandler {
         }
     }
 
-    private func _cancel() {
+    private func _cancel(reason: CancelReason) {
         switch self.state {
-        case .register(let bag):
-            self.state = .cancelled
-            bag.cancel()
+        case .register(let transaction):
+            self.state = .cancelled(reason)
+            self.cancelTransaction(transaction, for: reason)
         case .cancelled:
             break
         case .initialised:
-            self.state = .cancelled
+            self.state = .cancelled(reason)
         }
     }
 
-    nonisolated func cancel() {
+    nonisolated func cancel(reason: CancelReason) {
         Task {
-            await self._cancel()
+            await self._cancel(reason: reason)
         }
     }
 }

Sources/AsyncHTTPClient/AsyncAwait/Transaction+StateMachine.swift

@@ -674,6 +674,67 @@ extension Transaction {
                 preconditionFailure("Already received an eof or error before. Must not receive further events. Invalid state: \(self.state)")
             }
         }
+
+        enum DeadlineExceededAction {
+            case none
+            /// fail response before head received. scheduler and executor are exclusive here.
+            case cancel(
+                requestContinuation: CheckedContinuation<HTTPClientResponse, Error>,
+                scheduler: HTTPRequestScheduler?,
+                executor: HTTPRequestExecutor?,
+                bodyStreamContinuation: CheckedContinuation<Void, Error>?
+            )
+        }
+
+        mutating func deadlineExceeded() -> DeadlineExceededAction {
+            let error = HTTPClientError.deadlineExceeded
+            switch self.state {
+            case .initialized(let continuation):
+                self.state = .finished(error: error, nil)
+                return .cancel(
+                    requestContinuation: continuation,
+                    scheduler: nil,
+                    executor: nil,
+                    bodyStreamContinuation: nil
+                )
+
+            case .queued(let continuation, let scheduler):
+                self.state = .finished(error: error, nil)
+                return .cancel(
+                    requestContinuation: continuation,
+                    scheduler: scheduler,
+                    executor: nil,
+                    bodyStreamContinuation: nil
+                )
+
+            case .executing(let context, let requestStreamState, .waitingForResponseHead):
+                switch requestStreamState {
+                case .paused(continuation: .some(let continuation)):
+                    self.state = .finished(error: error, nil)
+                    return .cancel(
+                        requestContinuation: context.continuation,
+                        scheduler: nil,
+                        executor: context.executor,
+                        bodyStreamContinuation: continuation
+                    )
+                case .requestHeadSent, .finished, .producing, .paused(continuation: .none):
+                    self.state = .finished(error: error, nil)
+                    return .cancel(
+                        requestContinuation: context.continuation,
+                        scheduler: nil,
+                        executor: context.executor,
+                        bodyStreamContinuation: nil
+                    )
+                }
+
+            case .executing, .finished:
+                // The user specified deadline is only used until we received the response head.
+                // If we already received the head, we have also resumed the continuation and
+                // therefore return the HTTPClientResponse to the user. We do not want to cancel
+                // the request body streaming nor the response body streaming afterwards.
+                return .none
+            }
+        }
     }
 }
 

Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift

@@ -294,6 +294,26 @@ extension Transaction: HTTPExecutableRequest {
             bodyStreamContinuation.resume(throwing: error)
         }
     }
+
+    func deadlineExceeded() {
+        let action = self.stateLock.withLock {
+            self.state.deadlineExceeded()
+        }
+        self.performDeadlineExceededAction(action)
+    }
+
+    private func performDeadlineExceededAction(_ action: StateMachine.DeadlineExceededAction) {
+        switch action {
+        case .cancel(let requestContinuation, let scheduler, let executor, let bodyStreamContinuation):
+            requestContinuation.resume(throwing: HTTPClientError.deadlineExceeded)
+            scheduler?.cancelRequest(self)
+            executor?.cancelRequest(self)
+            bodyStreamContinuation?.resume(throwing: HTTPClientError.deadlineExceeded)
+
+        case .none:
+            break
+        }
+    }
 }
 
 @available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *)

Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests+XCTest.swift

@@ -37,7 +37,7 @@ extension AsyncAwaitEndToEndTests {
             ("testPostWithFragmentedAsyncSequenceOfLargeByteBuffers", testPostWithFragmentedAsyncSequenceOfLargeByteBuffers),
             ("testCanceling", testCanceling),
             ("testDeadline", testDeadline),
-            ("testConnectTimeout", testConnectTimeout),
+            ("testImmediateDeadline", testImmediateDeadline),
             ("testInvalidURL", testInvalidURL),
         ]
     }

Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift

@@ -346,7 +346,6 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
     }
 
     func testDeadline() throws {
-        try XCTSkipIf(true, "deadline is currently not correctly implemented. We only use it to timeout connection establishment. will be fixed in a follow up PR")
         #if compiler(>=5.5) && canImport(_Concurrency)
         guard #available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *) else { return }
         XCTAsyncTest(timeout: 5) {
@@ -358,15 +357,39 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
             let request = HTTPClientRequest(url: "https://localhost:\(bin.port)/wait")
 
             let task = Task<HTTPClientResponse, Error> { [request] in
-                try await client.execute(request, deadline: .now() + .seconds(1), logger: logger)
+                try await client.execute(request, deadline: .now() + .milliseconds(100), logger: logger)
             }
             await XCTAssertThrowsError(try await task.value) {
-                XCTAssertEqual($0 as? HTTPClientError, HTTPClientError.readTimeout)
+                XCTAssertEqual($0 as? HTTPClientError, HTTPClientError.deadlineExceeded)
             }
         }
         #endif
     }
 
+    func testImmediateDeadline() throws {
+        // does not work on nether Linux nor Darwin
+        try XCTSkipIf(true, "test times out because of a swift concurrency bug: https://bugs.swift.org/browse/SR-15592")
+        #if compiler(>=5.5) && canImport(_Concurrency)
+        guard #available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *) else { return }
+        XCTAsyncTest(timeout: 5) {
+            let bin = HTTPBin(.http2(compress: false))
+            defer { XCTAssertNoThrow(try bin.shutdown()) }
+            let client = makeDefaultHTTPClient()
+            defer { XCTAssertNoThrow(try client.syncShutdown()) }
+            let logger = Logger(label: "HTTPClient", factory: StreamLogHandler.standardOutput(label:))
+            let request = HTTPClientRequest(url: "https://localhost:\(bin.port)/wait")
+
+            let task = Task<HTTPClientResponse, Error> { [request] in
+                try await client.execute(request, deadline: .now(), logger: logger)
+            }
+            await XCTAssertThrowsError(try await task.value) {
+                XCTAssertEqual($0 as? HTTPClientError, HTTPClientError.deadlineExceeded)
+            }
+            print("done")
+        }
+        #endif
+    }
+
     func testInvalidURL() {
         #if compiler(>=5.5) && canImport(_Concurrency)
         guard #available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *) else { return }