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 + } + } +}