diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java index 2a7047b8..0af9b036 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java +++ b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java @@ -35,9 +35,15 @@ 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"; + + // See: https://github.com/aws/aws-xray-sdk-java/issues/251 + private static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader"; private static TraceHeader getTraceHeaderFromEnvironment() { - return TraceHeader.fromString(System.getenv(LAMBDA_TRACE_HEADER_KEY)); + String lambdaTraceHeaderKey = System.getenv(LAMBDA_TRACE_HEADER_KEY); + return TraceHeader.fromString(lambdaTraceHeaderKey != null && lambdaTraceHeaderKey.length() > 0 + ? lambdaTraceHeaderKey + : System.getProperty(LAMBDA_TRACE_HEADER_PROP)); } private static boolean isInitializing(TraceHeader traceHeader) { diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java index 61bd0201..ae696153 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java +++ b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java @@ -30,6 +30,7 @@ import org.junit.jupiter.api.TestInstance; import org.junit.jupiter.api.extension.ExtendWith; import org.junitpioneer.jupiter.SetEnvironmentVariable; +import org.junitpioneer.jupiter.SetSystemProperty; import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; @@ -117,6 +118,14 @@ void anotherInvocationGeneratesSegment() { assertThat(secondInvocation).isNotNull(); assertThat(secondInvocation.getParent()).isInstanceOf(FacadeSegment.class); } + + @Test + @SetSystemProperty(key = "com.amazonaws.xray.traceHeader", value = TRACE_HEADER) + void oneInvocationGeneratesSegmentUsingSystemProperty() { + Subsegment firstInvocation = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test"); + assertThat(firstInvocation).isNotNull(); + assertThat(firstInvocation.getParent()).isInstanceOf(FacadeSegment.class); + } } private static void testContextResultsInFacadeSegmentParent() {