[RFC] Adopts *Writable* BaggageContext (Better)

Package.swift

@@ -17,6 +17,8 @@ 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(name: "swift-context", path: "/Users/ktoso/code/gsoc-swift-baggage-context"), // TODO: remove development dep
         .package(url: "https://github.com/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
     ],
     targets: [
@@ -27,6 +29,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"),
         ]),

Sources/AWSLambdaRuntimeCore/LambdaContext.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import BaggageContext
 import Dispatch
 import Logging
 import NIO
@@ -49,40 +50,118 @@ 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 {
+        private var storage: _Storage
+
+        final class _Storage {
+            var baggage: Baggage
+
+            let invokedFunctionARN: String
+            let deadline: DispatchWallTime
+            let cognitoIdentity: String?
+            let clientContext: String?
+
+            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.baggage = newValue
+                } else {
+                    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
+        public var logger: Logger {
+            self.storage._logger.with(self.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 +172,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.background
+            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 +225,41 @@ 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
+        }
+    }
+}

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,