feat(core): Refactor propagation calls to the new inject API

DataDog · Jan 22, 2025 · 55a7ec9 · 55a7ec9
1 parent c73768e
commit 55a7ec9
Show file tree

Hide file tree

Showing 77 changed files with 190 additions and 112 deletions.
diff --git a/.../java/datadog/trace/bootstrap/instrumentation/httpurlconnection/HeadersInjectAdapter.java b/.../java/datadog/trace/bootstrap/instrumentation/httpurlconnection/HeadersInjectAdapter.java
@@ -3,6 +3,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import java.net.HttpURLConnection;
 
+@ParametersAreNonnullByDefault
 public class HeadersInjectAdapter implements AgentPropagation.Setter<HttpURLConnection> {
 
   public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter();

diff --git a/...t-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ContextPayload.java b/...t-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ContextPayload.java
@@ -1,7 +1,6 @@
 package datadog.trace.bootstrap.instrumentation.rmi;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
-
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.io.IOException;
@@ -33,7 +32,7 @@ public Map<String, String> getContext() {
 
   public static ContextPayload from(final AgentSpan span) {
     final ContextPayload payload = new ContextPayload();
-    propagate().inject(span, payload, SETTER);
+    Propagators.defaultPropagator().inject(span, payload, SETTER);
     return payload;
   }
 
@@ -54,6 +53,7 @@ public void write(final ObjectOutput out) throws IOException {
     out.writeObject(context);
   }
 
+  @ParametersAreNonnullByDefault
   public static class InjectAdapter implements AgentPropagation.Setter<ContextPayload> {
     @Override
     public void set(final ContextPayload carrier, final String key, final String value) {

diff --git a/...ty/src/main/java/datadog/trace/civisibility/domain/buildsystem/BuildSystemModuleImpl.java b/...ty/src/main/java/datadog/trace/civisibility/domain/buildsystem/BuildSystemModuleImpl.java
@@ -1,8 +1,7 @@
 package datadog.trace.civisibility.domain.buildsystem;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
-
 import datadog.communication.ddagent.TracerVersion;
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.civisibility.CIConstants;
@@ -33,10 +32,13 @@
 import datadog.trace.util.Strings;
 import java.net.InetSocketAddress;
 import java.nio.file.Path;
-import java.util.*;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Consumer;
-import javax.annotation.Nullable;
 
 public class BuildSystemModuleImpl extends AbstractTestModule implements BuildSystemModule {
 
@@ -106,6 +108,7 @@ public <T extends CoverageCalculator> BuildSystemModuleImpl(
     setTag(Tags.TEST_COMMAND, startCommand);
   }
 
+  @ParametersAreNonnullByDefault
   private static final class ChildProcessPropertiesPropagationSetter
       implements AgentPropagation.Setter<Map<String, String>> {
     static final AgentPropagation.Setter<Map<String, String>> INSTANCE =
@@ -213,7 +216,7 @@ private Map<String, String> getPropertiesPropagatedToChildProcess(
     }
 
     // propagate module span context to child processes
-    propagate()
+    Propagators.defaultPropagator()
         .inject(span, propagatedSystemProperties, ChildProcessPropertiesPropagationSetter.INSTANCE);
 
     return propagatedSystemProperties;

diff --git a/.../src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java b/.../src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java
@@ -3,10 +3,12 @@
 import static datadog.opentelemetry.shim.trace.OtelSpanContext.fromRemote;
 import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;
 
+import datadog.context.propagation.Propagators;
 import datadog.opentelemetry.shim.context.OtelContext;
 import datadog.opentelemetry.shim.trace.OtelExtractedContext;
 import datadog.opentelemetry.shim.trace.OtelSpan;
 import datadog.trace.api.TracePropagationStyle;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext.Extracted;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -19,7 +21,6 @@
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.context.propagation.TextMapSetter;
-import java.util.Collection;
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 
@@ -36,11 +37,7 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
     if (carrier == null) {
       return;
     }
-    Span span = Span.fromContext(context);
-    if (span.getSpanContext().isValid()) {
-      AgentSpanContext agentSpanContext = OtelExtractedContext.extract(context);
-      AgentTracer.propagate().inject(agentSpanContext, carrier, setter::set);
-    }
+    Propagators.defaultPropagator().inject(convertContext(context), carrier, setter::set);
   }
 
   @Override
@@ -66,15 +63,23 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
     }
   }
 
+  private static datadog.context.Context convertContext(Context context) {
+    // TODO Extract baggage too
+    // TODO Create fast path from OtelSpan --> AgentSpan delegate --> with() to inflate as full
+    // context if baggage
+    AgentSpanContext extract = OtelExtractedContext.extract(context);
+    return AgentSpan.fromSpanContext(extract);
+  }
+
   /**
    * Extracts tracestate if {@code tracestate} header is present and extracted context comes from
    * {@link TracePropagationStyle#TRACECONTEXT}
    *
    * @param extracted The extracted context.
    * @param carrier The context carrier.
    * @param getter The context getter.
-   * @return The extracted tracestate, or an empty tracestate otherwise.
    * @param <C> The carrier type.
+   * @return The extracted tracestate, or an empty tracestate otherwise.
    */
   private static <C> TraceState extractTraceState(
       Extracted extracted, C carrier, TextMapGetter<C> getter) {

diff --git a/...a-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpan.java b/...a-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpan.java
@@ -172,7 +172,7 @@ public AgentSpanContext getAgentSpanContext() {
 
   @Override
   public AgentSpan asAgentSpan() {
-    return delegate;
+    return this.delegate;
   }
 
   private static class NoopSpan implements Span {

diff --git a/...http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpClientHelpers.java b/...http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpClientHelpers.java
@@ -35,7 +35,7 @@ public Void apply(final Try<HttpResponse> result) {
   public static class AkkaHttpHeaders implements AgentPropagation.Setter<HttpRequest> {
     private HttpRequest request;
     // Did this request have a span when the AkkaHttpHeaders object was created?
-    private boolean hadSpan;
+    private final boolean hadSpan;
 
     public AkkaHttpHeaders(final HttpRequest request) {
       hadSpan = request != null && request.getHeader(HasSpanHeader.class).isPresent();
@@ -51,6 +51,7 @@ public boolean hadSpan() {
       return hadSpan;
     }
 
+    @ParametersAreNonnullByDefault
     @Override
     public void set(final HttpRequest carrier, final String key, final String value) {
       // Coerce a Scala trait Self type into the correct type

diff --git a/...ain/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java b/...ain/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java
@@ -14,6 +14,7 @@
 import akka.http.scaladsl.model.HttpRequest;
 import akka.http.scaladsl.model.HttpResponse;
 import com.google.auto.service.AutoService;
+import datadog.context.propagation.Propagators;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -78,7 +79,7 @@ public static AgentScope methodEnter(
       DECORATE.onRequest(span, request);
 
       if (request != null) {
-        propagate().inject(span, request, headers);
+        Propagators.defaultPropagator().inject(span, request, headers);
         propagate()
             .injectPathwayContext(
                 span, request, headers, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

diff --git a/...p-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java b/...p-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java
@@ -7,6 +7,7 @@
 import akka.http.scaladsl.HttpExt;
 import akka.http.scaladsl.model.HttpRequest;
 import akka.http.scaladsl.model.HttpResponse;
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
@@ -28,7 +29,7 @@ public static AgentScope methodEnter(
     AkkaHttpClientDecorator.DECORATE.onRequest(span, request);
 
     if (request != null) {
-      propagate().inject(span, request, headers);
+      Propagators.defaultPropagator().inject(span, request, headers);
       propagate()
           .injectPathwayContext(
               span, request, headers, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

diff --git a/...n/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java b/...n/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java
@@ -4,6 +4,7 @@
 import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
 import java.io.IOException;
@@ -34,7 +35,7 @@ public HttpRequest generateRequest() throws IOException, HttpException {
     final HttpRequest request = delegate.generateRequest();
     DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
 
-    propagate().inject(span, request, SETTER);
+    Propagators.defaultPropagator().inject(span, request, SETTER);
     propagate()
         .injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
 

diff --git a/...in/java/datadog/trace/instrumentation/apachehttpasyncclient/HttpHeadersInjectAdapter.java b/...in/java/datadog/trace/instrumentation/apachehttpasyncclient/HttpHeadersInjectAdapter.java
@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.apachehttpasyncclient;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import javax.annotation.ParametersAreNonnullByDefault;
 import org.apache.http.HttpRequest;
 
+@ParametersAreNonnullByDefault
 public class HttpHeadersInjectAdapter implements AgentPropagation.Setter<HttpRequest> {
 
   public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();

diff --git a/...pclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java b/...pclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java
@@ -7,6 +7,7 @@
 import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -45,7 +46,7 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) {
 
     // AWS calls are often signed, so we can't add headers without breaking the signature.
     if (!awsClientCall) {
-      propagate().inject(span, request, SETTER);
+      Propagators.defaultPropagator().inject(span, request, SETTER);
       propagate()
           .injectPathwayContext(
               span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

diff --git a/...rc/main/java/datadog/trace/instrumentation/apachehttpclient/HttpHeadersInjectAdapter.java b/...rc/main/java/datadog/trace/instrumentation/apachehttpclient/HttpHeadersInjectAdapter.java
@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.apachehttpclient;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import javax.annotation.ParametersAreNonnullByDefault;
 import org.apache.http.client.methods.HttpUriRequest;
 
+@ParametersAreNonnullByDefault
 public class HttpHeadersInjectAdapter implements AgentPropagation.Setter<HttpUriRequest> {
 
   public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();

diff --git a/...c/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java b/...c/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java
@@ -4,6 +4,7 @@
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
 import java.io.IOException;
@@ -27,7 +28,7 @@ public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpCo
       throws HttpException, IOException {
     DECORATE.onRequest(span, request);
 
-    propagate().inject(span, request, SETTER);
+    Propagators.defaultPropagator().inject(span, request, SETTER);
     propagate()
         .injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
     delegate.sendRequest(request, entityDetails, context);

diff --git a/...client-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java b/...client-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java
@@ -7,6 +7,7 @@
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -45,7 +46,7 @@ private static AgentScope activateHttpSpan(final HttpRequest request) {
     final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
     // AWS calls are often signed, so we can't add headers without breaking the signature.
     if (!awsClientCall) {
-      propagate().inject(span, request, SETTER);
+      Propagators.defaultPropagator().inject(span, request, SETTER);
       propagate()
           .injectPathwayContext(
               span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

diff --git a/...c/main/java/datadog/trace/instrumentation/apachehttpclient5/HttpHeadersInjectAdapter.java b/...c/main/java/datadog/trace/instrumentation/apachehttpclient5/HttpHeadersInjectAdapter.java
@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.apachehttpclient5;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import javax.annotation.ParametersAreNonnullByDefault;
 import org.apache.hc.core5.http.HttpRequest;
 
+@ParametersAreNonnullByDefault
 public class HttpHeadersInjectAdapter implements AgentPropagation.Setter<HttpRequest> {
 
   public static final HttpHeadersInjectAdapter SETTER = new HttpHeadersInjectAdapter();

diff --git a/...java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java b/...java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java
@@ -16,6 +16,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
+import datadog.context.propagation.Propagators;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.muzzle.Reference;
@@ -120,7 +121,7 @@ public static <T> AgentScope before(
       if (null != responseListener && null != headers) {
         span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call);
         if (null != span) {
-          propagate().inject(span, headers, SETTER);
+          Propagators.defaultPropagator().inject(span, headers, SETTER);
           propagate().injectPathwayContext(span, headers, SETTER, CLIENT_PATHWAY_EDGE_TAGS);
           return activateSpan(span);
         }

diff --git a/...pc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcInjectAdapter.java b/...pc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcInjectAdapter.java
@@ -2,7 +2,9 @@
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import io.grpc.Metadata;
+import javax.annotation.ParametersAreNonnullByDefault;
 
+@ParametersAreNonnullByDefault
 public final class GrpcInjectAdapter implements AgentPropagation.Setter<Metadata> {
 
   public static final GrpcInjectAdapter SETTER = new GrpcInjectAdapter();

diff --git a/...rc/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java b/...rc/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java
@@ -8,6 +8,7 @@
 import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
 import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.InstanceStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
@@ -86,7 +87,7 @@ private String getTraceContextToInject(
     jsonBuilder.append('{');
 
     // Inject trace context
-    propagate().inject(span, jsonBuilder, SETTER);
+    Propagators.defaultPropagator().inject(span, jsonBuilder, SETTER);
 
     if (traceConfig().isDataStreamsEnabled()) {
       propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(eventBusName));

diff --git a/.../src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/TextMapInjectAdapter.java b/.../src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/TextMapInjectAdapter.java
@@ -1,7 +1,9 @@
 package datadog.trace.instrumentation.aws.v2.eventbridge;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import javax.annotation.ParametersAreNonnullByDefault;
 
+@ParametersAreNonnullByDefault
 public class TextMapInjectAdapter implements AgentPropagation.Setter<StringBuilder> {
 
   public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter();

diff --git a/...s-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java b/...s-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java
@@ -37,7 +37,12 @@ private ByteBuffer getMessageAttributeValueToInject(
     final AgentSpan span = newSpan(request);
     StringBuilder jsonBuilder = new StringBuilder();
     jsonBuilder.append('{');
-    propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG);
+    propagate()
+        .inject(
+            span,
+            jsonBuilder,
+            SETTER,
+            TracePropagationStyle.DATADOG); // TODO Is forcing Datadog encoding on purpose?
     if (traceConfig().isDataStreamsEnabled()) {
       propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(snsTopicName));
     }

diff --git a/...s-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/...s-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java
@@ -8,6 +8,7 @@
 import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
 import static datadog.trace.instrumentation.aws.v2.sns.TextMapInjectAdapter.SETTER;
 
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.InstanceStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.nio.charset.StandardCharsets;
@@ -37,7 +38,7 @@ private SdkBytes getMessageAttributeValueToInject(
     final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
     StringBuilder jsonBuilder = new StringBuilder();
     jsonBuilder.append('{');
-    propagate().inject(span, jsonBuilder, SETTER);
+    Propagators.defaultPropagator().inject(span, jsonBuilder, SETTER);
     if (traceConfig().isDataStreamsEnabled()) {
       propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(snsTopicName));
     }