Respect deadline on new HTTPClient.execute for async/await (#529)

dnadoba · web-flow · commit d95277640f82 · 2021-12-17T09:36:41.000+01:00
* Schedule deadline timeout
* Add state machine tests and enable skipped test for http1
diff --git a/Sources/AsyncHTTPClient/AsyncAwait/HTTPClient+execute.swift b/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)
         }
     }
 }
diff --git a/Sources/AsyncHTTPClient/AsyncAwait/Transaction+StateMachine.swift b/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
+            }
+        }
     }
 }
 
diff --git a/Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift b/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, *)
diff --git a/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests+XCTest.swift b/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests+XCTest.swift
@@ -37,6 +37,7 @@ extension AsyncAwaitEndToEndTests {
             ("testPostWithFragmentedAsyncSequenceOfLargeByteBuffers", testPostWithFragmentedAsyncSequenceOfLargeByteBuffers),
             ("testCanceling", testCanceling),
             ("testDeadline", testDeadline),
+            ("testImmediateDeadline", testImmediateDeadline),
             ("testInvalidURL", testInvalidURL),
         ]
     }
diff --git a/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift b/Tests/AsyncHTTPClientTests/AsyncAwaitEndToEndTests.swift
@@ -318,18 +318,16 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
         #endif
     }
 
-    func testCanceling() throws {
-        #if os(Linux)
-        #else
-        try XCTSkipIf(true, "test times out because of a swift concurrency bug on macOS: https://bugs.swift.org/browse/SR-15592")
-        #endif
+    func testCanceling() {
         #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()
+            defer { XCTAssertNoThrow(try bin.shutdown()) }
             let client = makeDefaultHTTPClient()
             defer { XCTAssertNoThrow(try client.syncShutdown()) }
             let logger = Logger(label: "HTTPClient", factory: StreamLogHandler.standardOutput(label:))
-            var request = HTTPClientRequest(url: "https://localhost:45678/offline")
+            var request = HTTPClientRequest(url: "http://localhost:\(bin.port)/offline")
             request.method = .POST
             let streamWriter = AsyncSequenceWriter<ByteBuffer>()
             request.body = .stream(length: nil, streamWriter)
@@ -345,8 +343,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
         #endif
     }
 
-    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")
+    func testDeadline() {
         #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,11 +355,33 @@ 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.deadlineExceeded)
+            }
+        }
+        #endif
+    }
+
+    func testImmediateDeadline() {
+        #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()
+            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: "http://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.readTimeout)
+                XCTAssertEqual($0 as? HTTPClientError, HTTPClientError.deadlineExceeded)
             }
+            print("done")
         }
         #endif
     }
@@ -386,7 +405,7 @@ final class AsyncAwaitEndToEndTests: XCTestCase {
 
 #if compiler(>=5.5) && canImport(_Concurrency)
 extension AsyncSequence where Element == ByteBuffer {
-    func collect() async throws -> ByteBuffer {
+    func collect() async rethrows -> ByteBuffer {
         try await self.reduce(into: ByteBuffer()) { accumulatingBuffer, nextBuffer in
             var nextBuffer = nextBuffer
             accumulatingBuffer.writeBuffer(&nextBuffer)
diff --git a/Tests/AsyncHTTPClientTests/Transaction+StateMachineTests+XCTest.swift b/Tests/AsyncHTTPClientTests/Transaction+StateMachineTests+XCTest.swift
@@ -27,6 +27,9 @@ extension Transaction_StateMachineTests {
         return [
             ("testRequestWasQueuedAfterWillExecuteRequestWasCalled", testRequestWasQueuedAfterWillExecuteRequestWasCalled),
             ("testRequestBodyStreamWasPaused", testRequestBodyStreamWasPaused),
+            ("testQueuedRequestGetsRemovedWhenDeadlineExceeded", testQueuedRequestGetsRemovedWhenDeadlineExceeded),
+            ("testScheduledRequestGetsRemovedWhenDeadlineExceeded", testScheduledRequestGetsRemovedWhenDeadlineExceeded),
+            ("testRequestWithHeadReceivedGetNotCancelledWhenDeadlineExceeded", testRequestWithHeadReceivedGetNotCancelledWhenDeadlineExceeded),
         ]
     }
 }
diff --git a/Tests/AsyncHTTPClientTests/Transaction+StateMachineTests.swift b/Tests/AsyncHTTPClientTests/Transaction+StateMachineTests.swift
@@ -15,6 +15,7 @@
 @testable import AsyncHTTPClient
 import NIOCore
 import NIOEmbedded
+import NIOHTTP1
 import XCTest
 
 final class Transaction_StateMachineTests: XCTestCase {
@@ -70,6 +71,87 @@ final class Transaction_StateMachineTests: XCTestCase {
         }
         #endif
     }
+
+    func testQueuedRequestGetsRemovedWhenDeadlineExceeded() {
+        #if compiler(>=5.5) && canImport(_Concurrency)
+        guard #available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *) else { return }
+        XCTAsyncTest {
+            func workaround(_ continuation: CheckedContinuation<HTTPClientResponse, Error>) {
+                var state = Transaction.StateMachine(continuation)
+                let queuer = MockTaskQueuer()
+
+                state.requestWasQueued(queuer)
+
+                let failAction = state.deadlineExceeded()
+                guard case .cancel(let continuation, let scheduler, nil, nil) = failAction else {
+                    return XCTFail("Unexpected fail action: \(failAction)")
+                }
+                XCTAssertIdentical(scheduler as? MockTaskQueuer, queuer)
+
+                continuation.resume(throwing: HTTPClientError.deadlineExceeded)
+            }
+
+            await XCTAssertThrowsError(try await withCheckedThrowingContinuation(workaround))
+        }
+        #endif
+    }
+
+    func testScheduledRequestGetsRemovedWhenDeadlineExceeded() {
+        #if compiler(>=5.5) && canImport(_Concurrency)
+        guard #available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *) else { return }
+        let eventLoop = EmbeddedEventLoop()
+        XCTAsyncTest {
+            func workaround(_ continuation: CheckedContinuation<HTTPClientResponse, Error>) {
+                var state = Transaction.StateMachine(continuation)
+                let executor = MockRequestExecutor(eventLoop: eventLoop)
+                let queuer = MockTaskQueuer()
+
+                XCTAssertEqual(state.willExecuteRequest(executor), .none)
+                state.requestWasQueued(queuer)
+
+                let failAction = state.deadlineExceeded()
+                guard case .cancel(let continuation, nil, let rexecutor, nil) = failAction else {
+                    return XCTFail("Unexpected fail action: \(failAction)")
+                }
+                XCTAssertIdentical(rexecutor as? MockRequestExecutor, executor)
+
+                continuation.resume(throwing: HTTPClientError.deadlineExceeded)
+            }
+
+            await XCTAssertThrowsError(try await withCheckedThrowingContinuation(workaround))
+        }
+        #endif
+    }
+
+    func testRequestWithHeadReceivedGetNotCancelledWhenDeadlineExceeded() {
+        #if compiler(>=5.5) && canImport(_Concurrency)
+        guard #available(macOS 12.0, iOS 15.0, watchOS 8.0, tvOS 15.0, *) else { return }
+        let eventLoop = EmbeddedEventLoop()
+        XCTAsyncTest {
+            func workaround(_ continuation: CheckedContinuation<HTTPClientResponse, Error>) {
+                var state = Transaction.StateMachine(continuation)
+                let executor = MockRequestExecutor(eventLoop: eventLoop)
+                let queuer = MockTaskQueuer()
+
+                XCTAssertEqual(state.willExecuteRequest(executor), .none)
+                state.requestWasQueued(queuer)
+                let head = HTTPResponseHead(version: .http1_1, status: .ok)
+                let receiveResponseHeadAction = state.receiveResponseHead(head)
+                guard case .succeedResponseHead(head, let continuation) = receiveResponseHeadAction else {
+                    return XCTFail("Unexpected action: \(receiveResponseHeadAction)")
+                }
+
+                let failAction = state.deadlineExceeded()
+                guard case .none = failAction else {
+                    return XCTFail("Unexpected fail action: \(failAction)")
+                }
+                continuation.resume(throwing: HTTPClientError.deadlineExceeded)
+            }
+
+            await XCTAssertThrowsError(try await withCheckedThrowingContinuation(workaround))
+        }
+        #endif
+    }
 }
 
 #if compiler(>=5.5) && canImport(_Concurrency)

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ extension AsyncAwaitEndToEndTests {`
`37`	`37`	`("testPostWithFragmentedAsyncSequenceOfLargeByteBuffers", testPostWithFragmentedAsyncSequenceOfLargeByteBuffers),`
`38`	`38`	`("testCanceling", testCanceling),`
`39`	`39`	`("testDeadline", testDeadline),`
	`40`	`+ ("testImmediateDeadline", testImmediateDeadline),`
`40`	`41`	`("testInvalidURL", testInvalidURL),`
`41`	`42`	`]`
`42`	`43`	`}`
Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,9 @@ extension Transaction_StateMachineTests {`
`27`	`27`	`return [`
`28`	`28`	`("testRequestWasQueuedAfterWillExecuteRequestWasCalled", testRequestWasQueuedAfterWillExecuteRequestWasCalled),`
`29`	`29`	`("testRequestBodyStreamWasPaused", testRequestBodyStreamWasPaused),`
	`30`	`+ ("testQueuedRequestGetsRemovedWhenDeadlineExceeded", testQueuedRequestGetsRemovedWhenDeadlineExceeded),`
	`31`	`+ ("testScheduledRequestGetsRemovedWhenDeadlineExceeded", testScheduledRequestGetsRemovedWhenDeadlineExceeded),`
	`32`	`+ ("testRequestWithHeadReceivedGetNotCancelledWhenDeadlineExceeded", testRequestWithHeadReceivedGetNotCancelledWhenDeadlineExceeded),`
`30`	`33`	`]`
`31`	`34`	`}`
`32`	`35`	`}`