Add cancellation handling for nextInvocation()

Author: adam-fowler

Sources/AWSLambdaRuntimeCore/Lambda.swift

@@ -37,25 +37,31 @@ public enum Lambda {
     ) async throws where Handler: StreamingLambdaHandler {
         var handler = handler
 
-        while !Task.isCancelled {
-            let (invocation, writer) = try await runtimeClient.nextInvocation()
+        do {
+            while !Task.isCancelled {
+                let (invocation, writer) = try await runtimeClient.nextInvocation()
 
-            do {
-                try await handler.handle(
-                    invocation.event,
-                    responseWriter: writer,
-                    context: LambdaContext(
-                        requestID: invocation.metadata.requestID,
-                        traceID: invocation.metadata.traceID,
-                        invokedFunctionARN: invocation.metadata.invokedFunctionARN,
-                        deadline: DispatchWallTime(millisSinceEpoch: invocation.metadata.deadlineInMillisSinceEpoch),
-                        logger: logger
+                do {
+                    try await handler.handle(
+                        invocation.event,
+                        responseWriter: writer,
+                        context: LambdaContext(
+                            requestID: invocation.metadata.requestID,
+                            traceID: invocation.metadata.traceID,
+                            invokedFunctionARN: invocation.metadata.invokedFunctionARN,
+                            deadline: DispatchWallTime(
+                                millisSinceEpoch: invocation.metadata.deadlineInMillisSinceEpoch
+                            ),
+                            logger: logger
+                        )
                     )
-                )
-            } catch {
-                try await writer.reportError(error)
-                continue
+                } catch {
+                    try await writer.reportError(error)
+                    continue
+                }
             }
+        } catch is CancellationError {
+            // don't allow cancellation error to propagate further
         }
     }
 

Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift

@@ -406,6 +406,17 @@ private protocol LambdaChannelHandlerDelegate {
     func connectionErrorHappened(_ error: any Error, channel: any Channel)
 }
 
+struct UnsafeContext: @unchecked Sendable {
+    private let _context: ChannelHandlerContext
+    var context: ChannelHandlerContext {
+        self._context.eventLoop.preconditionInEventLoop()
+        return _context
+    }
+    init(_ context: ChannelHandlerContext) {
+        self._context = context
+    }
+}
+
 private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate> {
     let nextInvocationPath = Consts.invocationURLPrefix + Consts.getNextInvocationURLSuffix
 
@@ -465,10 +476,37 @@ 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 withCheckedThrowingContinuation {
-                (continuation: CheckedContinuation<Invocation, any Error>) in
-                self.state = .connected(context, .waitingForNextInvocation(continuation))
-                self.sendNextRequest(context: context)
+            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 = UnsafeContext(context)
+                    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.context
+                        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)")
+                }
             }
 
         case .connected(_, .sendingResponse),

Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTests.swift

@@ -86,4 +86,24 @@ 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()
+                    }
+                }
+                // wait a small amount to ensure we are waiting for continuation
+                try await Task.sleep(for: .milliseconds(100))
+                group.cancelAll()
+            }
+        }
+    }
 }