swift-server/swift-aws-lambda-runtime
1
0
Fork 0
mirror of https://github.com/swift-server/swift-aws-lambda-runtime.git synced 2026-06-02 07:27:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixes for Local Lambda Server (#486)

Browse Source
This commit is contained in:
Fabian Fett
2025-02-27 16:08:39 +01:00
committed by GitHub
parent 64b417926d
commit 5de00c96c8
2 changed files with 175 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Package.swift
+2 -1
View File
@@ -19,6 +19,7 @@ let package = Package(
dependencies: [
.package(url: "https://github.com/apple/swift-nio.git", from: "2.81.0"),
.package(url: "https://github.com/apple/swift-log.git", from: "1.5.4"),
.package(url: "https://github.com/apple/swift-collections.git", from: "1.1.4"),
],
targets: [
.target(
@@ -31,10 +32,10 @@ let package = Package(
.target(
name: "AWSLambdaRuntimeCore",
dependencies: [
.product(name: "DequeModule", package: "swift-collections"),
.product(name: "Logging", package: "swift-log"),
.product(name: "NIOHTTP1", package: "swift-nio"),
.product(name: "NIOCore", package: "swift-nio"),
.product(name: "NIOConcurrencyHelpers", package: "swift-nio"),
.product(name: "NIOPosix", package: "swift-nio"),
]
),
Sources/AWSLambdaRuntimeCore/Lambda+LocalServer.swift
+173 -80
View File
@@ -13,9 +13,9 @@
//===----------------------------------------------------------------------===//
#if DEBUG
import DequeModule
import Dispatch
import Logging
import NIOConcurrencyHelpers
import NIOCore
import NIOHTTP1
import NIOPosix
@@ -47,24 +47,15 @@ extension Lambda {
/// - note: This API is designed strictly for local testing and is behind a DEBUG flag
static func withLocalServer(
invocationEndpoint: String? = nil,
_ body: @escaping () async throws -> Void
_ body: sending @escaping () async throws -> Void
) async throws {
var logger = Logger(label: "LocalServer")
logger.logLevel = Lambda.env("LOG_LEVEL").flatMap(Logger.Level.init) ?? .info
// launch the local server and wait for it to be started before running the body
try await withThrowingTaskGroup(of: Void.self) { group in
// this call will return when the server calls continuation.resume()
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
group.addTask {
do {
try await LambdaHttpServer(invocationEndpoint: invocationEndpoint).start(
continuation: continuation
)
} catch {
continuation.resume(throwing: error)
}
}
}
// now that server is started, run the Lambda function itself
try await LambdaHTTPServer.withLocalServer(
invocationEndpoint: invocationEndpoint,
logger: logger
) {
try await body()
}
}
@@ -84,34 +75,46 @@ extension Lambda {
/// 1. POST /invoke - the client posts the event to the lambda function
///
/// This server passes the data received from /invoke POST request to the lambda function (GET /next) and then forwards the response back to the client.
private struct LambdaHttpServer {
private let logger: Logger
private let group: EventLoopGroup
private let host: String
private let port: Int
private struct LambdaHTTPServer {
private let invocationEndpoint: String
private let invocationPool = Pool<LocalServerInvocation>()
private let responsePool = Pool<LocalServerResponse>()
init(invocationEndpoint: String?) {
var logger = Logger(label: "LocalServer")
logger.logLevel = Lambda.env("LOG_LEVEL").flatMap(Logger.Level.init) ?? .info
self.logger = logger
self.group = MultiThreadedEventLoopGroup.singleton
self.host = "127.0.0.1"
self.port = 7000
private init(
invocationEndpoint: String?
) {
self.invocationEndpoint = invocationEndpoint ?? "/invoke"
}
func start(continuation: CheckedContinuation<Void, any Error>) async throws {
let channel = try await ServerBootstrap(group: self.group)
private enum TaskResult<Result: Sendable>: Sendable {
case closureResult(Swift.Result<Result, any Error>)
case serverReturned(Swift.Result<Void, any Error>)
}
struct UnsafeTransferBox<Value>: @unchecked Sendable {
let value: Value
init(value: sending Value) {
self.value = value
}
}
static func withLocalServer<Result: Sendable>(
invocationEndpoint: String?,
host: String = "127.0.0.1",
port: Int = 7000,
eventLoopGroup: MultiThreadedEventLoopGroup = .singleton,
logger: Logger,
_ closure: sending @escaping () async throws -> Result
) async throws -> Result {
let channel = try await ServerBootstrap(group: eventLoopGroup)
.serverChannelOption(.backlog, value: 256)
.serverChannelOption(.socketOption(.so_reuseaddr), value: 1)
.childChannelOption(.maxMessagesPerRead, value: 1)
.bind(
host: self.host,
port: self.port
host: host,
port: port
) { channel in
channel.eventLoop.makeCompletedFuture {
@@ -129,8 +132,6 @@ private struct LambdaHttpServer {
}
}
// notify the caller that the server is started
continuation.resume()
logger.info(
"Server started and listening",
metadata: [
@@ -139,30 +140,87 @@ private struct LambdaHttpServer {
]
)
// We are handling each incoming connection in a separate child task. It is important
// to use a discarding task group here which automatically discards finished child tasks.
// A normal task group retains all child tasks and their outputs in memory until they are
// consumed by iterating the group or by exiting the group. Since, we are never consuming
// the results of the group we need the group to automatically discard them; otherwise, this
// would result in a memory leak over time.
try await withThrowingDiscardingTaskGroup { group in
try await channel.executeThenClose { inbound in
for try await connectionChannel in inbound {
let server = LambdaHTTPServer(invocationEndpoint: invocationEndpoint)
group.addTask {
logger.trace("Handling a new connection")
await self.handleConnection(channel: connectionChannel)
logger.trace("Done handling the connection")
// Sadly the Swift compiler does not understand that the passed in closure will only be
// invoked once. Because of this we need an unsafe transfer box here. Buuuh!
let closureBox = UnsafeTransferBox(value: closure)
let result = await withTaskGroup(of: TaskResult<Result>.self, returning: Swift.Result<Result, any Error>.self) {
group in
group.addTask {
let c = closureBox.value
do {
let result = try await c()
return .closureResult(.success(result))
} catch {
return .closureResult(.failure(error))
}
}
group.addTask {
do {
// We are handling each incoming connection in a separate child task. It is important
// to use a discarding task group here which automatically discards finished child tasks.
// A normal task group retains all child tasks and their outputs in memory until they are
// consumed by iterating the group or by exiting the group. Since, we are never consuming
// the results of the group we need the group to automatically discard them; otherwise, this
// would result in a memory leak over time.
try await withThrowingDiscardingTaskGroup { taskGroup in
try await channel.executeThenClose { inbound in
for try await connectionChannel in inbound {
taskGroup.addTask {
logger.trace("Handling a new connection")
await server.handleConnection(channel: connectionChannel, logger: logger)
logger.trace("Done handling the connection")
}
}
}
}
return .serverReturned(.success(()))
} catch {
return .serverReturned(.failure(error))
}
}
// Now that the local HTTP server and LambdaHandler tasks are started, wait for the
// first of the two that will terminate.
// When the first task terminates, cancel the group and collect the result of the
// second task.
// collect and return the result of the LambdaHandler
let serverOrHandlerResult1 = await group.next()!
group.cancelAll()
switch serverOrHandlerResult1 {
case .closureResult(let result):
return result
case .serverReturned(let result):
logger.error(
"Server shutdown before closure completed",
metadata: [
"error": "\(result.maybeError != nil ? "\(result.maybeError!)" : "none")"
]
)
switch await group.next()! {
case .closureResult(let result):
return result
case .serverReturned:
fatalError("Only one task is a server, and only one can return `serverReturned`")
}
}
}
logger.info("Server shutting down")
return try result.get()
}
/// This method handles individual TCP connections
private func handleConnection(
channel: NIOAsyncChannel<HTTPServerRequestPart, HTTPServerResponsePart>
channel: NIOAsyncChannel<HTTPServerRequestPart, HTTPServerResponsePart>,
logger: Logger
) async {
var requestHead: HTTPRequestHead!
@@ -186,12 +244,14 @@ private struct LambdaHttpServer {
// process the request
let response = try await self.processRequest(
head: requestHead,
body: requestBody
body: requestBody,
logger: logger
)
// send the responses
try await self.sendResponse(
response: response,
outbound: outbound
outbound: outbound,
logger: logger
)
requestHead = nil
@@ -214,15 +274,19 @@ private struct LambdaHttpServer {
/// - body: the HTTP request body
/// - Throws:
/// - Returns: the response to send back to the client or the Lambda function
private func processRequest(head: HTTPRequestHead, body: ByteBuffer?) async throws -> LocalServerResponse {
private func processRequest(
head: HTTPRequestHead,
body: ByteBuffer?,
logger: Logger
) async throws -> LocalServerResponse {
if let body {
self.logger.trace(
logger.trace(
"Processing request",
metadata: ["URI": "\(head.method) \(head.uri)", "Body": "\(String(buffer: body))"]
)
} else {
self.logger.trace("Processing request", metadata: ["URI": "\(head.method) \(head.uri)"])
logger.trace("Processing request", metadata: ["URI": "\(head.method) \(head.uri)"])
}
switch (head.method, head.uri) {
@@ -237,7 +301,9 @@ private struct LambdaHttpServer {
}
// we always accept the /invoke request and push them to the pool
let requestId = "\(DispatchTime.now().uptimeNanoseconds)"
logger.trace("/invoke received invocation", metadata: ["requestId": "\(requestId)"])
var logger = logger
logger[metadataKey: "requestID"] = "\(requestId)"
logger.trace("/invoke received invocation")
await self.invocationPool.push(LocalServerInvocation(requestId: requestId, request: body))
// wait for the lambda function to process the request
@@ -273,9 +339,9 @@ private struct LambdaHttpServer {
case (.GET, let url) where url.hasSuffix(Consts.getNextInvocationURLSuffix):
// pop the tasks from the queue
self.logger.trace("/next waiting for /invoke")
logger.trace("/next waiting for /invoke")
for try await invocation in self.invocationPool {
self.logger.trace("/next retrieved invocation", metadata: ["requestId": "\(invocation.requestId)"])
logger.trace("/next retrieved invocation", metadata: ["requestId": "\(invocation.requestId)"])
// this call also stores the invocation requestId into the response
return invocation.makeResponse(status: .accepted)
}
@@ -322,12 +388,13 @@ private struct LambdaHttpServer {
private func sendResponse(
response: LocalServerResponse,
outbound: NIOAsyncChannelOutboundWriter<HTTPServerResponsePart>
outbound: NIOAsyncChannelOutboundWriter<HTTPServerResponsePart>,
logger: Logger
) async throws {
var headers = HTTPHeaders(response.headers ?? [])
headers.add(name: "Content-Length", value: "\(response.body?.readableBytes ?? 0)")
self.logger.trace("Writing response", metadata: ["requestId": "\(response.requestId ?? "")"])
logger.trace("Writing response", metadata: ["requestId": "\(response.requestId ?? "")"])
try await outbound.write(
HTTPServerResponsePart.head(
HTTPResponseHead(
@@ -350,44 +417,59 @@ private struct LambdaHttpServer {
private final class Pool<T>: AsyncSequence, AsyncIteratorProtocol, Sendable where T: Sendable {
typealias Element = T
private let _buffer = Mutex<CircularBuffer<T>>(.init())
private let _continuation = Mutex<CheckedContinuation<T, any Error>?>(nil)
/// retrieve the first element from the buffer
public func popFirst() async -> T? {
self._buffer.withLock { $0.popFirst() }
enum State: ~Copyable {
case buffer(Deque<T>)
case continuation(CheckedContinuation<T, any Error>?)
}
private let lock = Mutex<State>(.buffer([]))
/// enqueue an element, or give it back immediately to the iterator if it is waiting for an element
public func push(_ invocation: T) async {
// if the iterator is waiting for an element, give it to it
// otherwise, enqueue the element
if let continuation = self._continuation.withLock({ $0 }) {
self._continuation.withLock { $0 = nil }
continuation.resume(returning: invocation)
} else {
self._buffer.withLock { $0.append(invocation) }
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)
return nil
}
}
maybeContinuation?.resume(returning: invocation)
}
func next() async throws -> T? {
// exit the async for loop if the task is cancelled
guard !Task.isCancelled else {
return nil
}
if let element = await self.popFirst() {
return element
} else {
// we can't return nil if there is nothing to dequeue otherwise the async for loop will stop
// wait for an element to be enqueued
return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<T, any Error>) in
// store the continuation for later, when an element is enqueued
self._continuation.withLock {
$0 = continuation
return 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
} else {
state = .continuation(continuation)
return nil
}
case .continuation:
fatalError("Concurrent invocations to next(). This is illegal.")
}
}
guard let nextAction else { return }
continuation.resume(returning: nextAction)
}
}
@@ -432,3 +514,14 @@ private struct LambdaHttpServer {
}
}
#endif
extension Result {
var maybeError: Failure? {
switch self {
case .success:
return nil
case .failure(let error):
return error
}
}
}