feat(core): Migrate XRax inject calls to new API

DataDog · Jan 23, 2025 · ddb0406 · ddb0406
1 parent 09144b5
commit ddb0406
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 14 deletions.
diff --git a/...-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java b/...-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/...-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java b/...-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java
@@ -2,7 +2,6 @@
 
 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;
@@ -17,12 +16,12 @@
 import com.amazonaws.handlers.HandlerContextKey;
 import com.amazonaws.handlers.RequestHandler2;
 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;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
+import datadog.trace.core.propagation.XRayPropagator;
 import java.util.Date;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -69,7 +68,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);
+          XRayPropagator.INSTANCE.inject(span, request, DECORATE);
         } catch (Throwable e) {
           log.warn("Unable to inject trace header", e);
         }

diff --git a/...ava-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/...ava-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/...k-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/...k-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java
@@ -2,17 +2,16 @@
 
 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.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;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.core.propagation.XRayPropagator;
 import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -71,7 +70,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);
+          XRayPropagator.INSTANCE.inject(span, requestBuilder, DECORATE);
           return requestBuilder.build();
         }
       } catch (Throwable e) {

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java
@@ -1,5 +1,6 @@
 package datadog.trace.core.propagation;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
 import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext;
 
 import datadog.context.Context;
@@ -21,6 +22,8 @@ public class TracingPropagator implements Propagator {
   private final HttpCodec.Injector injector;
   private final HttpCodec.Extractor extractor;
 
+  // Create HttpCodecPropagator abstract class
+
   // TODO Javadoc
   public TracingPropagator(HttpCodec.Injector injector, HttpCodec.Extractor extractor) {
     this.injector = injector;
@@ -34,17 +37,18 @@ public static Concern concern() {
 
   @Override
   public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    AgentSpan span;
     //noinspection ConstantValue
-    if (context == null || carrier == null || setter == null) {
+    if (context == null
+        || carrier == null
+        || setter == null
+        || (span = fromContext(context)) == null) {
       return;
     }
-    AgentSpan span = AgentSpan.fromContext(context);
-    if (span != null) {
-      AgentSpanContext spanContext = span.context();
-      if (spanContext instanceof DDSpanContext) {
-        DDSpanContext ddSpanContext = (DDSpanContext) context;
-        injector.inject(ddSpanContext, carrier, setter::set);
-      }
+    AgentSpanContext spanContext = span.context();
+    if (spanContext instanceof DDSpanContext) {
+      DDSpanContext ddSpanContext = (DDSpanContext) context;
+      this.injector.inject(ddSpanContext, carrier, setter::set);
     }
   }
 

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,45 @@
+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;
+
+@ParametersAreNonnullByDefault
+public class XRayPropagator implements Propagator {
+  public static final Propagator INSTANCE = new XRayPropagator();
+  private final HttpCodec.Injector injector;
+
+  public XRayPropagator() {
+    this.injector = XRayHttpCodec.newInjector(Config.get().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) context;
+      this.injector.inject(ddSpanContext, carrier, setter::set);
+    }
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    return context;
+  }
+}