From 7f4308d00d32503f214db056a9978c8d0e5160d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Stormacq?= Date: Wed, 25 Feb 2026 12:50:28 +0100 Subject: [PATCH 1/9] Correctly propagate the trace id --- Sources/AWSLambdaRuntime/Lambda.swift | 78 ++++--- Sources/AWSLambdaRuntime/LambdaContext.swift | 20 ++ .../LambdaTraceIDPropagationTests.swift | 195 ++++++++++++++++++ 3 files changed, 252 insertions(+), 41 deletions(-) create mode 100644 Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift diff --git a/Sources/AWSLambdaRuntime/Lambda.swift b/Sources/AWSLambdaRuntime/Lambda.swift index 8bc014b84..25cd11f37 100644 --- a/Sources/AWSLambdaRuntime/Lambda.swift +++ b/Sources/AWSLambdaRuntime/Lambda.swift @@ -53,28 +53,6 @@ public enum Lambda { ) } - @available( - *, - deprecated, - message: - "This method will be removed in a future major version update. Use runLoop(runtimeClient:handler:loggingConfiguration:logger:isSingleConcurrencyMode:) instead." - ) - @inlinable - package static func runLoop( - runtimeClient: RuntimeClient, - handler: Handler, - loggingConfiguration: LoggingConfiguration, - logger: Logger - ) async throws where Handler: StreamingLambdaHandler { - try await self.runLoop( - runtimeClient: runtimeClient, - handler: handler, - loggingConfiguration: LoggingConfiguration(logger: logger), - logger: logger, - isSingleConcurrencyMode: true - ) - } - @inlinable package static func runLoop( runtimeClient: RuntimeClient, @@ -116,26 +94,44 @@ 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 TaskLocal scope so that + // LambdaContext.currentTraceID is available to all code + // in the handler's async task tree (e.g. OpenTelemetry instrumentation). + // In single-concurrency mode, also set the _X_AMZN_TRACE_ID env var + // for backward compatibility with legacy tooling. + try await LambdaContext.$currentTraceID.withValue(invocation.metadata.traceID) { + if isSingleConcurrencyMode { + setenv("_X_AMZN_TRACE_ID", invocation.metadata.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: invocation.metadata.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/LambdaContext.swift b/Sources/AWSLambdaRuntime/LambdaContext.swift index d14e16c6b..979367d6c 100644 --- a/Sources/AWSLambdaRuntime/LambdaContext.swift +++ b/Sources/AWSLambdaRuntime/LambdaContext.swift @@ -222,6 +222,26 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { "\(Self.self)(requestID: \(self.requestID), traceID: \(self.traceID), invokedFunctionARN: \(self.invokedFunctionARN), cognitoIdentity: \(self.cognitoIdentity ?? "nil"), clientContext: \(String(describing: self.clientContext)), deadline: \(self.deadline))" } + // MARK: - TaskLocal Trace ID + + /// The trace ID for the current Lambda invocation, available via Swift's `TaskLocal` mechanism. + /// + /// This enables OpenTelemetry instrumentation and other tracing libraries to discover + /// the current invocation's trace ID without requiring an explicit `LambdaContext` reference. + /// The value is automatically set by the runtime before calling the handler and is available + /// to all code running within the handler's async task tree. + /// + /// Returns `nil` when accessed outside of a Lambda invocation scope. + /// + /// ```swift + /// // Inside a Lambda handler or any code called from it: + /// if let traceID = LambdaContext.currentTraceID { + /// // Use traceID for downstream propagation + /// } + /// ``` + @TaskLocal + public static var currentTraceID: String? + /// This interface is not part of the public API and must not be used by adopters. This API is not part of semver versioning. /// The timeout is expressed relative to now package static func __forTestsOnly( diff --git a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift new file mode 100644 index 000000000..323c964a8 --- /dev/null +++ b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift @@ -0,0 +1,195 @@ +//===----------------------------------------------------------------------===// +// +// 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 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") +struct LambdaTraceIDPropagationTests { + + // MARK: - TaskLocal basic behavior + + @Test("currentTraceID returns nil outside invocation scope") + @available(LambdaSwift 2.0, *) + func currentTraceIDIsNilOutsideScope() async { + #expect(LambdaContext.currentTraceID == nil) + } + + @Test("currentTraceID returns value inside withValue scope") + @available(LambdaSwift 2.0, *) + func currentTraceIDAvailableInsideScope() async { + let expectedTraceID = "Root=1-abc-def123;Sampled=1" + + await LambdaContext.$currentTraceID.withValue(expectedTraceID) { + #expect(LambdaContext.currentTraceID == expectedTraceID) + } + + // After scope ends, should be nil again + #expect(LambdaContext.currentTraceID == nil) + } + + @Test("currentTraceID is isolated between concurrent tasks") + @available(LambdaSwift 2.0, *) + func currentTraceIDIsolatedBetweenConcurrentTasks() 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 { + await LambdaContext.$currentTraceID.withValue(traceID1) { + // Simulate some async work + try? await Task.sleep(for: .milliseconds(50)) + #expect(LambdaContext.currentTraceID == traceID1) + } + } + group.addTask { + await LambdaContext.$currentTraceID.withValue(traceID2) { + try? await Task.sleep(for: .milliseconds(50)) + #expect(LambdaContext.currentTraceID == traceID2) + } + } + group.addTask { + await LambdaContext.$currentTraceID.withValue(traceID3) { + try? await Task.sleep(for: .milliseconds(50)) + #expect(LambdaContext.currentTraceID == traceID3) + } + } + await group.waitForAll() + } + } + + @Test("currentTraceID propagates to child tasks") + @available(LambdaSwift 2.0, *) + func currentTraceIDPropagatesToChildTasks() async { + let expectedTraceID = "Root=1-child-test;Sampled=1" + + await LambdaContext.$currentTraceID.withValue(expectedTraceID) { + // Child task should inherit the TaskLocal value + await withTaskGroup(of: String?.self) { group in + group.addTask { + LambdaContext.currentTraceID + } + 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) + + // Simulate what the run loop does in single-concurrency mode + await LambdaContext.$currentTraceID.withValue(traceID) { + setenv("_X_AMZN_TRACE_ID", traceID, 1) + defer { unsetenv("_X_AMZN_TRACE_ID") } + + // During handler execution, env var should be set + #expect(Lambda.env("_X_AMZN_TRACE_ID") == traceID) + #expect(LambdaContext.currentTraceID == traceID) + } + + // After scope ends, env var should be cleared + #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" + + // Ensure it's not set before + unsetenv("_X_AMZN_TRACE_ID") + + // Simulate what the run loop does in multi-concurrency mode (isSingleConcurrencyMode = false) + // The env var should NOT be set, only the TaskLocal + await LambdaContext.$currentTraceID.withValue(traceID) { + // In multi-concurrency mode, we skip setenv entirely + // TaskLocal should still work + #expect(LambdaContext.currentTraceID == traceID) + // Env var should NOT be set + #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) + } + } + + // MARK: - Background task propagation + + @Test("currentTraceID remains available during simulated background work") + @available(LambdaSwift 2.0, *) + func currentTraceIDAvailableDuringBackgroundWork() async { + let traceID = "Root=1-background-test;Sampled=1" + + await LambdaContext.$currentTraceID.withValue(traceID) { + // Simulate sending response (the trace ID should still be available after) + #expect(LambdaContext.currentTraceID == traceID) + + // Simulate background work after response + try? await Task.sleep(for: .milliseconds(10)) + #expect(LambdaContext.currentTraceID == traceID) + + // Even deeper async work + await withTaskGroup(of: Void.self) { group in + group.addTask { + try? await Task.sleep(for: .milliseconds(10)) + #expect(LambdaContext.currentTraceID == traceID) + } + await group.waitForAll() + } + } + } + + // MARK: - Coexistence with instance property + + @Test("TaskLocal currentTraceID and instance traceID coexist independently") + @available(LambdaSwift 2.0, *) + func taskLocalAndInstancePropertyCoexist() async { + let taskLocalTraceID = "Root=1-tasklocal;Sampled=1" + let instanceTraceID = "Root=1-instance;Sampled=0" + + await LambdaContext.$currentTraceID.withValue(taskLocalTraceID) { + let context = 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") + ) + + // Instance property returns its own value + #expect(context.traceID == instanceTraceID) + // TaskLocal returns the TaskLocal value + #expect(LambdaContext.currentTraceID == taskLocalTraceID) + } + } +} From 103a0f9225857af62feac940c2b0dd316fe14e42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Stormacq?= Date: Wed, 25 Feb 2026 13:16:42 +0100 Subject: [PATCH 2/9] Serialize the tests that change env variables --- Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift index 323c964a8..0ef8f0ae1 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift @@ -25,7 +25,7 @@ import Glibc import Musl #endif -@Suite("Trace ID Propagation Tests") +@Suite("Trace ID Propagation Tests", .serialized) struct LambdaTraceIDPropagationTests { // MARK: - TaskLocal basic behavior From 6f3481d975d105a3ad58b65be9390fb08a7893a4 Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 11:33:45 +0100 Subject: [PATCH 3/9] make it explicit that LambdaContext.traceId and LambdaContext.currentTraceId are the same value --- Sources/AWSLambdaRuntime/Lambda.swift | 9 +++++---- .../LambdaTraceIDPropagationTests.swift | 20 +++++++++++++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/Sources/AWSLambdaRuntime/Lambda.swift b/Sources/AWSLambdaRuntime/Lambda.swift index 25cd11f37..69e75ae23 100644 --- a/Sources/AWSLambdaRuntime/Lambda.swift +++ b/Sources/AWSLambdaRuntime/Lambda.swift @@ -68,12 +68,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 @@ -99,9 +100,9 @@ public enum Lambda { // in the handler's async task tree (e.g. OpenTelemetry instrumentation). // In single-concurrency mode, also set the _X_AMZN_TRACE_ID env var // for backward compatibility with legacy tooling. - try await LambdaContext.$currentTraceID.withValue(invocation.metadata.traceID) { + try await LambdaContext.$currentTraceID.withValue(traceId) { if isSingleConcurrencyMode { - setenv("_X_AMZN_TRACE_ID", invocation.metadata.traceID, 1) + setenv("_X_AMZN_TRACE_ID", traceId, 1) } defer { if isSingleConcurrencyMode { @@ -115,7 +116,7 @@ public enum Lambda { responseWriter: writer, context: LambdaContext( requestID: invocation.metadata.requestID, - traceID: invocation.metadata.traceID, + traceID: traceId, tenantID: invocation.metadata.tenantID, invokedFunctionARN: invocation.metadata.invokedFunctionARN, deadline: LambdaClock.Instant( diff --git a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift index 0ef8f0ae1..8fce54fdf 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift @@ -192,4 +192,24 @@ struct LambdaTraceIDPropagationTests { #expect(LambdaContext.currentTraceID == taskLocalTraceID) } } + + @Test("TaskLocal currentTraceID and instance traceID match when set from the same source") + @available(LambdaSwift 2.0, *) + func taskLocalAndInstanceTraceIDMatchFromSameSource() async { + // Simulates what the run loop does: both are set from invocation.metadata.traceID + let traceID = "Root=1-65af3dc0-abc123def456;Sampled=1" + + await LambdaContext.$currentTraceID.withValue(traceID) { + let context = 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(context.traceID == LambdaContext.currentTraceID) + } + } } From 7dd9665eaacda0844b52fd58a648ea535310fc27 Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 11:46:22 +0100 Subject: [PATCH 4/9] add a deprecation warning on `package static func runLoop` to make API Breakage check happy --- Sources/AWSLambdaRuntime/Lambda.swift | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/Sources/AWSLambdaRuntime/Lambda.swift b/Sources/AWSLambdaRuntime/Lambda.swift index 69e75ae23..0745b9bd2 100644 --- a/Sources/AWSLambdaRuntime/Lambda.swift +++ b/Sources/AWSLambdaRuntime/Lambda.swift @@ -53,6 +53,28 @@ public enum Lambda { ) } + @available( + *, + deprecated, + message: + "This method will be removed in a future major version update. Use runLoop(runtimeClient:handler:loggingConfiguration:logger:isSingleConcurrencyMode:) instead." + ) + @inlinable + package static func runLoop( + runtimeClient: RuntimeClient, + handler: Handler, + loggingConfiguration: LoggingConfiguration, + logger: Logger + ) async throws where Handler: StreamingLambdaHandler { + try await self.runLoop( + runtimeClient: runtimeClient, + handler: handler, + loggingConfiguration: LoggingConfiguration(logger: logger), + logger: logger, + isSingleConcurrencyMode: true + ) + } + @inlinable package static func runLoop( runtimeClient: RuntimeClient, From dd1ad1aa7fef1012fbd73202d2264fbfb9f59897 Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 12:19:32 +0100 Subject: [PATCH 5/9] Use Swift Context to propage trace id to downstream libraries --- Package.swift | 3 + Sources/AWSLambdaRuntime/Lambda.swift | 11 +- Sources/AWSLambdaRuntime/LambdaContext.swift | 55 +++++--- .../LambdaTraceIDPropagationTests.swift | 131 ++++++++++-------- readme.md | 40 ++++++ 5 files changed, 156 insertions(+), 84 deletions(-) diff --git a/Package.swift b/Package.swift index 3f971016e..ca40c6234 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/Sources/AWSLambdaRuntime/Lambda.swift b/Sources/AWSLambdaRuntime/Lambda.swift index 0745b9bd2..ad9037386 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 @@ -117,12 +118,14 @@ public enum Lambda { metadata: metadata ) - // Wrap handler invocation in a TaskLocal scope so that - // LambdaContext.currentTraceID is available to all code - // in the handler's async task tree (e.g. OpenTelemetry instrumentation). + // 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. - try await LambdaContext.$currentTraceID.withValue(traceId) { + var serviceContext = ServiceContext.current ?? ServiceContext.topLevel + serviceContext.traceID = traceId + try await ServiceContext.withValue(serviceContext) { if isSingleConcurrencyMode { setenv("_X_AMZN_TRACE_ID", traceId, 1) } diff --git a/Sources/AWSLambdaRuntime/LambdaContext.swift b/Sources/AWSLambdaRuntime/LambdaContext.swift index 979367d6c..afd2fa169 100644 --- a/Sources/AWSLambdaRuntime/LambdaContext.swift +++ b/Sources/AWSLambdaRuntime/LambdaContext.swift @@ -15,6 +15,7 @@ import Logging import NIOCore +import ServiceContextModule // MARK: - Client Context @@ -222,26 +223,6 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { "\(Self.self)(requestID: \(self.requestID), traceID: \(self.traceID), invokedFunctionARN: \(self.invokedFunctionARN), cognitoIdentity: \(self.cognitoIdentity ?? "nil"), clientContext: \(String(describing: self.clientContext)), deadline: \(self.deadline))" } - // MARK: - TaskLocal Trace ID - - /// The trace ID for the current Lambda invocation, available via Swift's `TaskLocal` mechanism. - /// - /// This enables OpenTelemetry instrumentation and other tracing libraries to discover - /// the current invocation's trace ID without requiring an explicit `LambdaContext` reference. - /// The value is automatically set by the runtime before calling the handler and is available - /// to all code running within the handler's async task tree. - /// - /// Returns `nil` when accessed outside of a Lambda invocation scope. - /// - /// ```swift - /// // Inside a Lambda handler or any code called from it: - /// if let traceID = LambdaContext.currentTraceID { - /// // Use traceID for downstream propagation - /// } - /// ``` - @TaskLocal - public static var currentTraceID: String? - /// This interface is not part of the public API and must not be used by adopters. This API is not part of semver versioning. /// The timeout is expressed relative to now package static func __forTestsOnly( @@ -262,3 +243,37 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { ) } } + +// 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 index 8fce54fdf..23a1e7dcd 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift @@ -13,6 +13,7 @@ // //===----------------------------------------------------------------------===// +import ServiceContextModule import Testing @testable import AWSLambdaRuntime @@ -28,68 +29,78 @@ import Musl @Suite("Trace ID Propagation Tests", .serialized) struct LambdaTraceIDPropagationTests { - // MARK: - TaskLocal basic behavior + // MARK: - ServiceContext basic behavior - @Test("currentTraceID returns nil outside invocation scope") + @Test("ServiceContext traceID returns nil outside invocation scope") @available(LambdaSwift 2.0, *) - func currentTraceIDIsNilOutsideScope() async { - #expect(LambdaContext.currentTraceID == nil) + func traceIDIsNilOutsideScope() async { + #expect(ServiceContext.current?.traceID == nil) } - @Test("currentTraceID returns value inside withValue scope") + @Test("ServiceContext traceID returns value inside withValue scope") @available(LambdaSwift 2.0, *) - func currentTraceIDAvailableInsideScope() async { + func traceIDAvailableInsideScope() async { let expectedTraceID = "Root=1-abc-def123;Sampled=1" - await LambdaContext.$currentTraceID.withValue(expectedTraceID) { - #expect(LambdaContext.currentTraceID == expectedTraceID) + var context = ServiceContext.topLevel + context.traceID = expectedTraceID + + ServiceContext.withValue(context) { + #expect(ServiceContext.current?.traceID == expectedTraceID) } // After scope ends, should be nil again - #expect(LambdaContext.currentTraceID == nil) + #expect(ServiceContext.current?.traceID == nil) } - @Test("currentTraceID is isolated between concurrent tasks") + @Test("ServiceContext traceID is isolated between concurrent tasks") @available(LambdaSwift 2.0, *) - func currentTraceIDIsolatedBetweenConcurrentTasks() async { + 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 { - await LambdaContext.$currentTraceID.withValue(traceID1) { - // Simulate some async work + var ctx = ServiceContext.topLevel + ctx.traceID = traceID1 + await ServiceContext.withValue(ctx) { try? await Task.sleep(for: .milliseconds(50)) - #expect(LambdaContext.currentTraceID == traceID1) + #expect(ServiceContext.current?.traceID == traceID1) } } group.addTask { - await LambdaContext.$currentTraceID.withValue(traceID2) { + var ctx = ServiceContext.topLevel + ctx.traceID = traceID2 + await ServiceContext.withValue(ctx) { try? await Task.sleep(for: .milliseconds(50)) - #expect(LambdaContext.currentTraceID == traceID2) + #expect(ServiceContext.current?.traceID == traceID2) } } group.addTask { - await LambdaContext.$currentTraceID.withValue(traceID3) { + var ctx = ServiceContext.topLevel + ctx.traceID = traceID3 + await ServiceContext.withValue(ctx) { try? await Task.sleep(for: .milliseconds(50)) - #expect(LambdaContext.currentTraceID == traceID3) + #expect(ServiceContext.current?.traceID == traceID3) } } await group.waitForAll() } } - @Test("currentTraceID propagates to child tasks") + @Test("ServiceContext traceID propagates to child tasks") @available(LambdaSwift 2.0, *) - func currentTraceIDPropagatesToChildTasks() async { + func traceIDPropagatesToChildTasks() async { let expectedTraceID = "Root=1-child-test;Sampled=1" - await LambdaContext.$currentTraceID.withValue(expectedTraceID) { - // Child task should inherit the TaskLocal value + var ctx = ServiceContext.topLevel + ctx.traceID = expectedTraceID + + await ServiceContext.withValue(ctx) { await withTaskGroup(of: String?.self) { group in group.addTask { - LambdaContext.currentTraceID + ServiceContext.current?.traceID } for await childTraceID in group { #expect(childTraceID == expectedTraceID) @@ -109,17 +120,17 @@ struct LambdaTraceIDPropagationTests { unsetenv("_X_AMZN_TRACE_ID") #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) - // Simulate what the run loop does in single-concurrency mode - await LambdaContext.$currentTraceID.withValue(traceID) { + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + await ServiceContext.withValue(ctx) { setenv("_X_AMZN_TRACE_ID", traceID, 1) defer { unsetenv("_X_AMZN_TRACE_ID") } - // During handler execution, env var should be set #expect(Lambda.env("_X_AMZN_TRACE_ID") == traceID) - #expect(LambdaContext.currentTraceID == traceID) + #expect(ServiceContext.current?.traceID == traceID) } - // After scope ends, env var should be cleared #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) } @@ -128,40 +139,37 @@ struct LambdaTraceIDPropagationTests { func envVarNotSetInMultiConcurrency() async { let traceID = "Root=1-multi-test;Sampled=1" - // Ensure it's not set before unsetenv("_X_AMZN_TRACE_ID") - // Simulate what the run loop does in multi-concurrency mode (isSingleConcurrencyMode = false) - // The env var should NOT be set, only the TaskLocal - await LambdaContext.$currentTraceID.withValue(traceID) { - // In multi-concurrency mode, we skip setenv entirely - // TaskLocal should still work - #expect(LambdaContext.currentTraceID == traceID) - // Env var should NOT be set + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + await ServiceContext.withValue(ctx) { + #expect(ServiceContext.current?.traceID == traceID) #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) } } // MARK: - Background task propagation - @Test("currentTraceID remains available during simulated background work") + @Test("ServiceContext traceID remains available during simulated background work") @available(LambdaSwift 2.0, *) - func currentTraceIDAvailableDuringBackgroundWork() async { + func traceIDAvailableDuringBackgroundWork() async { let traceID = "Root=1-background-test;Sampled=1" - await LambdaContext.$currentTraceID.withValue(traceID) { - // Simulate sending response (the trace ID should still be available after) - #expect(LambdaContext.currentTraceID == traceID) + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + await ServiceContext.withValue(ctx) { + #expect(ServiceContext.current?.traceID == traceID) - // Simulate background work after response try? await Task.sleep(for: .milliseconds(10)) - #expect(LambdaContext.currentTraceID == traceID) + #expect(ServiceContext.current?.traceID == traceID) - // Even deeper async work await withTaskGroup(of: Void.self) { group in group.addTask { try? await Task.sleep(for: .milliseconds(10)) - #expect(LambdaContext.currentTraceID == traceID) + #expect(ServiceContext.current?.traceID == traceID) } await group.waitForAll() } @@ -170,14 +178,17 @@ struct LambdaTraceIDPropagationTests { // MARK: - Coexistence with instance property - @Test("TaskLocal currentTraceID and instance traceID coexist independently") + @Test("ServiceContext traceID and LambdaContext instance traceID coexist independently") @available(LambdaSwift 2.0, *) - func taskLocalAndInstancePropertyCoexist() async { - let taskLocalTraceID = "Root=1-tasklocal;Sampled=1" + func serviceContextAndInstancePropertyCoexist() async { + let serviceContextTraceID = "Root=1-tasklocal;Sampled=1" let instanceTraceID = "Root=1-instance;Sampled=0" - await LambdaContext.$currentTraceID.withValue(taskLocalTraceID) { - let context = LambdaContext.__forTestsOnly( + var ctx = ServiceContext.topLevel + ctx.traceID = serviceContextTraceID + + await ServiceContext.withValue(ctx) { + let lambdaContext = LambdaContext.__forTestsOnly( requestID: "test-request", traceID: instanceTraceID, tenantID: nil, @@ -186,21 +197,21 @@ struct LambdaTraceIDPropagationTests { logger: .init(label: "test") ) - // Instance property returns its own value - #expect(context.traceID == instanceTraceID) - // TaskLocal returns the TaskLocal value - #expect(LambdaContext.currentTraceID == taskLocalTraceID) + #expect(lambdaContext.traceID == instanceTraceID) + #expect(ServiceContext.current?.traceID == serviceContextTraceID) } } - @Test("TaskLocal currentTraceID and instance traceID match when set from the same source") + @Test("ServiceContext traceID and LambdaContext instance traceID match when set from the same source") @available(LambdaSwift 2.0, *) - func taskLocalAndInstanceTraceIDMatchFromSameSource() async { - // Simulates what the run loop does: both are set from invocation.metadata.traceID + func serviceContextAndInstanceTraceIDMatchFromSameSource() async { let traceID = "Root=1-65af3dc0-abc123def456;Sampled=1" - await LambdaContext.$currentTraceID.withValue(traceID) { - let context = LambdaContext.__forTestsOnly( + var ctx = ServiceContext.topLevel + ctx.traceID = traceID + + await ServiceContext.withValue(ctx) { + let lambdaContext = LambdaContext.__forTestsOnly( requestID: "test-request", traceID: traceID, tenantID: nil, @@ -209,7 +220,7 @@ struct LambdaTraceIDPropagationTests { logger: .init(label: "test") ) - #expect(context.traceID == LambdaContext.currentTraceID) + #expect(lambdaContext.traceID == ServiceContext.current?.traceID) } } } diff --git a/readme.md b/readme.md index e8be5c706..cf0df79e3 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. From 53095b46b31f45ef23d70eea20d0242179bd56ca Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 12:34:35 +0100 Subject: [PATCH 6/9] add Swift Service Context for Swift 6.0 and 6.1 --- Package@swift-6.0.swift | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Package@swift-6.0.swift b/Package@swift-6.0.swift index be01537f7..97975a555 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 ), From 7f4dee7509546bd8e6415cc7641076b5bda7e7b3 Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 12:47:12 +0100 Subject: [PATCH 7/9] Move service context code in a separate file --- Sources/AWSLambdaRuntime/LambdaContext.swift | 35 ------------------- .../ServiceContext+TraceId.swift | 35 +++++++++++++++++++ 2 files changed, 35 insertions(+), 35 deletions(-) create mode 100644 Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift diff --git a/Sources/AWSLambdaRuntime/LambdaContext.swift b/Sources/AWSLambdaRuntime/LambdaContext.swift index afd2fa169..d14e16c6b 100644 --- a/Sources/AWSLambdaRuntime/LambdaContext.swift +++ b/Sources/AWSLambdaRuntime/LambdaContext.swift @@ -15,7 +15,6 @@ import Logging import NIOCore -import ServiceContextModule // MARK: - Client Context @@ -243,37 +242,3 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { ) } } - -// 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/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift b/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift new file mode 100644 index 000000000..0183508fc --- /dev/null +++ b/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift @@ -0,0 +1,35 @@ +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 + } + } +} \ No newline at end of file From fb2a7dbcc5170f178df1f3b385f1e7f94bdcfff1 Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 13:19:39 +0100 Subject: [PATCH 8/9] remove warnings in tests --- .../LambdaTraceIDPropagationTests.swift | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift index 23a1e7dcd..d97b90b8c 100644 --- a/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift +++ b/Tests/AWSLambdaRuntimeTests/LambdaTraceIDPropagationTests.swift @@ -123,7 +123,7 @@ struct LambdaTraceIDPropagationTests { var ctx = ServiceContext.topLevel ctx.traceID = traceID - await ServiceContext.withValue(ctx) { + ServiceContext.withValue(ctx) { setenv("_X_AMZN_TRACE_ID", traceID, 1) defer { unsetenv("_X_AMZN_TRACE_ID") } @@ -144,7 +144,7 @@ struct LambdaTraceIDPropagationTests { var ctx = ServiceContext.topLevel ctx.traceID = traceID - await ServiceContext.withValue(ctx) { + ServiceContext.withValue(ctx) { #expect(ServiceContext.current?.traceID == traceID) #expect(Lambda.env("_X_AMZN_TRACE_ID") == nil) } @@ -187,7 +187,7 @@ struct LambdaTraceIDPropagationTests { var ctx = ServiceContext.topLevel ctx.traceID = serviceContextTraceID - await ServiceContext.withValue(ctx) { + ServiceContext.withValue(ctx) { let lambdaContext = LambdaContext.__forTestsOnly( requestID: "test-request", traceID: instanceTraceID, @@ -210,7 +210,7 @@ struct LambdaTraceIDPropagationTests { var ctx = ServiceContext.topLevel ctx.traceID = traceID - await ServiceContext.withValue(ctx) { + ServiceContext.withValue(ctx) { let lambdaContext = LambdaContext.__forTestsOnly( requestID: "test-request", traceID: traceID, From d457c2ffdfcca2d55bda4d1335a1c6f8a63f26a2 Mon Sep 17 00:00:00 2001 From: Sebastien Stormacq Date: Sun, 1 Mar 2026 20:42:09 +0100 Subject: [PATCH 9/9] fix format and license check --- .../ServiceContext+TraceId.swift | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift b/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift index 0183508fc..ab5f90c88 100644 --- a/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift +++ b/Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift @@ -1,3 +1,18 @@ +//===----------------------------------------------------------------------===// +// +// 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 @@ -32,4 +47,4 @@ extension ServiceContext { self[LambdaTraceIDKey.self] = newValue } } -} \ No newline at end of file +}