reverse branch merge

Author: Sebastien Stormacq

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -393,50 +393,35 @@ internal struct LambdaHTTPServer {
             self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body))
 
             // wait for the lambda function to process the request
-            // 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)
+            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
                     }
-                }
-                // 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)"
+                } else {
+                    logger.error(
+                        "Received response for a different requestId",
+                        metadata: ["response requestId": "\(response.requestId ?? "")"]
                     )
-                )
-                try await self.sendResponse(response, outbound: outbound, logger: logger)
+                    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")
 
         // client uses incorrect HTTP method
         case (_, let url) where url.hasSuffix(self.invocationEndpoint):
@@ -621,18 +606,19 @@ internal struct LambdaHTTPServer {
                             if let nextAction = state.actionQueue.popFirst() {
                                 return (nextAction, continuation)
                             } else {
-                                state = .continuation(continuation)
-                                return nil
+                                // there is no continuation and no action waiting,
+                                // enqueue the continuation for later usage
+                                state.continuationQueue.append(continuation)
+                                return (nil, continuation)
                             }
-
-                        case .continuation(_):
-                            fatalError("\(self.poolName) : Concurrent invocations to next(). This is not allowed.")
                         }
                     }
 
+                    // there is no next action, ignore
                     guard let nextAction else { return }
 
-                    continuation.resume(returning: nextAction)
+                    // we have a next action and a next continuation, resume it
+                    nextContinuation.resume(returning: nextAction)
                 }
             } onCancel: {
                 self.lock.withLock { state in

Tests/AWSLambdaRuntimeTests/PoolTests.swift

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