Skip to content

Correctly propagate the trace ID #652

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sebsto wants to merge 10 commits into
base: main
Choose a base branch
from
Open

Correctly propagate the trace ID #652

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
7f4308d
Correctly propagate the trace id
Feb 25, 2026
103a0f9
Serialize the tests that change env variables
Feb 25, 2026
6f3481d
make it explicit that LambdaContext.traceId and LambdaContext.current…
sebsto Mar 1, 2026
7dd9665
add a deprecation warning on `package static func runLoop` to make AP…
sebsto Mar 1, 2026
dd1ad1a
Use Swift Context to propage trace id to downstream libraries
sebsto Mar 1, 2026
53095b4
add Swift Service Context for Swift 6.0 and 6.1
sebsto Mar 1, 2026
7f4dee7
Move service context code in a separate file
sebsto Mar 1, 2026
fb2a7db
remove warnings in tests
sebsto Mar 1, 2026
b2c43fb
Merge branch 'main' into sebsto/trace-id
sebsto Mar 1, 2026
d457c2f
fix format and license check
sebsto Mar 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Package.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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"),
],
Expand Down
3 changes: 3 additions & 0 deletions Package@swift-6.0.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
Expand All @@ -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
),
Expand Down
62 changes: 42 additions & 20 deletions Sources/AWSLambdaRuntime/Lambda.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import Dispatch
import Logging
import NIOCore
import NIOPosix
import ServiceContextModule

#if os(macOS)
import Darwin.C
Expand Down Expand Up @@ -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
Expand All @@ -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")
Copy link
Contributor

@ktoso ktoso Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe you want to have a constant for this string somewhere?

}
}

do {
try await handler.handle(
invocation.event,
responseWriter: writer,
context: LambdaContext(
requestID: invocation.metadata.requestID,
traceID: traceId,
Copy link
Contributor

@ktoso ktoso Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nitpick: traceID for consistency

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 {
Expand Down
50 changes: 50 additions & 0 deletions Sources/AWSLambdaRuntime/ServiceContext+TraceId.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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? {
Copy link
Contributor

@ktoso ktoso Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is a bit problematic to declare the key in the lambda runtime, it'd be better if it was declared inside some other library that is "the xray library" and aws lambda runtime would depend on it for the key.

This way other libs which want to use xray specifically could do so without conflicting here...

I think what we may need to do in the short term -- unless we spin out a lib -- would be to call this awsLambdaXRayTraceID as long as it is, it won't cause conflicts in the future if there were some "xray library" which would be the right place to declare the proper key and aws lambda runtime would then use it as well 🤔

WDYT?

Also cc @slashmo for opinions

Copy link
Collaborator Author

@sebsto sebsto Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@maxday How runtimes are supporting this in other languages ?

Copy link
Collaborator Author

@sebsto sebsto Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jbelkins should we create a swift-xray-library that extends ServiceConfig with that key ? Can the AWS SDK use it ?

get {
self[LambdaTraceIDKey.self]
}
set {
self[LambdaTraceIDKey.self] = newValue
}
}
}
Loading