diff --git a/Package.swift b/Package.swift
index 0e5823f6..9c0bfedd 100644
--- a/Package.swift
+++ b/Package.swift
@@ -17,6 +17,7 @@ let package = Package(
dependencies: [
.package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.17.0")),
.package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.0.0")),
+ .package(name: "swift-context", url: "https://github.com/ktoso/gsoc-swift-baggage-context.git", .branch("simple-is-good-proposal")), // TODO: use main once merged
.package(url: "https://github.com/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
],
targets: [
@@ -27,6 +28,7 @@ let package = Package(
]),
.target(name: "AWSLambdaRuntimeCore", dependencies: [
.product(name: "Logging", package: "swift-log"),
+ .product(name: "BaggageContext", package: "swift-context"),
.product(name: "Backtrace", package: "swift-backtrace"),
.product(name: "NIOHTTP1", package: "swift-nio"),
]),
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
index ab30dd7b..738abe90 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
@@ -12,6 +12,7 @@
//
//===----------------------------------------------------------------------===//
+import BaggageContext
import Dispatch
import Logging
import NIO
@@ -49,40 +50,132 @@ extension Lambda {
extension Lambda {
/// Lambda runtime context.
/// The Lambda runtime generates and passes the `Context` to the Lambda handler as an argument.
- public final class Context: CustomDebugStringConvertible {
+ public struct Context: BaggageContext.Context, CustomDebugStringConvertible {
+ /// Used to store all contents of the context and implement CoW semantics for it.
+ private var storage: Storage
+
+ final class Storage {
+ var baggage: Baggage
+
+ let invokedFunctionARN: String
+ let deadline: DispatchWallTime
+ let cognitoIdentity: String?
+ let clientContext: String?
+
+ // Implementation note: This logger is the "user provided logger" that we will log to when `context.logger` is used.
+ // It must be updated with the latest metadata whenever the `baggage` changes.
+ var _logger: Logger
+
+ let eventLoop: EventLoop
+ let allocator: ByteBufferAllocator
+
+ init(
+ baggage: Baggage,
+ invokedFunctionARN: String,
+ deadline: DispatchWallTime,
+ cognitoIdentity: String?,
+ clientContext: String?,
+ logger: Logger,
+ eventLoop: EventLoop,
+ allocator: ByteBufferAllocator
+ ) {
+ self.baggage = baggage
+ self.invokedFunctionARN = invokedFunctionARN
+ self.deadline = deadline
+ self.cognitoIdentity = cognitoIdentity
+ self.clientContext = clientContext
+ self._logger = logger
+ self.eventLoop = eventLoop
+ self.allocator = allocator
+ }
+ }
+
+ /// Contains contextual metadata such as request and trace identifiers, along with other information which may
+ /// be carried throughout asynchronous and cross-node boundaries (e.g. through HTTPClient calls).
+ public var baggage: Baggage {
+ get {
+ self.storage.baggage
+ }
+ set {
+ if isKnownUniquelyReferenced(&self.storage) {
+ self.storage._logger.updateMetadata(previous: self.storage.baggage, latest: newValue)
+ self.storage.baggage = newValue
+ } else {
+ var logger = self.storage._logger
+ logger.updateMetadata(previous: self.storage.baggage, latest: newValue)
+ self.storage = Storage(
+ baggage: newValue,
+ invokedFunctionARN: self.storage.invokedFunctionARN,
+ deadline: self.storage.deadline,
+ cognitoIdentity: self.storage.cognitoIdentity,
+ clientContext: self.storage.clientContext,
+ logger: self.storage._logger,
+ eventLoop: self.storage.eventLoop,
+ allocator: self.storage.allocator
+ )
+ }
+ }
+ }
+
/// The request ID, which identifies the request that triggered the function invocation.
- public let requestID: String
+ public var requestID: String {
+ self.storage.baggage.lambdaRequestID
+ }
/// The AWS X-Ray tracing header.
- public let traceID: String
+ public var traceID: String {
+ self.storage.baggage.lambdaTraceID
+ }
/// The ARN of the Lambda function, version, or alias that's specified in the invocation.
- public let invokedFunctionARN: String
+ public var invokedFunctionARN: String {
+ self.storage.invokedFunctionARN
+ }
/// The timestamp that the function times out
- public let deadline: DispatchWallTime
+ public var deadline: DispatchWallTime {
+ self.storage.deadline
+ }
/// For invocations from the AWS Mobile SDK, data about the Amazon Cognito identity provider.
- public let cognitoIdentity: String?
+ public var cognitoIdentity: String? {
+ self.storage.cognitoIdentity
+ }
/// For invocations from the AWS Mobile SDK, data about the client application and device.
- public let clientContext: String?
+ public var clientContext: String? {
+ self.storage.clientContext
+ }
- /// `Logger` to log with
+ /// `Logger` to log with, it is automatically populated with `baggage` information (such as `traceID` and `requestID`).
///
- /// - note: The `LogLevel` can be configured using the `LOG_LEVEL` environment variable.
- public let logger: Logger
+ /// - note: The default `Logger.LogLevel` can be configured using the `LOG_LEVEL` environment variable.
+ public var logger: Logger {
+ get {
+ self.storage._logger
+ }
+ set {
+ if isKnownUniquelyReferenced(&self.storage) {
+ self.storage._logger = newValue
+ self.storage._logger.updateMetadata(previous: .topLevel, latest: self.storage.baggage)
+ }
+ }
+ }
/// The `EventLoop` the Lambda is executed on. Use this to schedule work with.
/// This is useful when implementing the `EventLoopLambdaHandler` protocol.
///
/// - note: The `EventLoop` is shared with the Lambda runtime engine and should be handled with extra care.
/// Most importantly the `EventLoop` must never be blocked.
- public let eventLoop: EventLoop
+ public var eventLoop: EventLoop {
+ self.storage.eventLoop
+ }
/// `ByteBufferAllocator` to allocate `ByteBuffer`
/// This is useful when implementing `EventLoopLambdaHandler`
- public let allocator: ByteBufferAllocator
+ public var allocator: ByteBufferAllocator {
+ self.storage.allocator
+ }
internal init(requestID: String,
traceID: String,
@@ -93,20 +186,20 @@ extension Lambda {
logger: Logger,
eventLoop: EventLoop,
allocator: ByteBufferAllocator) {
- self.requestID = requestID
- self.traceID = traceID
- self.invokedFunctionARN = invokedFunctionARN
- self.cognitoIdentity = cognitoIdentity
- self.clientContext = clientContext
- self.deadline = deadline
- // utility
- self.eventLoop = eventLoop
- self.allocator = allocator
- // mutate logger with context
- var logger = logger
- logger[metadataKey: "awsRequestID"] = .string(requestID)
- logger[metadataKey: "awsTraceID"] = .string(traceID)
- self.logger = logger
+ var baggage = Baggage.topLevel
+ baggage.lambdaRequestID = requestID
+ baggage.lambdaTraceID = traceID
+ self.storage = Storage(
+ baggage: baggage,
+ invokedFunctionARN: invokedFunctionARN,
+ deadline: deadline,
+ cognitoIdentity: cognitoIdentity,
+ clientContext: clientContext,
+ // utility
+ logger: logger,
+ eventLoop: eventLoop,
+ allocator: allocator
+ )
}
public func getRemainingTime() -> TimeAmount {
@@ -146,3 +239,42 @@ extension Lambda {
}
}
}
+
+// MARK: - Baggage Items
+
+extension Baggage {
+
+ // MARK: - Baggage: RequestID
+
+ enum LambdaRequestIDKey: Key {
+ typealias Value = String
+ static var name: String? { AmazonHeaders.requestID }
+ }
+
+ /// The request ID, which identifies the request that triggered the function invocation.
+ public internal(set) var lambdaRequestID: String {
+ get {
+ self[LambdaRequestIDKey.self]! // !-safe, the runtime guarantees to always set an identifier, even in testing
+ }
+ set {
+ self[LambdaRequestIDKey.self] = newValue
+ }
+ }
+
+ // MARK: - Baggage: TraceID
+
+ enum LambdaTraceIDKey: Key {
+ typealias Value = String
+ static var name: String? { AmazonHeaders.traceID }
+ }
+
+ /// The AWS X-Ray tracing header.
+ public internal(set) var lambdaTraceID: String {
+ get {
+ self[LambdaTraceIDKey.self]! // !-safe, the runtime guarantees to always set an identifier, even in testing
+ }
+ set {
+ self[LambdaTraceIDKey.self] = newValue
+ }
+ }
+}
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 8fc22de3..f24f6909 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -101,7 +101,7 @@ extension Lambda {
}
private extension Lambda.Context {
- convenience init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: Lambda.Invocation) {
+ init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: Lambda.Invocation) {
self.init(requestID: invocation.requestID,
traceID: invocation.traceID,
invokedFunctionARN: invocation.invokedFunctionARN,
diff --git a/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift b/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
index 9aa3f72a..a6676014 100644
--- a/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
+++ b/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
@@ -18,6 +18,7 @@ import Logging
import NIO
import NIOFoundationCompat
import XCTest
+import Baggage
class CodableLambdaTest: XCTestCase {
var eventLoopGroup: EventLoopGroup!
@@ -63,7 +64,22 @@ class CodableLambdaTest: XCTestCase {
XCTAssertEqual(response?.requestId, request.requestId)
}
- // convencience method
+ func testBaggageContextInteractions() {
+ var context = newContext()
+ context.baggage.testValue = "hello"
+
+ context.logger.info("Test")
+ XCTAssertEqual(context.baggage.testValue, "hello")
+ XCTAssertEqual(context.baggage.lambdaTraceID, "abc123")
+ XCTAssertEqual(context.baggage.lambdaTraceID, context.traceID)
+ XCTAssertEqual(context.baggage.lambdaRequestID, context.requestID)
+
+ XCTAssertEqual("\(context.logger[metadataKey: "LambdaTraceIDKey"]!)", context.traceID)
+ XCTAssertEqual("\(context.logger[metadataKey: "LambdaRequestIDKey"]!)", context.requestID)
+ XCTAssertEqual("\(context.logger[metadataKey: "TestKey"]!)", context.baggage.testValue)
+ }
+
+ // convenience method
func newContext() -> Lambda.Context {
Lambda.Context(requestID: UUID().uuidString,
traceID: "abc123",
@@ -90,3 +106,18 @@ private struct Response: Codable, Equatable {
self.requestId = requestId
}
}
+
+extension Baggage {
+ enum TestKey: Baggage.Key {
+ typealias Value = String
+
+ }
+ var testValue: String? {
+ get {
+ return self[TestKey.self]
+ }
+ set {
+ self[TestKey.self] = newValue
+ }
+ }
+}