Skip to content

Commit 85b938e

Browse files
aryan-25aryan-25
authored
Add new LambdaRuntime (#353)
1 parent 06763b8 commit 85b938e

File tree

6 files changed

+227
-25
lines changed

6 files changed

+227
-25
lines changed

‎Sources/AWSLambdaRuntime/Lambda+Codable.swift

+53
Original file line numberDiff line numberDiff line change
@@ -65,3 +65,56 @@ extension LambdaCodableAdapter {
6565
)
6666
}
6767
}
68+
69+
extension NewLambdaRuntime {
70+
/// Initialize an instance with a ``NewLambdaHandler`` defined in the form of a closure **with a non-`Void` return type**.
71+
/// - Parameter body: The handler in the form of a closure.
72+
/// - Parameter encoder: The encoder object that will be used to encode the generic ``Output`` into a ``ByteBuffer``. ``JSONEncoder()`` used as default.
73+
/// - Parameter decoder: The decoder object that will be used to decode the incoming ``ByteBuffer`` event into the generic ``Event`` type. ``JSONDecoder()`` used as default.
74+
package convenience init<Event: Decodable, Output>(
75+
body: @escaping (Event, NewLambdaContext) async throws -> Output,
76+
encoder: JSONEncoder = JSONEncoder(),
77+
decoder: JSONDecoder = JSONDecoder()
78+
)
79+
where
80+
Handler == LambdaCodableAdapter<
81+
LambdaHandlerAdapter<Event, Output, ClosureHandler<Event, Output>>,
82+
Event,
83+
Output,
84+
JSONDecoder,
85+
LambdaJSONOutputEncoder<Output>
86+
>
87+
{
88+
let handler = LambdaCodableAdapter(
89+
encoder: encoder,
90+
decoder: decoder,
91+
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
92+
)
93+
94+
self.init(handler: handler)
95+
}
96+
97+
/// Initialize an instance with a ``NewLambdaHandler`` defined in the form of a closure **with a `Void` return type**.
98+
/// - Parameter body: The handler in the form of a closure.
99+
/// - Parameter decoder: The decoder object that will be used to decode the incoming ``ByteBuffer`` event into the generic ``Event`` type. ``JSONDecoder()`` used as default.
100+
package convenience init<Event: Decodable>(
101+
body: @escaping (Event, NewLambdaContext) async throws -> Void,
102+
decoder: JSONDecoder = JSONDecoder()
103+
)
104+
where
105+
Handler == LambdaCodableAdapter<
106+
LambdaHandlerAdapter<Event, Void, ClosureHandler<Event, Void>>,
107+
Event,
108+
Void,
109+
JSONDecoder,
110+
VoidEncoder
111+
>
112+
{
113+
let handler = LambdaCodableAdapter(
114+
decoder: decoder,
115+
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
116+
)
117+
118+
self.init(handler: handler)
119+
}
120+
}

‎Sources/AWSLambdaRuntimeCore/NewLambda.swift

+4
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import Dispatch
1616
import Logging
17+
import NIOCore
1718

1819
extension Lambda {
1920
package static func runLoop<RuntimeClient: LambdaRuntimeClientProtocol, Handler>(
@@ -44,4 +45,7 @@ extension Lambda {
4445
}
4546
}
4647
}
48+
49+
/// The default EventLoop the Lambda is scheduled on.
50+
package static var defaultEventLoop: any EventLoop = NIOSingletons.posixEventLoopGroup.next()
4751
}

‎Sources/AWSLambdaRuntimeCore/NewLambdaHandlers.swift

+67
Original file line numberDiff line numberDiff line change
@@ -171,3 +171,70 @@ package struct ClosureHandler<Event: Decodable, Output>: NewLambdaHandler {
171171
try await self.body(event, context)
172172
}
173173
}
174+
175+
extension NewLambdaRuntime {
176+
/// Initialize an instance with a ``StreamingLambdaHandler`` in the form of a closure.
177+
/// - Parameter body: The handler in the form of a closure.
178+
package convenience init(
179+
body: @Sendable @escaping (ByteBuffer, LambdaResponseStreamWriter, NewLambdaContext) async throws -> Void
180+
) where Handler == StreamingClosureHandler {
181+
self.init(handler: StreamingClosureHandler(body: body))
182+
}
183+
184+
/// Initialize an instance with a ``NewLambdaHandler`` defined in the form of a closure **with a non-`Void` return type**, an encoder, and a decoder.
185+
/// - Parameter body: The handler in the form of a closure.
186+
/// - Parameter encoder: The encoder object that will be used to encode the generic ``Output`` into a ``ByteBuffer``.
187+
/// - Parameter decoder: The decoder object that will be used to decode the incoming ``ByteBuffer`` event into the generic ``Event`` type.
188+
package convenience init<
189+
Event: Decodable,
190+
Output: Encodable,
191+
Encoder: LambdaOutputEncoder,
192+
Decoder: LambdaEventDecoder
193+
>(
194+
encoder: Encoder,
195+
decoder: Decoder,
196+
body: @escaping (Event, NewLambdaContext) async throws -> Output
197+
)
198+
where
199+
Handler == LambdaCodableAdapter<
200+
LambdaHandlerAdapter<Event, Output, ClosureHandler<Event, Output>>,
201+
Event,
202+
Output,
203+
Decoder,
204+
Encoder
205+
>
206+
{
207+
let handler = LambdaCodableAdapter(
208+
encoder: encoder,
209+
decoder: decoder,
210+
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
211+
)
212+
213+
self.init(handler: handler)
214+
}
215+
216+
/// Initialize an instance with a ``NewLambdaHandler`` defined in the form of a closure **with a `Void` return type**, an encoder, and a decoder.
217+
/// - Parameter body: The handler in the form of a closure.
218+
/// - Parameter encoder: The encoder object that will be used to encode the generic ``Output`` into a ``ByteBuffer``.
219+
/// - Parameter decoder: The decoder object that will be used to decode the incoming ``ByteBuffer`` event into the generic ``Event`` type.
220+
package convenience init<Event: Decodable, Decoder: LambdaEventDecoder>(
221+
decoder: Decoder,
222+
body: @escaping (Event, NewLambdaContext) async throws -> Void
223+
)
224+
where
225+
Handler == LambdaCodableAdapter<
226+
LambdaHandlerAdapter<Event, Void, ClosureHandler<Event, Void>>,
227+
Event,
228+
Void,
229+
Decoder,
230+
VoidEncoder
231+
>
232+
{
233+
let handler = LambdaCodableAdapter(
234+
decoder: decoder,
235+
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
236+
)
237+
238+
self.init(handler: handler)
239+
}
240+
}

‎Sources/AWSLambdaRuntimeCore/NewLambdaRuntime.swift

+70
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the SwiftAWSLambdaRuntime open source project
4+
//
5+
// Copyright (c) 2024 Apple Inc. and the SwiftAWSLambdaRuntime project authors
6+
// Licensed under Apache License v2.0
7+
//
8+
// See LICENSE.txt for license information
9+
// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
10+
//
11+
// SPDX-License-Identifier: Apache-2.0
12+
//
13+
//===----------------------------------------------------------------------===//
14+
15+
import Foundation
16+
import Logging
17+
import NIOCore
18+
import NIOConcurrencyHelpers
19+
20+
// We need `@unchecked` Sendable here, as `NIOLockedValueBox` does not understand `sending` today.
21+
// We don't want to use `NIOLockedValueBox` here anyway. We would love to use Mutex here, but this
22+
// sadly crashes the compiler today.
23+
package final class NewLambdaRuntime<Handler>: @unchecked Sendable where Handler: StreamingLambdaHandler {
24+
// TODO: We want to change this to Mutex as soon as this doesn't crash the Swift compiler on Linux anymore
25+
let handlerMutex: NIOLockedValueBox<Optional<Handler>>
26+
let logger: Logger
27+
let eventLoop: EventLoop
28+
29+
package init(
30+
handler: sending Handler,
31+
eventLoop: EventLoop = Lambda.defaultEventLoop,
32+
logger: Logger = Logger(label: "LambdaRuntime")
33+
) {
34+
self.handlerMutex = NIOLockedValueBox(handler)
35+
self.eventLoop = eventLoop
36+
self.logger = logger
37+
}
38+
39+
package func run() async throws {
40+
guard let runtimeEndpoint = Lambda.env("AWS_LAMBDA_RUNTIME_API") else {
41+
throw NewLambdaRuntimeError(code: .missingLambdaRuntimeAPIEnvironmentVariable)
42+
}
43+
44+
let ipAndPort = runtimeEndpoint.split(separator: ":", maxSplits: 1)
45+
let ip = String(ipAndPort[0])
46+
guard let port = Int(ipAndPort[1]) else { throw NewLambdaRuntimeError(code: .invalidPort) }
47+
48+
let handler = self.handlerMutex.withLockedValue { handler in
49+
let result = handler
50+
handler = nil
51+
return result
52+
}
53+
54+
guard let handler else {
55+
throw NewLambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
56+
}
57+
58+
try await NewLambdaRuntimeClient.withRuntimeClient(
59+
configuration: .init(ip: ip, port: port),
60+
eventLoop: self.eventLoop,
61+
logger: self.logger
62+
) { runtimeClient in
63+
try await Lambda.runLoop(
64+
runtimeClient: runtimeClient,
65+
handler: handler,
66+
logger: self.logger
67+
)
68+
}
69+
}
70+
}

‎Sources/AWSLambdaRuntimeCore/NewLambdaRuntimeClient.swift

+30-25
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ final actor NewLambdaRuntimeClient: LambdaRuntimeClientProtocol {
303303
NIOHTTPClientResponseAggregator(maxContentLength: 6 * 1024 * 1024)
304304
)
305305
try channel.pipeline.syncOperations.addHandler(
306-
LambdaChannelHandler(delegate: self, logger: self.logger)
306+
LambdaChannelHandler(delegate: self, logger: self.logger, configuration: self.configuration)
307307
)
308308
return channel.eventLoop.makeSucceededFuture(())
309309
} catch {
@@ -433,10 +433,33 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
433433
private var reusableErrorBuffer: ByteBuffer?
434434
private let logger: Logger
435435
private let delegate: Delegate
436+
private let configuration: NewLambdaRuntimeClient.Configuration
436437

437-
init(delegate: Delegate, logger: Logger) {
438+
/// These are the default headers that must be sent along an invocation
439+
let defaultHeaders: HTTPHeaders
440+
/// These headers must be sent along an invocation or initialization error report
441+
let errorHeaders: HTTPHeaders
442+
/// These headers must be sent when streaming a response
443+
let streamingHeaders: HTTPHeaders
444+
445+
init(delegate: Delegate, logger: Logger, configuration: NewLambdaRuntimeClient.Configuration) {
438446
self.delegate = delegate
439447
self.logger = logger
448+
self.configuration = configuration
449+
self.defaultHeaders = [
450+
"host": "\(self.configuration.ip):\(self.configuration.port)",
451+
"user-agent": "Swift-Lambda/Unknown",
452+
]
453+
self.errorHeaders = [
454+
"host": "\(self.configuration.ip):\(self.configuration.port)",
455+
"user-agent": "Swift-Lambda/Unknown",
456+
"lambda-runtime-function-error-type": "Unhandled",
457+
]
458+
self.streamingHeaders = [
459+
"host": "\(self.configuration.ip):\(self.configuration.port)",
460+
"user-agent": "Swift-Lambda/Unknown",
461+
"transfer-encoding": "chunked",
462+
]
440463
}
441464

442465
func nextInvocation(isolation: isolated (any Actor)? = #isolation) async throws -> Invocation {
@@ -578,7 +601,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
578601
version: .http1_1,
579602
method: .POST,
580603
uri: url,
581-
headers: NewLambdaRuntimeClient.streamingHeaders
604+
headers: self.streamingHeaders
582605
)
583606

584607
context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
@@ -604,11 +627,12 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
604627
let headers: HTTPHeaders =
605628
if byteBuffer?.readableBytes ?? 0 < 6_000_000 {
606629
[
630+
"host": "\(self.configuration.ip):\(self.configuration.port)",
607631
"user-agent": "Swift-Lambda/Unknown",
608632
"content-length": "\(byteBuffer?.readableBytes ?? 0)",
609633
]
610634
} else {
611-
NewLambdaRuntimeClient.streamingHeaders
635+
self.streamingHeaders
612636
}
613637

614638
let httpRequest = HTTPRequestHead(
@@ -634,7 +658,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
634658
version: .http1_1,
635659
method: .GET,
636660
uri: self.nextInvocationPath,
637-
headers: NewLambdaRuntimeClient.defaultHeaders
661+
headers: self.defaultHeaders
638662
)
639663

640664
context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
@@ -650,7 +674,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
650674
version: .http1_1,
651675
method: .POST,
652676
uri: url,
653-
headers: NewLambdaRuntimeClient.errorHeaders
677+
headers: self.errorHeaders
654678
)
655679

656680
if self.reusableErrorBuffer == nil {
@@ -797,22 +821,3 @@ extension LambdaChannelHandler: ChannelInboundHandler {
797821
context.fireChannelInactive()
798822
}
799823
}
800-
801-
extension NewLambdaRuntimeClient {
802-
static let defaultHeaders: HTTPHeaders = [
803-
"user-agent": "Swift-Lambda/Unknown"
804-
]
805-
806-
/// These headers must be sent along an invocation or initialization error report
807-
static let errorHeaders: HTTPHeaders = [
808-
"user-agent": "Swift-Lambda/Unknown",
809-
"lambda-runtime-function-error-type": "Unhandled",
810-
]
811-
812-
/// These headers must be sent along an invocation or initialization error report
813-
static let streamingHeaders: HTTPHeaders = [
814-
"user-agent": "Swift-Lambda/Unknown",
815-
"transfer-encoding": "streaming",
816-
]
817-
818-
}

‎Sources/AWSLambdaRuntimeCore/NewLambdaRuntimeError.swift

+3
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ package struct NewLambdaRuntimeError: Error {
2929
case nextInvocationMissingHeaderDeadline
3030
case nextInvocationMissingHeaderInvokeFuctionARN
3131

32+
case missingLambdaRuntimeAPIEnvironmentVariable
33+
case runtimeCanOnlyBeStartedOnce
34+
case invalidPort
3235
}
3336

3437
package init(code: Code, underlying: (any Error)? = nil) {

0 commit comments

Comments
 (0)