Read _X_AMZN_TRACE_ID from environment & system properties (#252)

* Issue: #251 - Read _X_AMZN_TRACE_ID from system properties if environment variable is not set * Update aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java Co-authored-by: William Armiros <[email protected]>
aws · Jan 22, 2021 · 67172dc · 67172dc
1 parent 636e980
commit 67172dc
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/...ray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/...ray-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/...recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/...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() {