Move X-Ray Env Variable propagation to span link instead of parent

open-telemetry/opentelemetry-specification#3166 Per discussion in the FAAS SIG, we decided that the AWS X-Ray environment variable should be moved to a span link to avoid interfering with the configured propagators.
tylerbenson · Mar 3, 2023 · 9fdd241 · 9fdd241
1 parent ee781f5
commit 9fdd241
Show file tree

Hide file tree

Showing 10 changed files with 210 additions and 220 deletions.
diff --git a/.../io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ApiGatewayProxyRequest.java b/.../io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ApiGatewayProxyRequest.java
@@ -30,14 +30,7 @@ public abstract class ApiGatewayProxyRequest {
   private static boolean noHttpPropagationNeeded() {
     Collection<String> fields =
         GlobalOpenTelemetry.getPropagators().getTextMapPropagator().fields();
-    return fields.isEmpty() || xrayPropagationFieldsOnly(fields);
-  }
-
-  private static boolean xrayPropagationFieldsOnly(Collection<String> fields) {
-    // ugly but faster than typical convert-to-set-and-check-contains-only
-    return (fields.size() == 1)
-        && ParentContextExtractor.AWS_TRACE_HEADER_PROPAGATOR_KEY.equalsIgnoreCase(
-            fields.iterator().next());
+    return fields.isEmpty();
   }
 
   public static ApiGatewayProxyRequest forStream(InputStream source) {

diff --git a/...ntelemetry/instrumentation/awslambdacore/v1_0/internal/AwsLambdaFunctionInstrumenter.java b/...ntelemetry/instrumentation/awslambdacore/v1_0/internal/AwsLambdaFunctionInstrumenter.java
@@ -11,6 +11,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.internal.ContextPropagationDebug;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.AwsLambdaRequest;
+import java.util.Locale;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -46,15 +47,25 @@ public void end(
   }
 
   public Context extract(AwsLambdaRequest input) {
-    return ParentContextExtractor.extract(input.getHeaders(), this);
-  }
-
-  public Context extract(Map<String, String> headers, TextMapGetter<Map<String, String>> getter) {
     ContextPropagationDebug.debugContextLeakIfEnabled();
 
     return openTelemetry
         .getPropagators()
         .getTextMapPropagator()
-        .extract(Context.root(), headers, getter);
+        .extract(Context.root(), input.getHeaders(), MapGetter.INSTANCE);
+  }
+
+  private enum MapGetter implements TextMapGetter<Map<String, String>> {
+    INSTANCE;
+
+    @Override
+    public Iterable<String> keys(Map<String, String> map) {
+      return map.keySet();
+    }
+
+    @Override
+    public String get(Map<String, String> map, String s) {
+      return map.get(s.toLowerCase(Locale.ROOT));
+    }
   }
 }
diff --git a/...try/instrumentation/awslambdacore/v1_0/internal/AwsLambdaFunctionInstrumenterFactory.java b/...try/instrumentation/awslambdacore/v1_0/internal/AwsLambdaFunctionInstrumenterFactory.java
@@ -23,6 +23,7 @@ public static AwsLambdaFunctionInstrumenter createInstrumenter(OpenTelemetry ope
                 openTelemetry,
                 "io.opentelemetry.aws-lambda-core-1.0",
                 AwsLambdaFunctionInstrumenterFactory::spanName)
+            .addSpanLinksExtractor(new AwsXRayEnvSpanLinksExtractor())
             .addAttributesExtractor(new AwsLambdaFunctionAttributesExtractor())
             .buildInstrumenter(SpanKindExtractor.alwaysServer()));
   }

diff --git a/...entelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractor.java b/...entelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractor.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awslambdacore.v1_0.internal;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanLinksBuilder;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanLinksExtractor;
+import io.opentelemetry.instrumentation.awslambdacore.v1_0.AwsLambdaRequest;
+import java.util.Collections;
+import java.util.Locale;
+import java.util.Map;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public final class AwsXRayEnvSpanLinksExtractor implements SpanLinksExtractor<AwsLambdaRequest> {
+
+  private static final String AWS_TRACE_HEADER_ENV_KEY = "_X_AMZN_TRACE_ID";
+
+  // lower-case map getter used for extraction
+  private static final String AWS_TRACE_HEADER_PROPAGATOR_KEY = "x-amzn-trace-id";
+
+  @Override
+  public void extract(
+      SpanLinksBuilder spanLinks,
+      io.opentelemetry.context.Context parentContext,
+      AwsLambdaRequest awsLambdaRequest) {
+    extract(spanLinks);
+  }
+
+  public static void extract(SpanLinksBuilder spanLinks) {
+    String parentTraceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
+    if (parentTraceHeader == null || parentTraceHeader.isEmpty()) {
+      return;
+    }
+    Context parentCtx =
+        AwsXrayPropagator.getInstance()
+            .extract(
+                // see BaseTracer#extract() on why we're using root() here
+                Context.root(),
+                Collections.singletonMap(AWS_TRACE_HEADER_PROPAGATOR_KEY, parentTraceHeader),
+                MapGetter.INSTANCE);
+    SpanContext parent = Span.fromContext(parentCtx).getSpanContext();
+    if (parent.isValid()) {
+      spanLinks.addLink(parent);
+    }
+  }
+
+  private enum MapGetter implements TextMapGetter<Map<String, String>> {
+    INSTANCE;
+
+    @Override
+    public Iterable<String> keys(Map<String, String> map) {
+      return map.keySet();
+    }
+
+    @Override
+    public String get(Map<String, String> map, String s) {
+      return map.get(s.toLowerCase(Locale.ROOT));
+    }
+  }
+}
diff --git a/.../io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractor.java b/.../io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractor.java
diff --git a/...lemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractorTest.java b/...lemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractorTest.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awslambdacore.v1_0.internal;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanLinksBuilder;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
+import uk.org.webcompere.systemstubs.jupiter.SystemStub;
+import uk.org.webcompere.systemstubs.jupiter.SystemStubsExtension;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+@ExtendWith(SystemStubsExtension.class)
+class AwsXRayEnvSpanLinksExtractorTest {
+
+  @SystemStub final EnvironmentVariables environmentVariables = new EnvironmentVariables();
+
+  @Test
+  void shouldIgnoreIfEnvVarEmpty() {
+    // given
+    SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
+    environmentVariables.set("_X_AMZN_TRACE_ID", "");
+
+    // when
+    AwsXRayEnvSpanLinksExtractor.extract(spanLinksBuilder);
+    // then
+    verifyNoInteractions(spanLinksBuilder);
+  }
+
+  @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());
+    SpanContext spanContext = captor.getValue();
+    assertThat(spanContext.isValid()).isTrue();
+    assertThat(spanContext.isSampled()).isFalse();
+    assertThat(spanContext.getSpanId()).isEqualTo("0000000000000456");
+    assertThat(spanContext.getTraceId()).isEqualTo("8a3c60f7d188f8fa79d48a391a778fa6");
+  }
+
+  @Test
+  void shouldLinkAwsParentHeaderIfValidAndSampled() {
+    // given
+    SpanLinksBuilder spanLinksBuilder = mock(SpanLinksBuilder.class);
+    environmentVariables.set(
+        "_X_AMZN_TRACE_ID",
+        "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=0000000000000456;Sampled=1");
+
+    // when
+    AwsXRayEnvSpanLinksExtractor.extract(spanLinksBuilder);
+    // then
+    ArgumentCaptor<SpanContext> captor = ArgumentCaptor.forClass(SpanContext.class);
+    verify(spanLinksBuilder).addLink(captor.capture());
+    SpanContext spanContext = captor.getValue();
+    assertThat(spanContext.isValid()).isTrue();
+    assertThat(spanContext.isSampled()).isTrue();
+    assertThat(spanContext.getSpanId()).isEqualTo("0000000000000456");
+    assertThat(spanContext.getTraceId()).isEqualTo("8a3c60f7d188f8fa79d48a391a778fa6");
+  }
+}