Skip to content

Support extracting Trace Header via MDC in Lambda Segment Context #431

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

Closed
jj22ee wants to merge 1 commit into from
Closed

Support extracting Trace Header via MDC in Lambda Segment Context #431

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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 aws-xray-recorder-sdk-core/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ dependencies {
implementation("com.google.auto.value:auto-value-annotations:1.10.4")
implementation("com.google.auto.service:auto-service-annotations:1.1.1")
implementation("org.apache.httpcomponents:httpclient:4.5.14")
implementation("org.slf4j:slf4j-api:1.7.30")
Copy link
Contributor

@lukeina2z lukeina2z Aug 21, 2025

Choose a reason for hiding this comment

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

May I know why using this old version? Q complains this version has performance issue. https://mvnrepository.com/artifact/org.slf4j/slf4j-api?


annotationProcessor("com.google.auto.value:auto-value:1.10.4")

Expand All @@ -26,6 +27,8 @@ dependencies {
testImplementation("org.powermock:powermock-api-mockito2:2.0.7")
testImplementation("org.skyscreamer:jsonassert:1.3.0")
testImplementation("jakarta.servlet:jakarta.servlet-api:5.0.0")
testImplementation("org.apache.logging.log4j:log4j-slf4j-impl:2.17.2")
testImplementation("org.apache.logging.log4j:log4j-core:2.17.2")
}

tasks.jar {
Expand Down
18 changes: 13 additions & 5 deletions ...ray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,20 +30,28 @@
import java.util.Objects;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.slf4j.MDC;

public class LambdaSegmentContext implements SegmentContext {
private static final Log logger = LogFactory.getLog(LambdaSegmentContext.class);

private static final String LAMBDA_TRACE_HEADER_KEY = "_X_AMZN_TRACE_ID";

private static final String CONCURRENT_TRACE_ID_KEY = "AWS_LAMBDA_X_TRACE_ID";

// See: https://github.com/aws/aws-xray-sdk-java/issues/251
private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";

public static TraceHeader getTraceHeaderFromEnvironment() {
Copy link
Contributor

@lukeina2z lukeina2z Aug 21, 2025

Choose a reason for hiding this comment

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

Consider rename this function getTraceHeaderFromEnvironment()

String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY);
return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0
? lambdaTraceHeaderKey
: System.getProperty(LAMBDA_TRACE_HEADER_PROP));
String lambdaTraceHeaderKeyFromMdc = MDC.get(CONCURRENT_TRACE_ID_KEY);
Copy link
Contributor Author

@jj22ee jj22ee Aug 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

This slf4j API will return null if there is no implementation for logging (e.g. log4j-slf4j-impl, log4j-core).

I'm not exactly sure what guarantees (if it is guaranteed) the logging implementation to exist in Lambda Environments, but this is how AWS SDK is implementing this logic (which also only depends on slf4j API, not the implementation that is needed like in the unit test dependencies).

Copy link

@RanVaknin RanVaknin Aug 18, 2025

Choose a reason for hiding this comment

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

Hi there,

This is correct. For this to work one of the following has to provide this implementation:

  1. The Lambda runtime
  2. The AWS Java SDK v2 - not possible because it will break current customers who have their own implementation (logback etc)
  3. The Xray SDK
  4. The customer - probably not possible since this is supposed to work by default (?) and this would require customers to explicitly provide an implementation.

String lambdaTraceHeaderKeyFromEnvVar = System.getenv(LAMBDA_TRACE_HEADER_KEY);

if (lambdaTraceHeaderKeyFromMdc != null && lambdaTraceHeaderKeyFromMdc.length() > 0) {
return TraceHeader.fromString(lambdaTraceHeaderKeyFromMdc);
} else if (lambdaTraceHeaderKeyFromEnvVar != null && lambdaTraceHeaderKeyFromEnvVar.length() > 0) {
Copy link
Contributor

@lukeina2z lukeina2z Aug 21, 2025

Choose a reason for hiding this comment

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

nit: You don't need 'else' as you have 'return' in the previous block.

return TraceHeader.fromString(lambdaTraceHeaderKeyFromEnvVar);
} else {
return TraceHeader.fromString(System.getProperty(LAMBDA_TRACE_HEADER_PROP));
}
}

// SuppressWarnings is needed for passing Root TraceId to noOp segment
Expand Down
115 changes: 115 additions & 0 deletions ...recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
import org.slf4j.MDC;

@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.LENIENT)
Expand Down Expand Up @@ -174,4 +175,118 @@ private static void testContextResultsInNoOpSegmentParent() {
mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
assertThat(AWSXRay.getTraceEntity()).isNull();
}

@Test
@SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER)
void testSystemPropertyFallbackWithTraceValidation() {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
FacadeSegment parent = (FacadeSegment) subsegment.getParent();

// Verify system property values are used correctly
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
assertThat(parent.isSampled()).isTrue();

mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER)
void testEnvironmentVariableFallbackWithTraceValidation() {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
FacadeSegment parent = (FacadeSegment) subsegment.getParent();

// Verify system property values are used correctly
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
assertThat(parent.isSampled()).isTrue();

mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
}

@Test
@SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER_2)
void testMdcTakesPriorityOverSystemProperty() {
MDC.put("AWS_LAMBDA_X_TRACE_ID", TRACE_HEADER);

try {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
FacadeSegment parent = (FacadeSegment) subsegment.getParent();

// Verify MDC values are used (TRACE_HEADER), not system property values (TRACE_HEADER_2)
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
assertThat(parent.isSampled()).isTrue();

mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
} finally {
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
}
}

@Test
@SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER)
void testMdcWithEmptyStringFallsBackToSystemProperty() {
MDC.put("AWS_LAMBDA_X_TRACE_ID", "");

try {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
FacadeSegment parent = (FacadeSegment) subsegment.getParent();

// Verify system property values are used when MDC is empty
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
assertThat(parent.isSampled()).isTrue();

mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
} finally {
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
}
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER_2)
void testMdcTakesPriorityOverEnvironmentVariable() {
MDC.put("AWS_LAMBDA_X_TRACE_ID", TRACE_HEADER);

try {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
FacadeSegment parent = (FacadeSegment) subsegment.getParent();

// Verify MDC values are used (TRACE_HEADER), not environment variable values (TRACE_HEADER_2)
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
assertThat(parent.isSampled()).isTrue();

mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
} finally {
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
}
}

@Test
@SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER)
void testMdcWithEmptyStringFallsBackToEnvironmentVariable() {
MDC.put("AWS_LAMBDA_X_TRACE_ID", "");

try {
LambdaSegmentContext mockContext = new LambdaSegmentContext();
Subsegment subsegment = mockContext.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
FacadeSegment parent = (FacadeSegment) subsegment.getParent();

// Verify environment variable values are used when MDC is empty
assertThat(parent.getTraceId().toString()).isEqualTo("1-57ff426a-80c11c39b0c928905eb0828d");
assertThat(parent.getId()).isEqualTo("1234abcd1234abcd");
assertThat(parent.isSampled()).isTrue();

mockContext.endSubsegment(AWSXRay.getGlobalRecorder());
} finally {
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
}
}
}
Loading