Migrate context injection calls to context-first APIs (#8358)

* feat(core): Refactor propagation calls to the new inject API
* feat(core): Migrate X-Ray inject calls to new API
* feat(core): Implement standalone ASM propagator
* feat(aws-sns): Switch AWS SNS to default propagator (like AWS SNS 2)
* chore(ci): Increase limit of Muzzle

.circleci/config.continue.yml.j2

@@ -767,7 +767,7 @@ jobs:
 
   muzzle:
     <<: *defaults
-    resource_class: medium
+    resource_class: medium+
     parallelism: 4
     steps:
       - setup_code
@@ -798,7 +798,7 @@ jobs:
           name: Verify Muzzle
           command: >-
             SKIP_BUILDSCAN="true"
-            GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xmx2G -Xms2G -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp'"
+            GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xmx3G -Xms2G -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp'"
             ./gradlew `circleci tests split --split-by=timings workspace/build/muzzleTasks | xargs`
             << pipeline.parameters.gradle_flags >>
             --max-workers=4

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,6 +1,6 @@
 package datadog.trace.bootstrap.instrumentation.rmi;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
+import static datadog.context.propagation.Propagators.defaultPropagator;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -9,6 +9,7 @@
 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 +34,7 @@ public Map<String, String> getContext() {
 
   public static ContextPayload from(final AgentSpan span) {
     final ContextPayload payload = new ContextPayload();
-    propagate().inject(span, payload, SETTER);
+    defaultPropagator().inject(span, payload, SETTER);
     return payload;
   }
 
@@ -54,6 +55,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,6 +1,6 @@
 package datadog.trace.civisibility.domain.buildsystem;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
+import static datadog.context.propagation.Propagators.defaultPropagator;
 
 import datadog.communication.ddagent.TracerVersion;
 import datadog.trace.api.Config;
@@ -33,10 +33,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 +110,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 +227,7 @@ private Map<String, String> getPropertiesPropagatedToChildProcess(
     }
 
     // propagate module span context to child processes
-    propagate()
+    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

@@ -1,12 +1,14 @@
 package datadog.opentelemetry.shim.context.propagation;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.opentelemetry.shim.trace.OtelSpanContext.fromRemote;
 import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;
 
 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);
-    }
+    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

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.akkahttp;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
@@ -78,7 +79,7 @@ public static AgentScope methodEnter(
       DECORATE.onRequest(span, request);
 
       if (request != null) {
-        propagate().inject(span, request, headers);
+        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

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.akkahttp106;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
@@ -28,7 +29,7 @@ public static AgentScope methodEnter(
     AkkaHttpClientDecorator.DECORATE.onRequest(span, request);
 
     if (request != null) {
-      propagate().inject(span, request, headers);
+      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

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.apachehttpasyncclient;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
@@ -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);
+    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();

dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.apachehttpclient;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
@@ -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);
+      defaultPropagator().inject(span, request, SETTER);
       propagate()
           .injectPathwayContext(
               span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

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

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.apachehttpclient5;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
@@ -27,7 +28,7 @@ public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpCo
       throws HttpException, IOException {
     DECORATE.onRequest(span, request);
 
-    propagate().inject(span, request, SETTER);
+    defaultPropagator().inject(span, request, SETTER);
     propagate()
         .injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
     delegate.sendRequest(request, entityDetails, context);

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.apachehttpclient5;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
@@ -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);
+      defaultPropagator().inject(span, request, SETTER);
       propagate()
           .injectPathwayContext(
               span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);

dd-java-agent/instrumentation/apache-httpclient-5/src/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();

dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.armeria.grpc.client;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
@@ -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);
+          defaultPropagator().inject(span, headers, SETTER);
           propagate().injectPathwayContext(span, headers, SETTER, CLIENT_PATHWAY_EDGE_TAGS);
           return activateSpan(span);
         }

dd-java-agent/instrumentation/armeria-grpc/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();

dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.aws.v2.eventbridge;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
 import static datadog.trace.core.datastreams.TagsProcessor.BUS_TAG;
@@ -86,7 +87,7 @@ private String getTraceContextToInject(
     jsonBuilder.append('{');
 
     // Inject trace context
-    propagate().inject(span, jsonBuilder, SETTER);
+    defaultPropagator().inject(span, jsonBuilder, SETTER);
 
     if (traceConfig().isDataStreamsEnabled()) {
       propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(eventBusName));

dd-java-agent/instrumentation/aws-java-eventbridge-2.0/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();

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java

@@ -28,6 +28,7 @@
 import java.util.Locale;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response>
     implements AgentPropagation.Setter<Request<?>> {
@@ -332,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);