initial commit to support multi concurrency on Lambda Managed Instances

Sebastien Stormacq · Sebastien Stormacq · commit a4c04e93b739 · 2025-12-06T17:32:09.000-06:00
diff --git a/Sources/AWSLambdaRuntime/FoundationSupport/Lambda+JSON.swift b/Sources/AWSLambdaRuntime/FoundationSupport/Lambda+JSON.swift
@@ -137,7 +137,7 @@ extension LambdaRuntime {
         decoder: JSONDecoder = JSONDecoder(),
         encoder: JSONEncoder = JSONEncoder(),
         logger: Logger = Logger(label: "LambdaRuntime"),
-        body: sending @escaping (Event, LambdaContext) async throws -> Output
+        body: @Sendable @escaping (Event, LambdaContext) async throws -> Output
     )
     where
         Handler == LambdaCodableAdapter<
@@ -164,7 +164,7 @@ extension LambdaRuntime {
     public convenience init<Event: Decodable>(
         decoder: JSONDecoder = JSONDecoder(),
         logger: Logger = Logger(label: "LambdaRuntime"),
-        body: sending @escaping (Event, LambdaContext) async throws -> Void
+        body: @Sendable @escaping (Event, LambdaContext) async throws -> Void
     )
     where
         Handler == LambdaCodableAdapter<
diff --git a/Sources/AWSLambdaRuntime/Lambda+Codable.swift b/Sources/AWSLambdaRuntime/Lambda+Codable.swift
@@ -17,7 +17,7 @@ import NIOCore
 
 /// The protocol a decoder must conform to so that it can be used with ``LambdaCodableAdapter`` to decode incoming
 /// `ByteBuffer` events.
-public protocol LambdaEventDecoder {
+public protocol LambdaEventDecoder: Sendable {
     /// Decode the `ByteBuffer` representing the received event into the generic `Event` type
     /// the handler will receive.
     /// - Parameters:
@@ -29,7 +29,7 @@ public protocol LambdaEventDecoder {
 
 /// The protocol an encoder must conform to so that it can be used with ``LambdaCodableAdapter`` to encode the generic
 /// ``LambdaOutputEncoder/Output`` object into a `ByteBuffer`.
-public protocol LambdaOutputEncoder {
+public protocol LambdaOutputEncoder: Sendable {
     associatedtype Output
 
     /// Encode the generic type `Output` the handler has returned into a `ByteBuffer`.
diff --git a/Sources/AWSLambdaRuntime/LambdaHandlers.swift b/Sources/AWSLambdaRuntime/LambdaHandlers.swift
@@ -23,7 +23,7 @@ import NIOCore
 /// ``LambdaResponseStreamWriter/finish()`` or ``LambdaResponseStreamWriter/writeAndFinish(_:)``,
 /// the ``handle(_:responseWriter:context:)`` function is free to execute any background work.
 @available(LambdaSwift 2.0, *)
-public protocol StreamingLambdaHandler: _Lambda_SendableMetatype {
+public protocol StreamingLambdaHandler: Sendable, _Lambda_SendableMetatype {
     /// The handler function -- implement the business logic of the Lambda function here.
     /// - Parameters:
     ///   - event: The invocation's input data.
@@ -48,7 +48,7 @@ public protocol StreamingLambdaHandler: _Lambda_SendableMetatype {
 
 /// A writer object to write the Lambda response stream into. The HTTP response is started lazily.
 /// before the first call to ``write(_:)`` or ``writeAndFinish(_:)``.
-public protocol LambdaResponseStreamWriter {
+public protocol LambdaResponseStreamWriter: Sendable {
     /// Write a response part into the stream. Bytes written are streamed continually.
     /// - Parameter buffer: The buffer to write.
     /// - Parameter hasCustomHeaders: If `true`, the response will be sent with custom HTTP status code and headers.
@@ -69,7 +69,7 @@ public protocol LambdaResponseStreamWriter {
 /// - note: This handler protocol does not support response streaming because the output has to be encoded prior to it being sent, e.g. it is not possible to encode a partial/incomplete JSON string.
 /// This protocol also does not support the execution of background work after the response has been returned -- the ``LambdaWithBackgroundProcessingHandler`` protocol caters for such use-cases.
 @available(LambdaSwift 2.0, *)
-public protocol LambdaHandler {
+public protocol LambdaHandler: Sendable {
     /// Generic input type.
     /// The body of the request sent to Lambda will be decoded into this type for the handler to consume.
     associatedtype Event
@@ -92,7 +92,7 @@ public protocol LambdaHandler {
 /// ``LambdaWithBackgroundProcessingHandler/handle(_:outputWriter:context:)`` function is then
 /// free to implement any background work after the result has been sent to the AWS Lambda control plane.
 @available(LambdaSwift 2.0, *)
-public protocol LambdaWithBackgroundProcessingHandler {
+public protocol LambdaWithBackgroundProcessingHandler: Sendable {
     /// Generic input type.
     /// The body of the request sent to Lambda will be decoded into this type for the handler to consume.
     associatedtype Event
@@ -116,7 +116,7 @@ public protocol LambdaWithBackgroundProcessingHandler {
 /// Used with ``LambdaWithBackgroundProcessingHandler``.
 /// A mechanism to "return" an output from ``LambdaWithBackgroundProcessingHandler/handle(_:outputWriter:context:)`` without the function needing to
 /// have a return type and exit at that point. This allows for background work to be executed _after_ a response has been sent to the AWS Lambda response endpoint.
-public protocol LambdaResponseWriter<Output> {
+public protocol LambdaResponseWriter<Output>: Sendable {
     associatedtype Output
     /// Sends the generic ``LambdaResponseWriter/Output`` object (representing the computed result of the handler)
     /// to the AWS Lambda response endpoint.
@@ -157,18 +157,18 @@ public struct StreamingClosureHandler: StreamingLambdaHandler {
 /// A ``LambdaHandler`` conforming handler object that can be constructed with a closure.
 /// Allows for a handler to be defined in a clean manner, leveraging Swift's trailing closure syntax.
 @available(LambdaSwift 2.0, *)
-public struct ClosureHandler<Event: Decodable, Output>: LambdaHandler {
-    let body: (Event, LambdaContext) async throws -> Output
+public struct ClosureHandler<Event: Decodable & Sendable, Output>: LambdaHandler {
+    let body: @Sendable (Event, LambdaContext) async throws -> Output
 
     /// Initialize with a closure handler over generic `Input` and `Output` types.
     /// - Parameter body: The handler function written as a closure.
-    public init(body: sending @escaping (Event, LambdaContext) async throws -> Output) where Output: Encodable {
+    public init(body: @Sendable @escaping (Event, LambdaContext) async throws -> Output) where Output: Encodable {
         self.body = body
     }
 
     /// Initialize with a closure handler over a generic `Input` type, and a `Void` `Output`.
     /// - Parameter body: The handler function written as a closure.
-    public init(body: @escaping (Event, LambdaContext) async throws -> Void) where Output == Void {
+    public init(body: @Sendable @escaping (Event, LambdaContext) async throws -> Void) where Output == Void {
         self.body = body
     }
 
@@ -210,7 +210,7 @@ extension LambdaRuntime {
         encoder: sending Encoder,
         decoder: sending Decoder,
         logger: Logger = Logger(label: "LambdaRuntime"),
-        body: sending @escaping (Event, LambdaContext) async throws -> Output
+        body: @Sendable @escaping (Event, LambdaContext) async throws -> Output
     )
     where
         Handler == LambdaCodableAdapter<
@@ -240,7 +240,7 @@ extension LambdaRuntime {
     public convenience init<Event: Decodable, Decoder: LambdaEventDecoder>(
         decoder: sending Decoder,
         logger: Logger = Logger(label: "LambdaRuntime"),
-        body: sending @escaping (Event, LambdaContext) async throws -> Void
+        body: @Sendable @escaping (Event, LambdaContext) async throws -> Void
     )
     where
         Handler == LambdaCodableAdapter<
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntime.swift b/Sources/AWSLambdaRuntime/LambdaRuntime.swift
@@ -25,8 +25,7 @@ private let _isRunning = Atomic<Bool>(false)
 @available(LambdaSwift 2.0, *)
 public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLambdaHandler {
     @usableFromInline
-    /// we protect the handler behind a Mutex to ensure that we only ever have one copy of it
-    let handlerStorage: SendingStorage<Handler>
+    let handler: Handler
     @usableFromInline
     let logger: Logger
     @usableFromInline
@@ -37,7 +36,7 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
         eventLoop: EventLoop = Lambda.defaultEventLoop,
         logger: Logger = Logger(label: "LambdaRuntime")
     ) {
-        self.handlerStorage = SendingStorage(handler)
+        self.handler = handler
         self.eventLoop = eventLoop
 
         // by setting the log level here, we understand it can not be changed dynamically at runtime
@@ -74,45 +73,47 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
             _isRunning.store(false, ordering: .releasing)
         }
 
-        // The handler can be non-sendable, we want to ensure we only ever have one copy of it
-        let handler = try? self.handlerStorage.get()
-        guard let handler else {
-            throw LambdaRuntimeError(code: .handlerCanOnlyBeGetOnce)
-        }
-
         // are we running inside an AWS Lambda runtime environment ?
         // AWS_LAMBDA_RUNTIME_API is set when running on Lambda
         // https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html
         if let runtimeEndpoint = Lambda.env("AWS_LAMBDA_RUNTIME_API") {
 
-            let ipAndPort = runtimeEndpoint.split(separator: ":", maxSplits: 1)
-            let ip = String(ipAndPort[0])
-            guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeError(code: .invalidPort) }
-
-            do {
-                try await LambdaRuntimeClient.withRuntimeClient(
-                    configuration: .init(ip: ip, port: port),
+            // Get the max concurrency authorized by user when running on 
+            // Lambda Managed Instances
+            // This is not documented anywhere, except the NodeJS runtime
+            // https://github.com/aws/aws-lambda-nodejs-runtime-interface-client/blob/a4560c87426fa0a34756296a30d7add1388e575c/src/utils/env.ts#L34
+            // and
+            // https://github.com/aws/aws-lambda-nodejs-runtime-interface-client/blob/a4560c87426fa0a34756296a30d7add1388e575c/src/worker/ignition.ts#L12
+            let maxConcurrency = Int(Lambda.env("AWS_LAMBDA_MAX_CONCURRENCY") ?? "1") ?? 1
+
+            // when max concurrency is 1, do not pay the overhead of launching a Task
+            if maxConcurrency <= 1 {
+                self.logger.trace("Starting one Runtime Interface Client")
+                try await self.startRuntimeInterfaceClient(
+                    endpoint: runtimeEndpoint,
+                    handler: handler,
                     eventLoop: self.eventLoop,
                     logger: self.logger
-                ) { runtimeClient in
-                    try await Lambda.runLoop(
-                        runtimeClient: runtimeClient,
-                        handler: handler,
-                        logger: self.logger
-                    )
-                }
-            } catch {
-                // catch top level errors that have not been handled until now
-                // this avoids the runtime to crash and generate a backtrace
-                if let error = error as? LambdaRuntimeError,
-                    error.code != .connectionToControlPlaneLost
-                {
-                    // if the error is a LambdaRuntimeError but not a connection error,
-                    // we rethrow it to preserve existing behaviour
-                    self.logger.error("LambdaRuntime.run() failed with error", metadata: ["error": "\(error)"])
-                    throw error
-                } else {
-                    self.logger.trace("LambdaRuntime.run() connection lost")
+                )
+            } else {
+                try await withThrowingTaskGroup(of: Void.self) { group in
+
+                    self.logger.trace("Starting \(maxConcurrency) Runtime Interface Clients")
+                    for i in 0..<maxConcurrency {
+
+                        group.addTask {
+                            var logger = self.logger 
+                            logger[metadataKey: "RIC"] = "\(i)"
+                            try await self.startRuntimeInterfaceClient(
+                                endpoint: runtimeEndpoint,
+                                handler: self.handler,
+                                eventLoop: self.eventLoop,
+                                logger: logger
+                            )
+                        }
+                    }
+                    // Wait for all tasks to complete and propagate any errors
+                    try await group.waitForAll()                    
                 }
             }
 
@@ -141,7 +142,7 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
                 ) { runtimeClient in
                     try await Lambda.runLoop(
                         runtimeClient: runtimeClient,
-                        handler: handler,
+                        handler: self.handler,
                         logger: self.logger
                     )
                 }
@@ -152,4 +153,43 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
             #endif
         }
     }
+
+    internal func startRuntimeInterfaceClient(
+        endpoint: String,
+        handler: Handler,
+        eventLoop: EventLoop,
+        logger: Logger
+    ) async throws {
+
+        let ipAndPort = endpoint.split(separator: ":", maxSplits: 1)
+        let ip = String(ipAndPort[0])
+        guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeError(code: .invalidPort) }
+
+        do {
+            try await LambdaRuntimeClient.withRuntimeClient(
+                configuration: .init(ip: ip, port: port),
+                eventLoop: eventLoop,
+                logger: logger
+            ) { runtimeClient in
+                try await Lambda.runLoop(
+                    runtimeClient: runtimeClient,
+                    handler: handler,
+                    logger: logger
+                )
+            }
+        } catch {
+            // catch top level errors that have not been handled until now
+            // this avoids the runtime to crash and generate a backtrace
+            if let error = error as? LambdaRuntimeError,
+                error.code != .connectionToControlPlaneLost
+            {
+                // if the error is a LambdaRuntimeError but not a connection error,
+                // we rethrow it to preserve existing behaviour
+                logger.error("LambdaRuntime.run() failed with error", metadata: ["error": "\(error)"])
+                throw error
+            } else {
+                logger.trace("LambdaRuntime.run() connection lost")
+            }
+        }
+    }
 }
diff --git a/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift b/Sources/AWSLambdaRuntime/LambdaRuntimeError.swift
@@ -34,7 +34,6 @@ package struct LambdaRuntimeError: Error {
 
         case missingLambdaRuntimeAPIEnvironmentVariable
         case runtimeCanOnlyBeStartedOnce
-        case handlerCanOnlyBeGetOnce
         case invalidPort
     }
 
diff --git a/Sources/AWSLambdaRuntime/Utils.swift b/Sources/AWSLambdaRuntime/Utils.swift
@@ -107,33 +107,33 @@ extension AmazonHeaders {
     }
 }
 
-/// Temporary storage for value being sent from one isolation domain to another
-// use NIOLockedValueBox instead of Mutex to avoid compiler crashes on 6.0
-// see https://github.com/swiftlang/swift/issues/78048
-@usableFromInline
-struct SendingStorage<Value>: ~Copyable, @unchecked Sendable {
-    @usableFromInline
-    struct ValueAlreadySentError: Error {
-        @usableFromInline
-        init() {}
-    }
+// /// Temporary storage for value being sent from one isolation domain to another
+// // use NIOLockedValueBox instead of Mutex to avoid compiler crashes on 6.0
+// // see https://github.com/swiftlang/swift/issues/78048
+// @usableFromInline
+// struct SendingStorage<Value>: ~Copyable, @unchecked Sendable {
+//     @usableFromInline
+//     struct ValueAlreadySentError: Error {
+//         @usableFromInline
+//         init() {}
+//     }
 
-    @usableFromInline
-    // let storage: Mutex<Value?>
-    let storage: NIOLockedValueBox<Value?>
+//     @usableFromInline
+//     // let storage: Mutex<Value?>
+//     let storage: NIOLockedValueBox<Value?>
 
-    @inlinable
-    init(_ value: sending Value) {
-        self.storage = .init(value)
-    }
+//     @inlinable
+//     init(_ value: sending Value) {
+//         self.storage = .init(value)
+//     }
 
-    @inlinable
-    func get() throws -> Value {
-        // try self.storage.withLock {
-        try self.storage.withLockedValue {
-            guard let value = $0 else { throw ValueAlreadySentError() }
-            $0 = nil
-            return value
-        }
-    }
-}
+//     @inlinable
+//     func get() throws -> Value {
+//         // try self.storage.withLock {
+//         try self.storage.withLockedValue {
+//             guard let value = $0 else { throw ValueAlreadySentError() }
+//             $0 = nil
+//             return value
+//         }
+//     }
+// }

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,6 @@ package struct LambdaRuntimeError: Error {`
`34`	`34`
`35`	`35`	`case missingLambdaRuntimeAPIEnvironmentVariable`
`36`	`36`	`case runtimeCanOnlyBeStartedOnce`
`37`		`- case handlerCanOnlyBeGetOnce`
`38`	`37`	`case invalidPort`
`39`	`38`	`}`
`40`	`39`