Merge branch 'sebsto/fix_584' into sebsto/multiple_continuations

Author: Sebastien Stormacq

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -393,35 +393,50 @@ internal struct LambdaHTTPServer {
             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 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)
                 }
+                // 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):
@@ -606,19 +621,18 @@ internal struct LambdaHTTPServer {
                             if let nextAction = state.actionQueue.popFirst() {
                                 return (nextAction, continuation)
                             } else {
-                                // there is no continuation and no action waiting,
-                                // enqueue the continuation for later usage
-                                state.continuationQueue.append(continuation)
-                                return (nil, continuation)
+                                state = .continuation(continuation)
+                                return nil
                             }
+
+                        case .continuation(_):
+                            fatalError("\(self.poolName) : Concurrent invocations to next(). This is not allowed.")
                         }
                     }
 
-                    // there is no next action, ignore
                     guard let nextAction else { return }
 
-                    // we have a next action and a next continuation, resume it
-                    nextContinuation.resume(returning: nextAction)
+                    continuation.resume(returning: nextAction)
                 }
             } onCancel: {
                 self.lock.withLock { state in

Tests/AWSLambdaRuntimeTests/PoolTests.swift

@@ -158,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()
+            }
+        }
+    }
+
 }