Add xray propagators that prioritizes xray environment variable (#1032)

aws-xray-propagator/build.gradle.kts

@@ -10,4 +10,8 @@ otelJava.moduleName.set("io.opentelemetry.contrib.awsxray.propagator")
 dependencies {
   api("io.opentelemetry:opentelemetry-api")
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-trace")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-testing")
+  testImplementation("uk.org.webcompere:system-stubs-jupiter:2.0.2")
 }

aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayLambdaConfigurablePropagator.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.awsxray.propagator;
+
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.autoconfigure.spi.ConfigurablePropagatorProvider;
+
+/**
+ * A {@link ConfigurablePropagatorProvider} which allows enabling the {@link
+ * AwsXrayLambdaPropagator} with the propagator name {@code xray-lambda}.
+ */
+public final class AwsXrayLambdaConfigurablePropagator implements ConfigurablePropagatorProvider {
+  @Override
+  public TextMapPropagator getPropagator(ConfigProperties config) {
+    return AwsXrayLambdaPropagator.getInstance();
+  }
+
+  @Override
+  public String getName() {
+    return "xray-lambda";
+  }
+}

aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayLambdaPropagator.java

@@ -0,0 +1,100 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.awsxray.propagator;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.context.propagation.TextMapSetter;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nullable;
+
+/**
+ * Implementation of the AWS X-Ray Trace Header propagation protocol but with special handling for
+ * Lambda's {@code _X_AMZN_TRACE_ID} environment variable and {@code com.amazonaws.xray.traceHeader}
+ * system property.
+ *
+ * <p>To register the X-Ray propagator together with default propagator when using the SDK:
+ *
+ * <pre>{@code
+ * OpenTelemetrySdk.builder()
+ *   .setPropagators(
+ *     ContextPropagators.create(
+ *         TextMapPropagator.composite(
+ *             W3CTraceContextPropagator.getInstance(),
+ *             AwsXrayLambdaPropagator.getInstance())))
+ *    .build();
+ * }</pre>
+ */
+public final class AwsXrayLambdaPropagator implements TextMapPropagator {
+
+  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";
+  private final AwsXrayPropagator xrayPropagator = AwsXrayPropagator.getInstance();
+  private static final AwsXrayLambdaPropagator INSTANCE = new AwsXrayLambdaPropagator();
+
+  private AwsXrayLambdaPropagator() {
+    // singleton
+  }
+
+  public static AwsXrayLambdaPropagator getInstance() {
+    return INSTANCE;
+  }
+
+  @Override
+  public List<String> fields() {
+    return xrayPropagator.fields();
+  }
+
+  @Override
+  public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> setter) {
+    xrayPropagator.inject(context, carrier, setter);
+  }
+
+  @Override
+  public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C> getter) {
+    Context xrayContext = xrayPropagator.extract(context, carrier, getter);
+
+    if (Span.fromContext(context).getSpanContext().isValid()) {
+      return xrayContext;
+    }
+
+    String traceHeader = System.getProperty(AWS_TRACE_HEADER_PROP);
+    if (isEmptyOrNull(traceHeader)) {
+      traceHeader = System.getenv(AWS_TRACE_HEADER_ENV_KEY);
+    }
+    if (isEmptyOrNull(traceHeader)) {
+      return xrayContext;
+    }
+    return xrayPropagator.extract(
+        xrayContext,
+        Collections.singletonMap(AwsXrayPropagator.TRACE_HEADER_KEY, traceHeader),
+        MapGetter.INSTANCE);
+  }
+
+  private static boolean isEmptyOrNull(@Nullable String value) {
+    return value == null || value.isEmpty();
+  }
+
+  private enum MapGetter implements TextMapGetter<Map<String, String>> {
+    INSTANCE;
+
+    @Override
+    public Set<String> keys(Map<String, String> map) {
+      return map.keySet();
+    }
+
+    @Override
+    @Nullable
+    public String get(@Nullable Map<String, String> map, String s) {
+      return map == null ? null : map.get(s);
+    }
+  }
+}

aws-xray-propagator/src/main/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayPropagator.java

@@ -38,7 +38,7 @@
  *     ContextPropagators.create(
  *         TextMapPropagator.composite(
  *             W3CTraceContextPropagator.getInstance(),
- *             AWSXrayPropagator.getInstance())))
+ *             AwsXrayPropagator.getInstance())))
  *    .build();
  * }</pre>
  */

aws-xray-propagator/src/test/java/io/opentelemetry/contrib/awsxray/propagator/AwsXrayLambdaPropagatorTest.java

@@ -0,0 +1,141 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.contrib.awsxray.propagator;
+
+import static io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator.TRACE_HEADER_KEY;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import java.util.Collections;
+import java.util.Map;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+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;
+
+@ExtendWith(SystemStubsExtension.class)
+class AwsXrayLambdaPropagatorTest extends AwsXrayPropagatorTest {
+
+  @SystemStub final EnvironmentVariables environmentVariables = new EnvironmentVariables();
+  @SystemStub final SystemProperties systemProperties = new SystemProperties();
+
+  private Tracer tracer;
+
+  @Override
+  TextMapPropagator propagator() {
+    return AwsXrayLambdaPropagator.getInstance();
+  }
+
+  @BeforeEach
+  public void setup() {
+    tracer = SdkTracerProvider.builder().build().get("awsxray");
+  }
+
+  @Test
+  void extract_fromEnvironmentVariable() {
+    environmentVariables.set(
+        "_X_AMZN_TRACE_ID",
+        "Root=1-00000001-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=1;Foo=Bar");
+
+    assertThat(
+            getSpanContext(propagator().extract(Context.current(), Collections.emptyMap(), GETTER)))
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                "00000001d188f8fa79d48a391a778fa6",
+                SPAN_ID,
+                TraceFlags.getSampled(),
+                TraceState.getDefault()));
+  }
+
+  @Test
+  void extract_fromSystemProperty() {
+    systemProperties.set(
+        "com.amazonaws.xray.traceHeader",
+        "Root=1-00000001-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=1;Foo=Bar");
+
+    assertThat(
+            getSpanContext(propagator().extract(Context.current(), Collections.emptyMap(), GETTER)))
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                "00000001d188f8fa79d48a391a778fa6",
+                SPAN_ID,
+                TraceFlags.getSampled(),
+                TraceState.getDefault()));
+  }
+
+  @Test
+  void extract_systemPropertyBeforeEnvironmentVariable() {
+    environmentVariables.set(
+        "_X_AMZN_TRACE_ID",
+        "Root=1-00000001-240000000000000000000001;Parent=1600000000000001;Sampled=1;Foo=Bar");
+    systemProperties.set(
+        "com.amazonaws.xray.traceHeader",
+        "Root=1-00000002-240000000000000000000002;Parent=1600000000000002;Sampled=1;Foo=Baz");
+
+    assertThat(
+            getSpanContext(propagator().extract(Context.current(), Collections.emptyMap(), GETTER)))
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                "00000002240000000000000000000002",
+                "1600000000000002",
+                TraceFlags.getSampled(),
+                TraceState.getDefault()));
+  }
+
+  @Test
+  void addLink_SystemProperty() {
+    Map<String, String> carrier =
+        Collections.singletonMap(
+            TRACE_HEADER_KEY,
+            "Root=1-00000001-240000000000000000000001;Parent=1600000000000001;Sampled=1");
+    environmentVariables.set(
+        "_X_AMZN_TRACE_ID",
+        "Root=1-00000002-240000000000000000000002;Parent=1600000000000002;Sampled=1;Foo=Bar");
+    systemProperties.set(
+        "com.amazonaws.xray.traceHeader",
+        "Root=1-00000003-240000000000000000000003;Parent=1600000000000003;Sampled=1;Foo=Baz");
+
+    Context extract = propagator().extract(Context.current(), carrier, GETTER);
+    ReadableSpan span =
+        (ReadableSpan)
+            tracer
+                .spanBuilder("test")
+                .setParent(extract)
+                .addLink(
+                    Span.fromContext(propagator().extract(extract, carrier, GETTER))
+                        .getSpanContext())
+                .startSpan();
+    assertThat(span.getParentSpanContext())
+        .isEqualTo(
+            SpanContext.createFromRemoteParent(
+                "00000003240000000000000000000003",
+                "1600000000000003",
+                TraceFlags.getSampled(),
+                TraceState.getDefault()));
+
+    assertThat(span.toSpanData().getLinks())
+        .isEqualTo(
+            Collections.singletonList(
+                LinkData.create(
+                    SpanContext.createFromRemoteParent(
+                        "00000001240000000000000000000001",
+                        "1600000000000001",
+                        TraceFlags.getSampled(),
+                        TraceState.getDefault()))));
+  }
+}