read tracing info from system properties (#8368)

open-telemetry · May 19, 2023 · 6c34365 · 6c34365
1 parent 474a6d0
commit 6c34365
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 14 deletions.
diff --git a/...entelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXrayEnvSpanLinksExtractor.java b/...entelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXrayEnvSpanLinksExtractor.java
@@ -26,7 +26,7 @@
 final class AwsXrayEnvSpanLinksExtractor implements SpanLinksExtractor<AwsLambdaRequest> {
 
   private static final String AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID";
-
+  private static final String AWS_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";
   // lower-case map getter used for extraction
   private static final String AWS_TRACE_HEADER_PROPAGATOR_KEY = "x-amzn-trace-id";
 
@@ -42,22 +42,32 @@ public void extract(
   }
 
   public static void extract(SpanLinksBuilder spanLinks) {
-    String parentTraceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
-    if (parentTraceHeader == null || parentTraceHeader.isEmpty()) {
+    Map<String, String> contextMap = getTraceHeaderMap();
+    if (contextMap.isEmpty()) {
       return;
     }
     Context xrayContext =
-        AwsXrayPropagator.getInstance()
-            .extract(
-                Context.root(),
-                Collections.singletonMap(AWS_TRACE_HEADER_PROPAGATOR_KEY, parentTraceHeader),
-                MapGetter.INSTANCE);
+        AwsXrayPropagator.getInstance().extract(Context.root(), contextMap, MapGetter.INSTANCE);
     SpanContext envVarSpanCtx = Span.fromContext(xrayContext).getSpanContext();
     if (envVarSpanCtx.isValid()) {
       spanLinks.addLink(envVarSpanCtx, LINK_ATTRIBUTES);
     }
   }
 
+  private static Map<String, String> getTraceHeaderMap() {
+    String traceHeader = System.getProperty(AWS_TRACE_HEADER_PROP);
+    if (isEmptyOrNull(traceHeader)) {
+      traceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
+    }
+    return isEmptyOrNull(traceHeader)
+        ? Collections.emptyMap()
+        : Collections.singletonMap(AWS_TRACE_HEADER_PROPAGATOR_KEY, traceHeader);
+  }
+
+  private static boolean isEmptyOrNull(String value) {
+    return value == null || value.isEmpty();
+  }
+
   private enum MapGetter implements TextMapGetter<Map<String, String>> {
     INSTANCE;
 

diff --git a/...lemetry/instrumentation/awslambdacore/v1_0/internal/AwsXrayEnvSpanLinksExtractorTest.java b/...lemetry/instrumentation/awslambdacore/v1_0/internal/AwsXrayEnvSpanLinksExtractorTest.java
@@ -21,6 +21,7 @@
 import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
 import uk.org.webcompere.systemstubs.jupiter.SystemStub;
 import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension;
+import uk.org.webcompere.systemstubs.properties.SystemProperties;
 
 /**
  * This class is internal and is hence not for public use. Its APIs are unstable and can change at
@@ -32,27 +33,68 @@ class AwsXrayEnvSpanLinksExtractorTest {
       Attributes.of(AttributeKey.stringKey("source"), "x-ray-env");
 
   @SystemStub final EnvironmentVariables environmentVariables = new EnvironmentVariables();
+  @SystemStub final SystemProperties systemProperties = new SystemProperties();
 
   @Test
-  void shouldIgnoreIfEnvVarEmpty() {
+  void shouldIgnoreIfEnvVarAndSystemPropertyEmpty() {
     // given
     SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
     environmentVariables.set("_X_AMZN_TRACE_ID", "");
-
+    systemProperties.set("com.amazonaws.xray.traceHeader", "");
     // when
     AwsXrayEnvSpanLinksExtractor.extract(spanLinksBuilder);
     // then
     verifyNoInteractions(spanLinksBuilder);
   }
 
+  @Test
+  void shouldLinkAwsParentHeaderAndChooseSystemPropertyIfValidAndNotSampled() {
+    // given
+    SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
+    environmentVariables.set(
+        "_X_AMZN_TRACE_ID",
+        "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=0000000000000456;Sampled=0");
+    systemProperties.set(
+        "com.amazonaws.xray.traceHeader",
+        "Root=1-8a3c60f7-d188f8fa79d48a391a778fa7;Parent=0000000000000789;Sampled=0");
+    // when
+    AwsXrayEnvSpanLinksExtractor.extract(spanLinksBuilder);
+    // then
+    ArgumentCaptor<SpanContext> captor = ArgumentCaptor.forClass(SpanContext.class);
+    verify(spanLinksBuilder).addLink(captor.capture(), eq(EXPECTED_LINK_ATTRIBUTES));
+    SpanContext spanContext = captor.getValue();
+    assertThat(spanContext.isValid()).isTrue();
+    assertThat(spanContext.isSampled()).isFalse();
+    assertThat(spanContext.getSpanId()).isEqualTo("0000000000000789");
+    assertThat(spanContext.getTraceId()).isEqualTo("8a3c60f7d188f8fa79d48a391a778fa7");
+  }
+
   @Test
   void shouldLinkAwsParentHeaderIfValidAndNotSampled() {
     // given
     SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
     environmentVariables.set(
         "_X_AMZN_TRACE_ID",
         "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=0000000000000456;Sampled=0");
+    // when
+    AwsXrayEnvSpanLinksExtractor.extract(spanLinksBuilder);
+    // then
+    ArgumentCaptor<SpanContext> captor = ArgumentCaptor.forClass(SpanContext.class);
+    verify(spanLinksBuilder).addLink(captor.capture(), eq(EXPECTED_LINK_ATTRIBUTES));
+    SpanContext spanContext = captor.getValue();
+    assertThat(spanContext.isValid()).isTrue();
+    assertThat(spanContext.isSampled()).isFalse();
+    assertThat(spanContext.getSpanId()).isEqualTo("0000000000000456");
+    assertThat(spanContext.getTraceId()).isEqualTo("8a3c60f7d188f8fa79d48a391a778fa6");
+  }
 
+  @Test
+  void shouldLinkAwsParentHeaderIfValidAndNotSampledSystemProperty() {
+    // given
+    SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
+    systemProperties.set(
+        "com.amazonaws.xray.traceHeader",
+        "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=0000000000000456;Sampled=0");
     // when
     AwsXrayEnvSpanLinksExtractor.extract(spanLinksBuilder);
     // then
@@ -66,13 +108,12 @@ void shouldLinkAwsParentHeaderIfValidAndNotSampled() {
   }
 
   @Test
-  void shouldLinkAwsParentHeaderIfValidAndSampled() {
+  void shouldLinkAwsParentHeaderIfValidAndSampledSystemProperty() {
     // given
     SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
-    environmentVariables.set(
-        "_X_AMZN_TRACE_ID",
+    systemProperties.set(
+        "com.amazonaws.xray.traceHeader",
         "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=0000000000000456;Sampled=1");
-
     // when
     AwsXrayEnvSpanLinksExtractor.extract(spanLinksBuilder);
     // then