diff --git a/Package.swift b/Package.swift index 3f971016..ca40c623 100644 --- a/Package.swift +++ b/Package.swift @@ -36,6 +36,7 @@ let package = Package( .package(url: "https://github.com/apple/swift-log.git", from: "1.10.1"), .package(url: "https://github.com/apple/swift-collections.git", from: "1.3.0"), .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.10.1"), + .package(url: "https://github.com/apple/swift-service-context.git", from: "1.3.0"), ], targets: [ .target( @@ -44,6 +45,7 @@ let package = Package( .product(name: "NIOCore", package: "swift-nio"), .product(name: "DequeModule", package: "swift-collections"), .product(name: "Logging", package: "swift-log"), + .product(name: "ServiceContextModule", package: "swift-service-context"), .product(name: "NIOHTTP1", package: "swift-nio"), .product(name: "NIOPosix", package: "swift-nio"), .product( @@ -74,6 +76,7 @@ let package = Package( name: "AWSLambdaRuntimeTests", dependencies: [ .byName(name: "AWSLambdaRuntime"), + .product(name: "ServiceContextModule", package: "swift-service-context"), .product(name: "NIOTestUtils", package: "swift-nio"), .product(name: "NIOFoundationCompat", package: "swift-nio"), ], diff --git a/Package@swift-6.0.swift b/Package@swift-6.0.swift index be01537f..97975a55 100644 --- a/Package@swift-6.0.swift +++ b/Package@swift-6.0.swift @@ -25,6 +25,7 @@ let package = Package( .package(url: "https://github.com/apple/swift-log.git", from: "1.8.0"), .package(url: "https://github.com/apple/swift-collections.git", from: "1.3.0"), .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.9.0"), + .package(url: "https://github.com/apple/swift-service-context.git", from: "1.3.0"), ], targets: [ .target( @@ -35,6 +36,7 @@ let package = Package( .product(name: "Logging", package: "swift-log"), .product(name: "NIOHTTP1", package: "swift-nio"), .product(name: "NIOPosix", package: "swift-nio"), + .product(name: "ServiceContextModule", package: "swift-service-context"), .product(name: "ServiceLifecycle", package: "swift-service-lifecycle"), ], swiftSettings: defaultSwiftSettings @@ -61,6 +63,7 @@ let package = Package( .byName(name: "AWSLambdaRuntime"), .product(name: "NIOTestUtils", package: "swift-nio"), .product(name: "NIOFoundationCompat", package: "swift-nio"), + .product(name: "ServiceContextModule", package: "swift-service-context"), ], swiftSettings: defaultSwiftSettings ), diff --git a/Sources/AWSLambdaRuntime/Lambda.swift b/Sources/AWSLambdaRuntime/Lambda.swift index 10e3522c..29cdd3fa 100644 --- a/Sources/AWSLambdaRuntime/Lambda.swift +++ b/Sources/AWSLambdaRuntime/Lambda.swift @@ -17,6 +17,7 @@ import Dispatch import Logging import NIOCore import NIOPosix +import ServiceContextModule #if os(macOS) import Darwin.C @@ -90,12 +91,13 @@ public enum Lambda { logger.trace("Waiting for next invocation") let (invocation, writer) = try await runtimeClient.nextInvocation() + let traceId = invocation.metadata.traceID // Create a per-request logger with request-specific metadata let requestLogger = loggingConfiguration.makeLogger( label: "Lambda", requestID: invocation.metadata.requestID, - traceID: invocation.metadata.traceID + traceID: traceId ) // when log level is trace or lower, print the first 6 Mb of the payload @@ -116,26 +118,46 @@ public enum Lambda { metadata: metadata ) - do { - try await handler.handle( - invocation.event, - responseWriter: writer, - context: LambdaContext( - requestID: invocation.metadata.requestID, - traceID: invocation.metadata.traceID, - tenantID: invocation.metadata.tenantID, - invokedFunctionARN: invocation.metadata.invokedFunctionARN, - deadline: LambdaClock.Instant( - millisecondsSinceEpoch: invocation.metadata.deadlineInMillisSinceEpoch - ), - logger: requestLogger + // Wrap handler invocation in a ServiceContext scope so that + // downstream libraries can access the trace ID via + // ServiceContext.current?.traceID without depending on AWSLambdaRuntime. + // In single-concurrency mode, also set the _X_AMZN_TRACE_ID env var + // for backward compatibility with legacy tooling. + var serviceContext = ServiceContext.current ?? ServiceContext.topLevel + serviceContext.traceID = traceId + try await ServiceContext.withValue(serviceContext) { + if isSingleConcurrencyMode { + setenv("_X_AMZN_TRACE_ID", traceId, 1) + } + defer { + if isSingleConcurrencyMode { + unsetenv("_X_AMZN_TRACE_ID") + } + } + + do { + try await handler.handle( + invocation.event, + responseWriter: writer, + context: LambdaContext( + requestID: invocation.metadata.requestID, + traceID: traceId, + tenantID: invocation.metadata.tenantID, + invokedFunctionARN: invocation.metadata.invokedFunctionARN, + deadline: LambdaClock.Instant( + millisecondsSinceEpoch: invocation.metadata.deadlineInMillisSinceEpoch + ), + logger: requestLogger + ) + ) + requestLogger.trace("Handler finished processing invocation") + } catch { + requestLogger.trace( + "Handler failed processing invocation", + metadata: ["Handler error": "\(error)"] ) - ) - requestLogger.trace("Handler finished processing invocation") - } catch { - requestLogger.trace("Handler failed processing invocation", metadata: ["Handler error": "\(error)"]) - try await writer.reportError(error) - continue + try await writer.reportError(error) + } } } } catch is CancellationError { diff --git a/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift b/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift new file mode 100644 index 00000000..ab5f90c8 --- /dev/null +++ b/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift @@ -0,0 +1,50 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftAWSLambdaRuntime open source project +// +// Copyright SwiftAWSLambdaRuntime project authors +// Copyright (c) Amazon.com, Inc. or its affiliates. +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +import ServiceContextModule + +// MARK: - ServiceContext integration + +/// A ``ServiceContextKey`` for the AWS X-Ray trace ID. +/// +/// This allows downstream libraries that depend on `swift-service-context` +/// (but not on `AWSLambdaRuntime`) to access the current trace ID via +/// `ServiceContext.current?.traceID`. +private enum LambdaTraceIDKey: ServiceContextKey { + typealias Value = String + static var nameOverride: String? { AmazonHeaders.traceID } +} + +extension ServiceContext { + /// The AWS X-Ray trace ID for the current Lambda invocation, if available. + /// + /// This value is automatically set by the Lambda runtime before calling the handler + /// and is available to all code running within the handler's async task tree. + /// + /// Downstream libraries can read this without depending on `AWSLambdaRuntime`: + /// ```swift + /// if let traceID = ServiceContext.current?.traceID { + /// // propagate traceID to outgoing HTTP requests, etc. + /// } + /// ``` + public var traceID: String? { + get { + self[LambdaTraceIDKey.self] + } + set { + self[LambdaTraceIDKey.self] = newValue + } + } +} diff --git a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift new file mode 100644 index 00000000..d97b90b8 --- /dev/null +++ b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift @@ -0,0 +1,226 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftAWSLambdaRuntime open source project +// +// Copyright SwiftAWSLambdaRuntime project authors +// Copyright (c) Amazon.com, Inc. or its affiliates. +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +import ServiceContextModule +import Testing + +@testable import AWSLambdaRuntime + +#if os(macOS) +import Darwin.C +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif + +@Suite("Trace ID Propagation Tests", .serialized) +struct LambdaTraceIDPropagationTests { + + // MARK: - ServiceContext basic behavior + + @Test("ServiceContext traceID returns nil outside invocation scope") + @available(LambdaSwift 2.0, *) + func traceIDIsNilOutsideScope() async { + #expect(ServiceContext.current?.traceID == nil) + } + + @Test("ServiceContext traceID returns value inside withValue scope") + @available(LambdaSwift 2.0, *) + func traceIDAvailableInsideScope() async { + let expectedTraceID = "Root=1-abc-def123;Sampled=1" + + var context = ServiceContext.topLevel + context.traceID = expectedTraceID + + ServiceContext.withValue(context) { + #expect(ServiceContext.current?.traceID == expectedTraceID) + } + + // After scope ends, should be nil again + #expect(ServiceContext.current?.traceID == nil) + } + + @Test("ServiceContext traceID is isolated between concurrent tasks") + @available(LambdaSwift 2.0, *) + func traceIDIsolatedBetweenConcurrentTasks() async { + let traceID1 = "Root=1-aaa-111;Sampled=1" + let traceID2 = "Root=1-bbb-222;Sampled=1" + let traceID3 = "Root=1-ccc-333;Sampled=1" + + await withTaskGroup(of: Void.self) { group in + group.addTask { + var ctx = ServiceContext.topLevel + ctx.traceID = traceID1 + await ServiceContext.withValue(ctx) { + try? await Task.sleep(for: .milliseconds(50)) + #expect(ServiceContext.current?.traceID == traceID1) + } + } + group.addTask { + var ctx = ServiceContext.topLevel + ctx.traceID = traceID2 + await ServiceContext.withValue(ctx) { + try? await Task.sleep(for: .milliseconds(50)) + #expect(ServiceContext.current?.traceID == traceID2) + } + } + group.addTask { + var ctx = ServiceContext.topLevel + ctx.traceID = traceID3 + await ServiceContext.withValue(ctx) { + try? await Task.sleep(for: .milliseconds(50)) + #expect(ServiceContext.current?.traceID == traceID3) + } + } + await group.waitForAll() + } + } + + @Test("ServiceContext traceID propagates to child tasks") + @available(LambdaSwift 2.0, *) + func traceIDPropagatesToChildTasks() async { + let expectedTraceID = "Root=1-child-test;Sampled=1" + + var ctx = ServiceContext.topLevel + ctx.traceID = expectedTraceID + + await ServiceContext.withValue(ctx) { + await withTaskGroup(of: String?.self) { group in + group.addTask { + ServiceContext.current?.traceID + } + for await childTraceID in group { + #expect(childTraceID == expectedTraceID) + } + } + } + } + + // MARK: - Environment variable behavior + + @Test("_X_AMZN_TRACE_ID env var is set and cleared in single-concurrency simulation") + @available(LambdaSwift 2.0, *) + func envVarSetAndClearedInSingleConcurrency() async { + let traceID = "Root=1-envvar-test;Sampled=1" + + // Ensure it's not set before + unsetenv("_X_AMZN_TRACE_ID") + #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) + + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + ServiceContext.withValue(ctx) { + setenv("_X_AMZN_TRACE_ID", traceID, 1) + defer { unsetenv("_X_AMZN_TRACE_ID") } + + #expect(Lambda.env("_X_AMZN_TRACE_ID") == traceID) + #expect(ServiceContext.current?.traceID == traceID) + } + + #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) + } + + @Test("_X_AMZN_TRACE_ID env var is NOT set in multi-concurrency simulation") + @available(LambdaSwift 2.0, *) + func envVarNotSetInMultiConcurrency() async { + let traceID = "Root=1-multi-test;Sampled=1" + + unsetenv("_X_AMZN_TRACE_ID") + + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + ServiceContext.withValue(ctx) { + #expect(ServiceContext.current?.traceID == traceID) + #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) + } + } + + // MARK: - Background task propagation + + @Test("ServiceContext traceID remains available during simulated background work") + @available(LambdaSwift 2.0, *) + func traceIDAvailableDuringBackgroundWork() async { + let traceID = "Root=1-background-test;Sampled=1" + + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + await ServiceContext.withValue(ctx) { + #expect(ServiceContext.current?.traceID == traceID) + + try? await Task.sleep(for: .milliseconds(10)) + #expect(ServiceContext.current?.traceID == traceID) + + await withTaskGroup(of: Void.self) { group in + group.addTask { + try? await Task.sleep(for: .milliseconds(10)) + #expect(ServiceContext.current?.traceID == traceID) + } + await group.waitForAll() + } + } + } + + // MARK: - Coexistence with instance property + + @Test("ServiceContext traceID and LambdaContext instance traceID coexist independently") + @available(LambdaSwift 2.0, *) + func serviceContextAndInstancePropertyCoexist() async { + let serviceContextTraceID = "Root=1-tasklocal;Sampled=1" + let instanceTraceID = "Root=1-instance;Sampled=0" + + var ctx = ServiceContext.topLevel + ctx.traceID = serviceContextTraceID + + ServiceContext.withValue(ctx) { + let lambdaContext = LambdaContext.__forTestsOnly( + requestID: "test-request", + traceID: instanceTraceID, + tenantID: nil, + invokedFunctionARN: "arn:aws:lambda:us-east-1:123456789:function:test", + timeout: .seconds(30), + logger: .init(label: "test") + ) + + #expect(lambdaContext.traceID == instanceTraceID) + #expect(ServiceContext.current?.traceID == serviceContextTraceID) + } + } + + @Test("ServiceContext traceID and LambdaContext instance traceID match when set from the same source") + @available(LambdaSwift 2.0, *) + func serviceContextAndInstanceTraceIDMatchFromSameSource() async { + let traceID = "Root=1-65af3dc0-abc123def456;Sampled=1" + + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + ServiceContext.withValue(ctx) { + let lambdaContext = LambdaContext.__forTestsOnly( + requestID: "test-request", + traceID: traceID, + tenantID: nil, + invokedFunctionARN: "arn:aws:lambda:us-east-1:123456789:function:test", + timeout: .seconds(30), + logger: .init(label: "test") + ) + + #expect(lambdaContext.traceID == ServiceContext.current?.traceID) + } + } +} diff --git a/readme.md b/readme.md index e8be5c70..cf0df79e 100644 --- a/readme.md +++ b/readme.md @@ -397,6 +397,46 @@ try await runtime.run() You can learn how to deploy and invoke this function in [the API Gateway example README file](Examples/APIGatewayV2/README.md). +### Trace ID Propagation + +The runtime automatically propagates the AWS X-Ray trace ID to your handler's async task tree via [ServiceContext](https://github.com/apple/swift-service-context). This is the standard context propagation mechanism in the Swift server ecosystem, allowing any downstream library to access the trace ID without depending on `AWSLambdaRuntime`. + +#### Accessing the trace ID from a Lambda handler + +```swift +import AWSLambdaRuntime + +let runtime = LambdaRuntime { + (event: String, context: LambdaContext) in + + // from the context instance + context.logger.info("Trace ID: \(context.traceID)") + + return "OK" +} + +try await runtime.run() +``` + +#### Accessing the trace ID from a downstream library + +Libraries that don't depend on `AWSLambdaRuntime` can read the trace ID through `ServiceContext`: + +```swift +import ServiceContextModule + +func makeHTTPRequest() async throws { + if let traceID = ServiceContext.current?.traceID { + // Add X-Amzn-Trace-Id header to outgoing requests + } +} +``` + +This enables tracing libraries like [swift-distributed-tracing](https://github.com/apple/swift-distributed-tracing) and OpenTelemetry-Swift to discover the trace ID automatically without coupling to the Lambda runtime. + +> [!NOTE] +> In single-concurrency mode, the runtime also sets the `_X_AMZN_TRACE_ID` environment variable for backward compatibility with legacy tooling that reads the trace ID from the process environment. This environment variable is not set in multi-concurrency mode, as environment variables are shared across the process and would be subject to race conditions. + ### Integration with Swift Service LifeCycle The Swift AWS Lambda Runtime provides built-in support for [Swift Service Lifecycle](https://github.com/swift-server/swift-service-lifecycle), allowing you to manage the lifecycle of your Lambda runtime alongside other services like database clients, HTTP clients, or any other resources that need proper initialization and cleanup.