WIP LocalServer : allow concurrent invocations of POST /invoke

Author: Sebastien Stormacq

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -390,19 +390,7 @@ internal struct LambdaHTTPServer {
             logger[metadataKey: "requestId"] = "\(requestId)"
 
             logger.trace("/invoke received invocation, pushing it to the pool and wait for a lambda response")
-            // detect concurrent invocations of POST and gently decline the requests while we're processing one.
-            if await !self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body)) {
-                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)
-                return
-            }
+            await self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body))
 
             // wait for the lambda function to process the request
             for try await response in self.responsePool {
@@ -571,39 +559,26 @@ internal struct LambdaHTTPServer {
 
         typealias Element = T
 
-        enum State: ~Copyable {
-            case buffer(Deque<T>)
-            case continuation(CheckedContinuation<T, any Error>?)
+        struct State: ~Copyable {
+            var actions : Deque<T>
+            var continuationQueue: Deque<CheckedContinuation<T, any Error>>
         }
-
-        private let lock = Mutex<State>(.buffer([]))
+        private let lock = Mutex<State>(State(actions: [], continuationQueue: []))
 
         /// enqueue an element, or give it back immediately to the iterator if it is waiting for an element
-        /// Returns true when we receive a element and the pool was in "waiting for continuation" state, false otherwise
-        @discardableResult
-        public func push(_ invocation: T) async -> Bool {
+        public func push(_ action: T) async {
 
-            // if the iterator is waiting for an element on `next()``, give it to it
+            // if the iterator is waiting for an element on `next()`, give it to it (resume the continuation)
             // otherwise, enqueue the element
             let maybeContinuation = self.lock.withLock { state -> CheckedContinuation<T, any Error>? in
-                switch consume state {
-                case .continuation(let continuation):
-                    state = .buffer([])
-                    return continuation
-
-                case .buffer(var buffer):
-                    buffer.append(invocation)
-                    state = .buffer(buffer)
+                if let nextContinuation = state.continuationQueue.popFirst() {
+                    return nextContinuation
+                } else {
+                    state.actions.append(action)
                     return nil
                 }
             }
-
-            if let maybeContinuation {
-                maybeContinuation.resume(returning: invocation)
-                return true
-            } else {
-                return false
-            }
+            maybeContinuation?.resume(returning: action)
         }
 
         func next() async throws -> T? {
@@ -614,34 +589,41 @@ internal struct LambdaHTTPServer {
 
             return try await withTaskCancellationHandler {
                 try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<T, any Error>) in
-                    let nextAction = self.lock.withLock { state -> T? in
-                        switch consume state {
-                        case .buffer(var buffer):
-                            if let first = buffer.popFirst() {
-                                state = .buffer(buffer)
-                                return first
+                    let (nextAction, nextContinuation) = self.lock.withLock { state -> (T?, CheckedContinuation<T, any Error>) in
+                        if let nextContinuation = state.continuationQueue.popFirst() {
+                            if let nextAction = state.actions.popFirst() {
+                                return (nextAction, nextContinuation)
                             } else {
-                                state = .continuation(continuation)
-                                return nil
+                                // next is called when there is no action to process, but there is a continuation waiting
+                                // we need to enqueue the continuation to be resumed when the next action is available
+                                state.continuationQueue.append(nextContinuation)
+                                state.continuationQueue.append(continuation)
+                                return (nil, continuation)
+                            }
+                        } else {
+                            // we have no waiting continuation, let's consume the one we just received
+                            if let nextAction = state.actions.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)
                             }
-
-                        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
-                    switch consume state {
-                    case .buffer(let buffer):
-                        state = .buffer(buffer)
-                    case .continuation(let continuation):
-                        continuation?.resume(throwing: CancellationError())
-                        state = .buffer([])
+                    // when a cancellation is received, resume all pending continuations
+                    for continuation in state.continuationQueue {
+                        continuation.resume(throwing: CancellationError())
                     }
                 }
             }