feat(core): Migrate X-Ray inject calls to new API

PerfectSlayer · PerfectSlayer · commit bda8a1da94d7 · 2025-02-07T17:43:42.000+01:00
diff --git a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java
@@ -333,6 +333,7 @@ protected int status(final Response response) {
     return response.getHttpResponse().getStatusCode();
   }
 
+  @ParametersAreNonnullByDefault
   @Override
   public void set(Request<?> carrier, String key, String value) {
     carrier.addHeader(key, value);
diff --git a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java
@@ -1,8 +1,8 @@
 package datadog.trace.instrumentation.aws.v0;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_IN;
 import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
@@ -16,8 +16,8 @@
 import com.amazonaws.Response;
 import com.amazonaws.handlers.HandlerContextKey;
 import com.amazonaws.handlers.RequestHandler2;
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
-import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentDataStreamsMonitoring;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -69,7 +69,7 @@ public void beforeRequest(final Request<?> request) {
       request.addHandlerContext(SPAN_CONTEXT_KEY, span);
       if (Config.get().isAwsPropagationEnabled()) {
         try {
-          propagate().inject(span, request, DECORATE, TracePropagationStyle.XRAY);
+          Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, request, DECORATE);
         } catch (Throwable e) {
           log.warn("Unable to inject trace header", e);
         }
diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java
@@ -37,6 +37,7 @@
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nonnull;
+import javax.annotation.ParametersAreNonnullByDefault;
 import software.amazon.awssdk.awscore.AwsResponse;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.SdkField;
@@ -443,6 +444,7 @@ protected int status(final SdkHttpResponse response) {
     return response.statusCode();
   }
 
+  @ParametersAreNonnullByDefault
   @Override
   public void set(SdkHttpRequest.Builder carrier, String key, String value) {
     carrier.putHeader(key, value);
diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java
@@ -1,14 +1,14 @@
 package datadog.trace.instrumentation.aws.v2;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_LEGACY_TRACING;
 import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.DECORATE;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
-import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstanceStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -71,7 +71,7 @@ public SdkHttpRequest modifyHttpRequest(
         final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
         if (span != null) {
           SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder();
-          propagate().inject(span, requestBuilder, DECORATE, TracePropagationStyle.XRAY);
+          Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, requestBuilder, DECORATE);
           return requestBuilder.build();
         }
       } catch (Throwable e) {
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
@@ -6,6 +6,7 @@
 import static datadog.trace.api.DDTags.DSM_ENABLED;
 import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
 import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator;
 import static datadog.trace.util.AgentThreadFactory.AGENT_THREAD_GROUP;
 import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap;
@@ -90,6 +91,7 @@
 import datadog.trace.core.propagation.HttpCodec;
 import datadog.trace.core.propagation.PropagationTags;
 import datadog.trace.core.propagation.TracingPropagator;
+import datadog.trace.core.propagation.XRayPropagator;
 import datadog.trace.core.scopemanager.ContinuableScopeManager;
 import datadog.trace.core.taginterceptor.RuleFlags;
 import datadog.trace.core.taginterceptor.TagInterceptor;
@@ -724,6 +726,7 @@ private CoreTracer(
         new CorePropagation(builtExtractor, injector, injectors, dataStreamContextInjector);
 
     Propagators.register(TRACING_CONCERN, new TracingPropagator(injector, extractor));
+    Propagators.register(XRAY_TRACING_CONCERN, new XRayPropagator(config), false);
 
     this.tagInterceptor =
         null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor;
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayPropagator.java
@@ -0,0 +1,47 @@
+package datadog.trace.core.propagation;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
+
+import datadog.context.Context;
+import datadog.context.propagation.CarrierSetter;
+import datadog.context.propagation.CarrierVisitor;
+import datadog.context.propagation.Propagator;
+import datadog.trace.api.Config;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.core.DDSpanContext;
+import javax.annotation.ParametersAreNonnullByDefault;
+
+// TODO Javadoc
+@ParametersAreNonnullByDefault
+public class XRayPropagator implements Propagator {
+  private final HttpCodec.Injector injector;
+
+  // TODO Javadoc
+  public XRayPropagator(Config config) {
+    this.injector = XRayHttpCodec.newInjector(config.getBaggageMapping());
+  }
+
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    AgentSpan span;
+    //noinspection ConstantValue
+    if (context == null
+        || carrier == null
+        || setter == null
+        || (span = fromContext(context)) == null) {
+      return;
+    }
+    AgentSpanContext spanContext = span.context();
+    if (spanContext instanceof DDSpanContext) {
+      DDSpanContext ddSpanContext = (DDSpanContext) spanContext;
+      ddSpanContext.getTraceCollector().setSamplingPriorityIfNecessary();
+      this.injector.inject(ddSpanContext, carrier, setter::set);
+    }
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    return context;
+  }
+}
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/TracingPropagatorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/TracingPropagatorTest.groovy
@@ -11,6 +11,7 @@ import datadog.trace.core.test.DDCoreSpecification
 
 import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP
 import static datadog.trace.api.sampling.PrioritySampling.USER_DROP
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN
 
 class TracingPropagatorTest extends DDCoreSpecification {
   HttpCodec.Injector injector
@@ -136,4 +137,23 @@ class TracingPropagatorTest extends DDCoreSpecification {
     root.finish()
     tracer.close()
   }
+
+  def 'test AWS X-Ray propagator'() {
+    setup:
+    def tracer = tracerBuilder().build()
+    def span = tracer.buildSpan('test', 'operation').start()
+    def propagator = Propagators.forConcerns(XRAY_TRACING_CONCERN)
+    def setter = Mock(CarrierSetter)
+    def carrier = new Object()
+
+    when:
+    propagator.inject(span, carrier, setter)
+
+    then:
+    1 * setter.set(carrier, 'X-Amzn-Trace-Id', _)
+
+    cleanup:
+    span.finish()
+    tracer.close()
+  }
 }
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java
@@ -10,6 +10,7 @@
 
 public interface AgentPropagation {
   Concern TRACING_CONCERN = Concern.named("tracing");
+  Concern XRAY_TRACING_CONCERN = Concern.named("tracing-xray");
 
   <C> void inject(AgentSpan span, C carrier, Setter<C> setter);
 

Original file line number	Diff line number	Diff line change
`@@ -333,6 +333,7 @@ protected int status(final Response response) {`
`333`	`333`	`return response.getHttpResponse().getStatusCode();`
`334`	`334`	`}`
`335`	`335`
	`336`	`+ @ParametersAreNonnullByDefault`
`336`	`337`	`@Override`
`337`	`338`	`public void set(Request<?> carrier, String key, String value) {`
`338`	`339`	`carrier.addHeader(key, value);`