Move cancellation to LambdaRuntimeClient.nextInvocation

Author: adam-fowler

Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift

@@ -145,22 +145,28 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
     }
 
     func nextInvocation() async throws -> (Invocation, Writer) {
-        switch self.lambdaState {
-        case .idle:
-            self.lambdaState = .waitingForNextInvocation
-            let handler = try await self.makeOrGetConnection()
-            let invocation = try await handler.nextInvocation()
-            guard case .waitingForNextInvocation = self.lambdaState else {
+        try await withTaskCancellationHandler {
+            switch self.lambdaState {
+            case .idle:
+                self.lambdaState = .waitingForNextInvocation
+                let handler = try await self.makeOrGetConnection()
+                let invocation = try await handler.nextInvocation()
+                guard case .waitingForNextInvocation = self.lambdaState else {
+                    fatalError("Invalid state: \(self.lambdaState)")
+                }
+                self.lambdaState = .waitingForResponse(requestID: invocation.metadata.requestID)
+                return (invocation, Writer(runtimeClient: self))
+
+            case .waitingForNextInvocation,
+                .waitingForResponse,
+                .sendingResponse,
+                .sentResponse:
                 fatalError("Invalid state: \(self.lambdaState)")
             }
-            self.lambdaState = .waitingForResponse(requestID: invocation.metadata.requestID)
-            return (invocation, Writer(runtimeClient: self))
-
-        case .waitingForNextInvocation,
-            .waitingForResponse,
-            .sendingResponse,
-            .sentResponse:
-            fatalError("Invalid state: \(self.lambdaState)")
+        } onCancel: {
+            Task {
+                await self.close()
+            }
         }
     }
 
@@ -469,37 +475,10 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
     func nextInvocation(isolation: isolated (any Actor)? = #isolation) async throws -> Invocation {
         switch self.state {
         case .connected(let context, .idle):
-            return try await withTaskCancellationHandler {
-                try Task.checkCancellation()
-                return try await withCheckedThrowingContinuation {
-                    (continuation: CheckedContinuation<Invocation, any Error>) in
-                    self.state = .connected(context, .waitingForNextInvocation(continuation))
-
-                    let unsafeContext = NIOLoopBound(context, eventLoop: context.eventLoop)
-                    context.eventLoop.execute { [nextInvocationPath, defaultHeaders] in
-                        // Send next request. The function `sendNextRequest` requires `self` which is not
-                        // Sendable so just inlined the code instead
-                        let httpRequest = HTTPRequestHead(
-                            version: .http1_1,
-                            method: .GET,
-                            uri: nextInvocationPath,
-                            headers: defaultHeaders
-                        )
-                        let context = unsafeContext.value
-                        context.write(Self.wrapOutboundOut(.head(httpRequest)), promise: nil)
-                        context.write(Self.wrapOutboundOut(.end(nil)), promise: nil)
-                        context.flush()
-                    }
-                }
-            } onCancel: {
-                switch self.state {
-                case .connected(_, .waitingForNextInvocation(let continuation)):
-                    continuation.resume(throwing: CancellationError())
-                case .connected(_, .idle):
-                    break
-                default:
-                    fatalError("Invalid state: \(self.state)")
-                }
+            return try await withCheckedThrowingContinuation {
+                (continuation: CheckedContinuation<Invocation, any Error>) in
+                self.state = .connected(context, .waitingForNextInvocation(continuation))
+                self.sendNextRequest(context: context)
             }
 
         case .connected(_, .sendingResponse),
@@ -846,6 +825,12 @@ extension LambdaChannelHandler: ChannelInboundHandler {
 
     func channelInactive(context: ChannelHandlerContext) {
         // fail any pending responses with last error or assume peer disconnected
+        switch self.state {
+        case .connected(_, .waitingForNextInvocation(let continuation)):
+            continuation.resume(throwing: self.lastError ?? ChannelError.ioOnClosedChannel)
+        default:
+            break
+        }
 
         // we don't need to forward channelInactive to the delegate, as the delegate observes the
         // closeFuture

Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTests.swift

@@ -89,20 +89,52 @@ struct LambdaRuntimeClientTests {
 
     @Test
     func testCancellation() async throws {
-        try await LambdaRuntimeClient.withRuntimeClient(
-            configuration: .init(ip: "127.0.0.1", port: 7000),
-            eventLoop: NIOSingletons.posixEventLoopGroup.next(),
-            logger: self.logger
-        ) { runtimeClient in
-            try await withThrowingTaskGroup(of: Void.self) { group in
-                group.addTask {
-                    while true {
-                        _ = try await runtimeClient.nextInvocation()
+        struct HappyBehavior: LambdaServerBehavior {
+            let requestId = UUID().uuidString
+            let event = "hello"
+
+            func getInvocation() -> GetInvocationResult {
+                .success((self.requestId, self.event))
+            }
+
+            func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
+                #expect(self.requestId == requestId)
+                #expect(self.event == response)
+                return .success(())
+            }
+
+            func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+                Issue.record("should not report error")
+                return .failure(.internalServerError)
+            }
+
+            func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+                Issue.record("should not report init error")
+                return .failure(.internalServerError)
+            }
+        }
+
+        try await withMockServer(behaviour: HappyBehavior()) { port in
+            try await LambdaRuntimeClient.withRuntimeClient(
+                configuration: .init(ip: "127.0.0.1", port: port),
+                eventLoop: NIOSingletons.posixEventLoopGroup.next(),
+                logger: self.logger
+            ) { runtimeClient in
+                try await withThrowingTaskGroup(of: Void.self) { group in
+                    group.addTask {
+                        while true {
+                            print("Waiting")
+                            let (_, writer) = try await runtimeClient.nextInvocation()
+                            try await Task {
+                                try await writer.write(ByteBuffer(string: "hello"))
+                                try await writer.finish()
+                            }.value
+                        }
                     }
+                    // wait a small amount to ensure we are waiting for continuation
+                    try await Task.sleep(for: .milliseconds(100))
+                    group.cancelAll()
                 }
-                // wait a small amount to ensure we are waiting for continuation
-                try await Task.sleep(for: .milliseconds(100))
-                group.cancelAll()
             }
         }
     }