[core] Implement X-Ray Trace ID Propagation via Environment Variable

[sebsto]

The Swift AWS Lambda Runtime currently receives the X-Ray trace ID from the Lambda Runtime API via the Lambda-Runtime-Trace-Id header and stores it in the LambdaContext.traceID property. However, according to the AWS Lambda Runtime API documentation, the runtime should also set the _X_AMZN_TRACE_ID environment variable with this value to enable the X-Ray SDK to automatically trace requests.

"The runtime should set the _X_AMZN_TRACE_ID with the value of the header. The X-Ray SDK reads this to get the IDs and determine whether to trace the request."

This feature is currently missing from the Swift runtime, preventing seamless integration with X-Ray tracing SDKs.

Current Behavior

1. The runtime receives the Lambda-Runtime-Trace-Id header from the Lambda Runtime API
2. The trace ID is stored in LambdaContext.traceID
3. The _X_AMZN_TRACE_ID environment variable is not set
4. X-Ray SDKs cannot automatically discover the trace context

Expected Behavior

1. The runtime receives the Lambda-Runtime-Trace-Id header
2. The trace ID is stored in LambdaContext.traceID
3. The _X_AMZN_TRACE_ID environment variable is set with the trace ID value
4. X-Ray SDKs can automatically read the trace context from the environment variable

Implementation Considerations

Standard Lambda Functions

For standard Lambda functions, setting the environment variable is straightforward using setenv() or similar mechanisms.

Lambda Managed Instances (LMI)

Lambda Managed Instances introduce complexity because:

• Multiple concurrent invocations may be processed by the same runtime instance
• Environment variables are process-global and shared across all invocations
• Setting _X_AMZN_TRACE_ID globally would cause trace ID conflicts between concurrent invocations

How Other Runtimes Handle This

Python Runtime:

• Uses os.environ['_X_AMZN_TRACE_ID'] to set the trace ID before each handler invocation
• Reference: bootstrap.py#L434
• Boto3 SDK reads from environment variable: botocore handlers.py#L110

Java Runtime:

• Uses thread-local storage via LambdaSegmentContext to avoid environment variable conflicts
• Reference: AWSLambda.java#L241
• X-Ray SDK uses thread-local context: LambdaSegmentContext.java#L46

Node.js Runtime:

• Sets process.env._X_AMZN_TRACE_ID before each invocation
• Reference: InvokeContext.js#L115
• AWS SDK middleware injects trace ID: middleware-recursion-detection index.ts#L41

Solution: Task-Local Values

Swift's task-local values provide the ideal solution for managing per-invocation trace IDs in Lambda Managed Instances. This approach:

• Isolates trace IDs to individual invocation tasks
• Avoids global environment variable conflicts
• Aligns with Swift's structured concurrency model
• Mirrors Java's thread-local approach but leverages Swift's modern concurrency primitives

Implementation Plan

1. Define Task-Local Trace ID Storage

Create a task-local value to store the X-Ray trace ID for the current invocation:

// In Sources/AWSLambdaRuntime/LambdaContext.swift or new file

@TaskLocal
public static var xrayTraceID: String?

2. Set Task-Local Value During Invocation

Modify the runtime to set the task-local value when processing each invocation:

File: Sources/AWSLambdaRuntime/Runtime/LambdaRuntime.swift

In the invocation handling code, wrap the handler execution with the task-local value:

await $xrayTraceID.withValue(context.traceID) {
    // Execute handler with trace ID available in task-local storage
    try await handler(event, context)
}

3. Provide Public API for X-Ray SDK Integration

Add a public accessor that X-Ray SDKs can use to retrieve the current trace ID:

File: Sources/AWSLambdaRuntime/LambdaContext.swift

extension LambdaContext {
    /// Returns the X-Ray trace ID for the current invocation.
    /// This value is available via task-local storage and is safe to use in concurrent environments.
    public static var currentTraceID: String? {
        xrayTraceID
    }
}

4. Maintain Backward Compatibility with Environment Variable

For standard Lambda functions (non-LMI), also set the environment variable to maintain compatibility with existing X-Ray SDK implementations:

File: Sources/AWSLambdaRuntime/Runtime/LambdaRuntime.swift

// Set environment variable for backward compatibility
// This is safe in standard Lambda but will be overridden per-invocation in LMI
setenv("_X_AMZN_TRACE_ID", context.traceID, 1)

Files to Modify

1. Sources/AWSLambdaRuntime/LambdaContext.swift

   • Add @TaskLocal static var xrayTraceID: String?
   • Add public static var currentTraceID: String? accessor

2. Sources/AWSLambdaRuntime/Runtime/LambdaRuntime.swift

   • Set task-local value before handler execution
   • Set environment variable for backward compatibility

3. Documentation

   • Update documentation to explain X-Ray integration
   • Document the task-local API for X-Ray SDK developers
   • Add example showing how to access trace ID from within handlers

Testing Requirements

• Unit tests for task-local trace ID storage and retrieval
• Integration tests with X-Ray SDK using task-local API
• Tests for concurrent invocations (LMI scenario) verifying trace ID isolation
• Verify no trace ID leakage between concurrent invocations
• Test environment variable setting for backward compatibility
• Performance tests to ensure minimal overhead from task-local storage

References

• AWS Lambda Runtime API Documentation
• X-Ray Tracing Header Documentation
• Swift Task-Local Values Documentation

Other Runtime Implementations

• Python Runtime: bootstrap.py#L434
• Java Runtime: LambdaSegmentContext.java
• Node.js Runtime: InvokeContext.js#L115

Labels

• enhancement
• x-ray
• tracing
• observability
• lambda-managed-instances