Accept multiple POST /invoke requests to allow parallel testing

Examples/Streaming/README.md

@@ -82,7 +82,7 @@ You can test the function locally before deploying:
 swift run 
 
 # In another terminal, test with curl:
-curl -v \
+curl -v --output response.txt \
   --header "Content-Type: application/json" \
   --data '"this is not used"' \
   http://127.0.0.1:7000/invoke

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -95,8 +95,8 @@ extension Lambda {
 internal struct LambdaHTTPServer {
     private let invocationEndpoint: String
 
-    private let invocationPool = Pool<LocalServerInvocation>()
-    private let responsePool = Pool<LocalServerResponse>()
+    private let invocationPool = Pool<LocalServerInvocation>(name: "Invocation Pool")
+    private let responsePool = Pool<LocalServerResponse>(name: "Response Pool")
 
     private init(
         invocationEndpoint: String?
@@ -272,7 +272,7 @@ internal struct LambdaHTTPServer {
 
                             // for streaming requests, push a partial head response
                             if self.isStreamingResponse(requestHead) {
-                                await self.responsePool.push(
+                                self.responsePool.push(
                                     LocalServerResponse(
                                         id: requestId,
                                         status: .ok
@@ -286,7 +286,7 @@ internal struct LambdaHTTPServer {
                             // if this is a request from a Streaming Lambda Handler,
                             // stream the response instead of buffering it
                             if self.isStreamingResponse(requestHead) {
-                                await self.responsePool.push(
+                                self.responsePool.push(
                                     LocalServerResponse(id: requestId, body: body)
                                 )
                             } else {
@@ -298,7 +298,7 @@ internal struct LambdaHTTPServer {
 
                             if self.isStreamingResponse(requestHead) {
                                 // for streaming response, send the final response
-                                await self.responsePool.push(
+                                self.responsePool.push(
                                     LocalServerResponse(id: requestId, final: true)
                                 )
                             } else {
@@ -388,34 +388,55 @@ internal struct LambdaHTTPServer {
             // we always accept the /invoke request and push them to the pool
             let requestId = "\(DispatchTime.now().uptimeNanoseconds)"
             logger[metadataKey: "requestId"] = "\(requestId)"
+
             logger.trace("/invoke received invocation, pushing it to the pool and wait for a lambda response")
-            await self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body))
+            self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body))
 
             // wait for the lambda function to process the request
-            for try await response in self.responsePool {
-                logger[metadataKey: "response requestId"] = "\(response.requestId ?? "nil")"
-                logger.trace("Received response to return to client")
-                if response.requestId == requestId {
-                    logger.trace("/invoke requestId is valid, sending the response")
-                    // send the response to the client
-                    // if the response is final, we can send it and return
-                    // if the response is not final, we can send it and wait for the next response
-                    try await self.sendResponse(response, outbound: outbound, logger: logger)
-                    if response.final == true {
-                        logger.trace("/invoke returning")
-                        return  // if the response is final, we can return and close the connection
+            // when POST /invoke is called multiple times before a response is processed,
+            // the `for try await ... in` loop will throw an error and we will return a 400 error to the client
+            do {
+                for try await response in self.responsePool {
+                    logger[metadataKey: "response requestId"] = "\(response.requestId ?? "nil")"
+                    logger.trace("Received response to return to client")
+                    if response.requestId == requestId {
+                        logger.trace("/invoke requestId is valid, sending the response")
+                        // send the response to the client
+                        // if the response is final, we can send it and return
+                        // if the response is not final, we can send it and wait for the next response
+                        try await self.sendResponse(response, outbound: outbound, logger: logger)
+                        if response.final == true {
+                            logger.trace("/invoke returning")
+                            return  // if the response is final, we can return and close the connection
+                        }
+                    } else {
+                        logger.error(
+                            "Received response for a different requestId",
+                            metadata: ["response requestId": "\(response.requestId ?? "")"]
+                        )
+                        let response = LocalServerResponse(
+                            id: requestId,
+                            status: .badRequest,
+                            body: ByteBuffer(string: "The responseId is not equal to the requestId.")
+                        )
+                        try await self.sendResponse(response, outbound: outbound, logger: logger)
                     }
-                } else {
-                    logger.error(
-                        "Received response for a different request id",
-                        metadata: ["response requestId": "\(response.requestId ?? "")"]
-                    )
-                    // should we return an error here ? Or crash as this is probably a programming error?
                 }
+                // What todo when there is no more responses to process?
+                // This should not happen as the async iterator blocks until there is a response to process
+                fatalError("No more responses to process - the async for loop should not return")
+            } catch is LambdaHTTPServer.Pool<LambdaHTTPServer.LocalServerResponse>.PoolError {
+                // detect concurrent invocations of POST and gently decline the requests while we're processing one.
+                let response = LocalServerResponse(
+                    id: requestId,
+                    status: .badRequest,
+                    body: ByteBuffer(
+                        string:
+                            "It is not allowed to invoke multiple Lambda function executions in parallel. (The Lambda runtime environment on AWS will never do that)"
+                    )
+                )
+                try await self.sendResponse(response, outbound: outbound, logger: logger)
             }
-            // What todo when there is no more responses to process?
-            // This should not happen as the async iterator blocks until there is a response to process
-            fatalError("No more responses to process - the async for loop should not return")
 
         // client uses incorrect HTTP method
         case (_, let url) where url.hasSuffix(self.invocationEndpoint):
@@ -457,7 +478,7 @@ internal struct LambdaHTTPServer {
             }
             // enqueue the lambda function response to be served as response to the client /invoke
             logger.trace("/:requestId/response received response", metadata: ["requestId": "\(requestId)"])
-            await self.responsePool.push(
+            self.responsePool.push(
                 LocalServerResponse(
                     id: requestId,
                     status: .accepted,
@@ -488,7 +509,7 @@ internal struct LambdaHTTPServer {
             }
             // enqueue the lambda function response to be served as response to the client /invoke
             logger.trace("/:requestId/response received response", metadata: ["requestId": "\(requestId)"])
-            await self.responsePool.push(
+            self.responsePool.push(
                 LocalServerResponse(
                     id: requestId,
                     status: .internalServerError,
@@ -548,18 +569,21 @@ internal struct LambdaHTTPServer {
     /// This data structure is shared between instances of the HTTPHandler
     /// (one instance to serve requests from the Lambda function and one instance to serve requests from the client invoking the lambda function).
     internal final class Pool<T>: AsyncSequence, AsyncIteratorProtocol, Sendable where T: Sendable {
+        private let poolName: String
+        internal init(name: String = "Pool") { self.poolName = name }
+
         typealias Element = T
 
         enum State: ~Copyable {
             case buffer(Deque<T>)
-            case continuation(CheckedContinuation<T, any Error>?)
+            case continuation(CheckedContinuation<T, any Error>)
         }
 
         private let lock = Mutex<State>(.buffer([]))
 
         /// enqueue an element, or give it back immediately to the iterator if it is waiting for an element
-        public func push(_ invocation: T) async {
-            // if the iterator is waiting for an element, give it to it
+        public func push(_ invocation: T) {
+            // if the iterator is waiting for an element on `next()``, give it to it
             // otherwise, enqueue the element
             let maybeContinuation = self.lock.withLock { state -> CheckedContinuation<T, any Error>? in
                 switch consume state {
@@ -574,6 +598,7 @@ internal struct LambdaHTTPServer {
                 }
             }
 
+            // Resume continuation outside the lock to prevent potential deadlocks
             maybeContinuation?.resume(returning: invocation)
         }
 
@@ -585,42 +610,71 @@ internal struct LambdaHTTPServer {
 
             return try await withTaskCancellationHandler {
                 try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<T, any Error>) in
-                    let nextAction = self.lock.withLock { state -> T? in
+                    let nextAction: Result<T, PoolError>? = self.lock.withLock { state -> Result<T, PoolError>? in
                         switch consume state {
                         case .buffer(var buffer):
                             if let first = buffer.popFirst() {
                                 state = .buffer(buffer)
-                                return first
+                                return .success(first)
                             } else {
                                 state = .continuation(continuation)
                                 return nil
                             }
 
-                        case .continuation:
-                            fatalError("Concurrent invocations to next(). This is illegal.")
+                        case .continuation(let previousContinuation):
+                            state = .buffer([])
+                            return .failure(PoolError(cause: .nextCalledTwice(previousContinuation)))
                         }
                     }
 
-                    guard let nextAction else { return }
-
-                    continuation.resume(returning: nextAction)
+                    switch nextAction {
+                    case .success(let action):
+                        continuation.resume(returning: action)
+                    case .failure(let error):
+                        if case let .nextCalledTwice(prevContinuation) = error.cause {
+                            prevContinuation.resume(throwing: error)
+                        }
+                        continuation.resume(throwing: error)
+                    case .none:
+                        // do nothing - continuation is stored in state
+                        break
+                    }
                 }
             } onCancel: {
-                self.lock.withLock { state in
+                // Ensure we properly handle cancellation by checking if we have a stored continuation
+                let continuationToCancel = self.lock.withLock { state -> CheckedContinuation<T, any Error>? in
                     switch consume state {
                     case .buffer(let buffer):
                         state = .buffer(buffer)
+                        return nil
                     case .continuation(let continuation):
-                        continuation?.resume(throwing: CancellationError())
                         state = .buffer([])
+                        return continuation
                     }
                 }
+
+                // Resume the continuation outside the lock to avoid potential deadlocks
+                continuationToCancel?.resume(throwing: CancellationError())
             }
         }
 
         func makeAsyncIterator() -> Pool {
             self
         }
+
+        struct PoolError: Error {
+            let cause: Cause
+            var message: String {
+                switch self.cause {
+                case .nextCalledTwice:
+                    return "Concurrent invocations to next(). This is not allowed."
+                }
+            }
+
+            enum Cause {
+                case nextCalledTwice(CheckedContinuation<T, any Error>)
+            }
+        }
     }
 
     private struct LocalServerResponse: Sendable {

Tests/AWSLambdaRuntimeTests/PoolTests.swift

@@ -24,8 +24,8 @@ struct PoolTests {
         let pool = LambdaHTTPServer.Pool<String>()
 
         // Push values
-        await pool.push("first")
-        await pool.push("second")
+        pool.push("first")
+        pool.push("second")
 
         // Iterate and verify order
         var values = [String]()
@@ -53,7 +53,11 @@ struct PoolTests {
         task.cancel()
 
         // This should complete without receiving any values
-        try await task.value
+        do {
+            try await task.value
+        } catch is CancellationError {
+            // this might happen depending on the order on which the cancellation is handled
+        }
     }
 
     @Test
@@ -78,7 +82,7 @@ struct PoolTests {
         try await withThrowingTaskGroup(of: Void.self) { group in
             for i in 0..<iterations {
                 group.addTask {
-                    await pool.push(i)
+                    pool.push(i)
                 }
             }
             try await group.waitForAll()
@@ -110,7 +114,7 @@ struct PoolTests {
         try await Task.sleep(nanoseconds: 100_000_000)  // 0.1 seconds
 
         // Push a value
-        await pool.push(expectedValue)
+        pool.push(expectedValue)
 
         // Wait for consumer to complete
         try await consumer.value
@@ -140,7 +144,7 @@ struct PoolTests {
             for p in 0..<producerCount {
                 group.addTask {
                     for i in 0..<messagesPerProducer {
-                        await pool.push(p * messagesPerProducer + i)
+                        pool.push(p * messagesPerProducer + i)
                     }
                 }
             }
@@ -154,4 +158,32 @@ struct PoolTests {
         #expect(receivedValues.count == producerCount * messagesPerProducer)
         #expect(Set(receivedValues).count == producerCount * messagesPerProducer)
     }
+
+    @Test
+    @available(LambdaSwift 2.0, *)
+    func testConcurrentNext() async throws {
+        let pool = LambdaHTTPServer.Pool<String>()
+
+        // Create two tasks that will both wait for elements to be available
+        await #expect(throws: LambdaHTTPServer.Pool<Swift.String>.PoolError.self) {
+            try await withThrowingTaskGroup(of: Void.self) { group in
+
+                // one of the two task will throw a PoolError
+
+                group.addTask {
+                    for try await _ in pool {
+                    }
+                    Issue.record("Loop 1 should not complete")
+                }
+
+                group.addTask {
+                    for try await _ in pool {
+                    }
+                    Issue.record("Loop 2 should not complete")
+                }
+                try await group.waitForAll()
+            }
+        }
+    }
+
 }