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

Author: PerfectSlayer

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/httpurlconnection/HeadersInjectAdapter.java

@@ -2,7 +2,9 @@
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import java.net.HttpURLConnection;
+import javax.annotation.ParametersAreNonnullByDefault;
 
+@ParametersAreNonnullByDefault
 public class HeadersInjectAdapter implements AgentPropagation.Setter<HttpURLConnection> {
 
   public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter();

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ContextPayload.java

@@ -1,14 +1,14 @@
 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;
 import java.io.ObjectInput;
 import java.io.ObjectOutput;
 import java.util.HashMap;
 import java.util.Map;
+import javax.annotation.ParametersAreNonnullByDefault;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,7 +33,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 +54,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) {

dd-java-agent/agent-ci-visibility/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,15 @@
 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;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 public class BuildSystemModuleImpl extends AbstractTestModule implements BuildSystemModule {
 
@@ -105,6 +109,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 =
@@ -221,7 +226,7 @@ private Map<String, String> getPropertiesPropagatedToChildProcess(
     }
 
     // propagate module span context to child processes
-    propagate()
+    Propagators.defaultPropagator()
         .inject(span, propagatedSystemProperties, ChildProcessPropertiesPropagationSetter.INSTANCE);
 
     return propagatedSystemProperties;

dd-java-agent/agent-otel/otel-shim/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;
@@ -36,11 +38,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 +64,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) {

dd-java-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 {

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpClientHelpers.java

@@ -8,6 +8,7 @@
 import akka.http.scaladsl.model.headers.CustomHeader;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import javax.annotation.ParametersAreNonnullByDefault;
 import scala.runtime.AbstractFunction1;
 import scala.util.Try;
 
@@ -35,7 +36,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 +52,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

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/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);

dd-java-agent/instrumentation/akka-http/akka-http-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);

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/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);
 

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/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();