fix(java-wrapper): link wrapper spans to invocation X-Ray trace

java/build.gradle.kts

@@ -7,6 +7,10 @@ plugins {
 base.archivesBaseName = "aws-otel-lambda-java-extensions"
 group = "software.amazon.opentelemetry.lambda"
 
+val openTelemetryVersion = "1.32.0"
+val openTelemetryAlphaVersion = "1.32.0-alpha"
+val junitVersion = "5.10.1"
+
 repositories {
     mavenCentral()
     mavenLocal()
@@ -28,12 +32,26 @@ val javaagentDependency by configurations.creating {
 }
 
 dependencies {
-    compileOnly(platform("io.opentelemetry:opentelemetry-bom:1.32.0"))
-    compileOnly(platform("io.opentelemetry:opentelemetry-bom-alpha:1.32.0-alpha"))
+    compileOnly(platform("io.opentelemetry:opentelemetry-bom:$openTelemetryVersion"))
+    compileOnly(platform("io.opentelemetry:opentelemetry-bom-alpha:$openTelemetryAlphaVersion"))
+    // opentelemetry-api and opentelemetry-context are already provided by the wrapper layer at
+    // runtime, so they are compileOnly here.
+    compileOnly("io.opentelemetry:opentelemetry-api")
+    compileOnly("io.opentelemetry:opentelemetry-context")
     // Already included in wrapper so compileOnly
     compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi")
     compileOnly("io.opentelemetry:opentelemetry-sdk-extension-aws")
     javaagentDependency("software.amazon.opentelemetry:aws-opentelemetry-agent:1.32.0-adot-lambda1")
+
+    testImplementation(platform("io.opentelemetry:opentelemetry-bom:$openTelemetryVersion"))
+    testImplementation("io.opentelemetry:opentelemetry-api")
+    testImplementation("io.opentelemetry:opentelemetry-context")
+    testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi")
+    // Supplies the X-Ray propagator the wrapper layer configures at runtime; needed here only to
+    // exercise the fallback path in tests.
+    testImplementation(
+        "io.opentelemetry.contrib:opentelemetry-aws-xray-propagator:$openTelemetryAlphaVersion")
+    testImplementation("org.junit.jupiter:junit-jupiter:$junitVersion")
 }
 
 tasks.register<Copy>("download") {
@@ -44,3 +62,7 @@ tasks.register<Copy>("download") {
 tasks.named("build") {
     dependsOn("download")
 }
+
+tasks.test {
+    useJUnitPlatform()
+}

java/src/main/java/software/amazon/opentelemetry/lambda/AwsLambdaXrayContextAutoConfigurationCustomizerProvider.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.opentelemetry.lambda;
+
+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 io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
+import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Locale;
+import java.util.Map;
+
+/**
+ * Links spans created by the OTel Lambda SDK "wrapper" layer to the X-Ray trace of the enclosing
+ * Lambda invocation.
+ *
+ * <p>On the managed Java 17+ Lambda runtimes the per-invocation X-Ray trace header is exposed via
+ * the {@code com.amazonaws.xray.traceHeader} system property rather than the {@code
+ * _X_AMZN_TRACE_ID} environment variable that older OTel Lambda instrumentation reads. Without this
+ * customizer the wrapper root span gets a fresh (X-Ray-compatible) trace id from {@link
+ * software.amazon.opentelemetry.lambda.AwsOtelTracerProviderConfigurer}, so downstream spans land
+ * in a trace that is disconnected from the invocation's X-Ray trace.
+ *
+ * <p>This customizer wraps every configured propagator. During {@code extract} it first runs the
+ * normal propagation from the event carrier (for example API Gateway headers). If that yields a
+ * valid remote parent it is preserved. Otherwise it reads the X-Ray trace-header system property
+ * and presents it to the same configured propagators as a synthetic {@code X-Amzn-Trace-Id} header,
+ * delegating the actual parsing to the already-configured X-Ray propagator. This adds no new
+ * runtime dependency: the wrapper layer supplies and (by default) configures the X-Ray propagator
+ * via {@code OTEL_PROPAGATORS}.
+ *
+ * <p>This extension jar is only packaged into the wrapper layer, so the javaagent layer (which
+ * wires X-Ray propagation through bytecode instrumentation) is unaffected.
+ */
+public final class AwsLambdaXrayContextAutoConfigurationCustomizerProvider
+    implements AutoConfigurationCustomizerProvider {
+
+  static final String TRACE_HEADER_PROPERTY = "com.amazonaws.xray.traceHeader";
+  private static final String TRACE_HEADER_KEY = "x-amzn-trace-id";
+
+  @Override
+  public void customize(AutoConfigurationCustomizer autoConfiguration) {
+    autoConfiguration.addPropagatorCustomizer(
+        (propagator, config) -> new LambdaXraySystemPropertyPropagator(propagator));
+  }
+
+  static final class LambdaXraySystemPropertyPropagator implements TextMapPropagator {
+    private static final TextMapGetter<Map<String, String>> MAP_GETTER =
+        new TextMapGetter<Map<String, String>>() {
+          @Override
+          public Iterable<String> keys(Map<String, String> carrier) {
+            return carrier.keySet();
+          }
+
+          @Override
+          public String get(Map<String, String> carrier, String key) {
+            if (carrier == null) {
+              return null;
+            }
+            return carrier.get(key.toLowerCase(Locale.ROOT));
+          }
+        };
+
+    private final TextMapPropagator delegate;
+
+    LambdaXraySystemPropertyPropagator(TextMapPropagator delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public Collection<String> fields() {
+      return delegate.fields();
+    }
+
+    @Override
+    public <C> void inject(Context context, C carrier, TextMapSetter<C> setter) {
+      delegate.inject(context, carrier, setter);
+    }
+
+    @Override
+    public <C> Context extract(Context context, C carrier, TextMapGetter<C> getter) {
+      Context base = context == null ? Context.root() : context;
+      Context extracted = delegate.extract(base, carrier, getter);
+
+      // Preserve an already-valid parent extracted from the event carrier (e.g. API Gateway
+      // headers) so explicit upstream propagation always wins.
+      if (Span.fromContext(extracted).getSpanContext().isValid()) {
+        return extracted;
+      }
+
+      // Java 17+ managed Lambda runtimes place the per-invocation X-Ray header here.
+      String traceHeader = System.getProperty(TRACE_HEADER_PROPERTY);
+      if (traceHeader == null || traceHeader.trim().isEmpty()) {
+        return extracted;
+      }
+
+      return delegate.extract(
+          extracted, Collections.singletonMap(TRACE_HEADER_KEY, traceHeader), MAP_GETTER);
+    }
+  }
+}

java/src/main/resources/META-INF/services/io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider

@@ -0,0 +1 @@
+software.amazon.opentelemetry.lambda.AwsLambdaXrayContextAutoConfigurationCustomizerProvider

java/src/test/java/software/amazon/opentelemetry/lambda/AwsLambdaXrayContextAutoConfigurationCustomizerProviderTest.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.opentelemetry.lambda;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator;
+import java.util.Collections;
+import java.util.Map;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import software.amazon.opentelemetry.lambda.AwsLambdaXrayContextAutoConfigurationCustomizerProvider.LambdaXraySystemPropertyPropagator;
+
+class AwsLambdaXrayContextAutoConfigurationCustomizerProviderTest {
+
+  private static final String XRAY_HEADER_KEY = "X-Amzn-Trace-Id";
+
+  // Trace context that the Lambda runtime exposes via the system property.
+  private static final String PROPERTY_HEADER =
+      "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=1";
+  private static final String PROPERTY_TRACE_ID = "8a3c60f7d188f8fa79d48a391a778fa6";
+  private static final String PROPERTY_SPAN_ID = "53995c3f42cd8ad8";
+
+  // A different, valid trace context that arrives on the event carrier (e.g. API Gateway headers).
+  private static final String EVENT_HEADER =
+      "Root=1-11111111-222222222222222222222222;Parent=3333333333333333;Sampled=1";
+  private static final String EVENT_TRACE_ID = "11111111222222222222222222222222";
+
+  // A second distinct system-property value, used to prove the property is re-read per extraction.
+  private static final String SECOND_PROPERTY_HEADER =
+      "Root=1-44444444-555555555555555555555555;Parent=6666666666666666;Sampled=1";
+  private static final String SECOND_PROPERTY_TRACE_ID = "44444444555555555555555555555555";
+
+  private static final TextMapGetter<Map<String, String>> EXACT_KEY_GETTER =
+      new TextMapGetter<Map<String, String>>() {
+        @Override
+        public Iterable<String> keys(Map<String, String> carrier) {
+          return carrier.keySet();
+        }
+
+        @Override
+        public String get(Map<String, String> carrier, String key) {
+          return carrier == null ? null : carrier.get(key);
+        }
+      };
+
+  @AfterEach
+  void clearTraceHeader() {
+    System.clearProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY);
+  }
+
+  private static TextMapPropagator xrayFallbackPropagator() {
+    return new LambdaXraySystemPropertyPropagator(AwsXrayPropagator.getInstance());
+  }
+
+  private static SpanContext extract(TextMapPropagator propagator, Map<String, String> carrier) {
+    Context result = propagator.extract(Context.root(), carrier, EXACT_KEY_GETTER);
+    return Span.fromContext(result).getSpanContext();
+  }
+
+  @Test
+  void emptyHeadersWithValidProperty_extractsPropertyContext() {
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY,
+        PROPERTY_HEADER);
+
+    SpanContext spanContext = extract(xrayFallbackPropagator(), Collections.emptyMap());
+
+    assertTrue(spanContext.isValid(), "expected a valid parent from the system property");
+    assertEquals(PROPERTY_TRACE_ID, spanContext.getTraceId());
+    assertEquals(PROPERTY_SPAN_ID, spanContext.getSpanId());
+    assertTrue(spanContext.isSampled());
+    assertTrue(spanContext.isRemote());
+  }
+
+  @Test
+  void validEventHeaderWithDifferentProperty_eventHeaderWins() {
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY,
+        PROPERTY_HEADER);
+
+    Map<String, String> carrier = Collections.singletonMap(XRAY_HEADER_KEY, EVENT_HEADER);
+    SpanContext spanContext = extract(xrayFallbackPropagator(), carrier);
+
+    assertTrue(spanContext.isValid());
+    assertEquals(EVENT_TRACE_ID, spanContext.getTraceId(), "event header must take precedence");
+  }
+
+  @Test
+  void missingProperty_behaviorUnchanged() {
+    System.clearProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY);
+
+    SpanContext spanContext = extract(xrayFallbackPropagator(), Collections.emptyMap());
+
+    assertFalse(spanContext.isValid(), "no header and no property must yield no parent");
+  }
+
+  @Test
+  void malformedProperty_noExceptionAndNoParent() {
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY,
+        "this-is-not-a-valid-xray-trace-header");
+
+    SpanContext spanContext = extract(xrayFallbackPropagator(), Collections.emptyMap());
+
+    assertFalse(spanContext.isValid(), "a malformed property must not produce a valid parent");
+  }
+
+  @Test
+  void propertyChangesBetweenExtractions_eachExtractionReadsCurrentValue() {
+    TextMapPropagator propagator = xrayFallbackPropagator();
+
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY,
+        PROPERTY_HEADER);
+    SpanContext first = extract(propagator, Collections.emptyMap());
+
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY,
+        SECOND_PROPERTY_HEADER);
+    SpanContext second = extract(propagator, Collections.emptyMap());
+
+    assertEquals(PROPERTY_TRACE_ID, first.getTraceId());
+    assertEquals(SECOND_PROPERTY_TRACE_ID, second.getTraceId());
+  }
+
+  @Test
+  void blankProperty_treatedAsAbsent() {
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY, "   ");
+
+    SpanContext spanContext = extract(xrayFallbackPropagator(), Collections.emptyMap());
+
+    assertFalse(spanContext.isValid(), "a whitespace-only property must be treated as absent");
+  }
+
+  @Test
+  void delegateWithoutXrayPropagator_propertyIgnored() {
+    System.setProperty(
+        AwsLambdaXrayContextAutoConfigurationCustomizerProvider.TRACE_HEADER_PROPERTY,
+        PROPERTY_HEADER);
+
+    // A delegate that does not understand X-Amzn-Trace-Id must leave the property unused, honoring
+    // an explicit OTEL_PROPAGATORS configuration that excludes xray.
+    TextMapPropagator propagator =
+        new LambdaXraySystemPropertyPropagator(W3CTraceContextPropagator.getInstance());
+
+    SpanContext spanContext = extract(propagator, Collections.emptyMap());
+
+    assertFalse(spanContext.isValid(), "non-xray delegate must ignore the X-Ray property");
+  }
+}