Clean up DSM context injection

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpClientDecorator.java

@@ -26,21 +26,12 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.BitSet;
-import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.BiFunction;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends UriBasedClientDecorator {
-  public static final LinkedHashMap<String, String> CLIENT_PATHWAY_EDGE_TAGS;
-
-  static {
-    CLIENT_PATHWAY_EDGE_TAGS = new LinkedHashMap<>(2);
-    // TODO: Refactor TagsProcessor to move it into a package that we can link the constants for.
-    CLIENT_PATHWAY_EDGE_TAGS.put("direction", "out");
-    CLIENT_PATHWAY_EDGE_TAGS.put("type", "http");
-  }
 
   private static final Logger log = LoggerFactory.getLogger(HttpClientDecorator.class);
 

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -2,7 +2,7 @@
 
 import static datadog.context.Context.root;
 import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS;
-import static datadog.trace.api.datastreams.DataStreamsContext.fromTags;
+import static datadog.trace.api.datastreams.DataStreamsContext.forHttpServer;
 import static datadog.trace.api.gateway.Events.EVENTS;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
@@ -38,14 +38,14 @@
 import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver;
 import java.net.InetAddress;
 import java.util.BitSet;
-import java.util.LinkedHashMap;
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -62,15 +62,6 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
   public static final String DD_RESPONSE_ATTRIBUTE = "datadog.response";
   public static final String DD_IGNORE_COMMIT_ATTRIBUTE = "datadog.commit.ignore";
 
-  public static final LinkedHashMap<String, String> SERVER_PATHWAY_EDGE_TAGS;
-
-  static {
-    SERVER_PATHWAY_EDGE_TAGS = new LinkedHashMap<>(2);
-    // TODO: Refactor TagsProcessor to move it into a package that we can link the constants for.
-    SERVER_PATHWAY_EDGE_TAGS.put("direction", "in");
-    SERVER_PATHWAY_EDGE_TAGS.put("type", "http");
-  }
-
   private static final UTF8BytesString DEFAULT_RESOURCE_NAME = UTF8BytesString.create("/");
   protected static final UTF8BytesString NOT_FOUND_RESOURCE_NAME = UTF8BytesString.create("404");
   protected static final boolean SHOULD_SET_404_RESOURCE_NAME =
@@ -165,7 +156,7 @@ public AgentSpan startSpan(
     }
     AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
     if (null != carrier && null != getter) {
-      tracer().getDataStreamsMonitoring().setCheckpoint(span, fromTags(SERVER_PATHWAY_EDGE_TAGS));
+      tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
     }
     return span;
   }
@@ -191,7 +182,7 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final CONNECTION connection,
       final REQUEST request,
-      final AgentSpanContext.Extracted context) {
+      @Nullable final AgentSpanContext.Extracted context) {
     Config config = Config.get();
     boolean clientIpResolverEnabled =
         config.isClientIpEnabled()

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/UriBasedClientDecorator.java

@@ -1,13 +1,18 @@
 package datadog.trace.bootstrap.instrumentation.decorator;
 
+import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
+
+import datadog.context.Context;
+import datadog.context.propagation.CarrierSetter;
 import datadog.trace.api.Config;
+import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import java.net.URI;
 import javax.annotation.Nonnull;
 
 public abstract class UriBasedClientDecorator extends ClientDecorator {
-
   public void onURI(@Nonnull final AgentSpan span, @Nonnull final URI uri) {
     String host = uri.getHost();
     int port = uri.getPort();
@@ -21,4 +26,13 @@ public void onURI(@Nonnull final AgentSpan span, @Nonnull final URI uri) {
       }
     }
   }
+
+  public <C> void injectContext(Context context, final C request, CarrierSetter<C> setter) {
+    // Add additional default DSM context for HTTP clients if missing but DSM is enabled
+    if (DataStreamsContext.fromContext(context) == null && traceConfig().isDataStreamsEnabled()) {
+      context = context.with(DataStreamsContext.forHttpClient());
+    }
+    // Inject context into carrier
+    defaultPropagator().inject(context, request, setter);
+  }
 }

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

@@ -1,11 +1,9 @@
 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.startSpan;
 import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
 import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
 import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientHelpers.AkkaHttpHeaders;
@@ -18,7 +16,6 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
-import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
@@ -75,14 +72,12 @@ public static AgentScope methodEnter(
         return null;
       }
 
-      final AgentSpan span = startSpan(AKKA_CLIENT_REQUEST);
+      final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request);
 
       if (request != null) {
-        DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-        defaultPropagator()
-            .inject(getCurrentContext().with(span).with(dsmContext), request, headers);
+        DECORATE.injectContext(getCurrentContext().with(span), request, headers);
         // Request is immutable, so we have to assign new value once we update headers
         request = headers.getRequest();
       }

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

@@ -1,10 +1,11 @@
 package datadog.trace.instrumentation.akkahttp;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
+
 import akka.http.scaladsl.model.HttpRequest;
 import akka.http.scaladsl.model.HttpResponse;
 import akka.http.scaladsl.util.FastFuture$;
 import akka.stream.Materializer;
-import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -28,8 +29,7 @@ public DatadogAsyncHandlerWrapper(
   @Override
   public Future<HttpResponse> apply(final HttpRequest request) {
     final ContextScope scope = DatadogWrapperHelper.createSpan(request);
-    Context context = scope.context();
-    final AgentSpan span = AgentSpan.fromContext(context);
+    final AgentSpan span = fromContext(scope.context());
     Future<HttpResponse> futureResponse;
 
     // handle blocking in the beginning of the request

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

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.akkahttp;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 
 import akka.http.scaladsl.model.HttpRequest;
@@ -67,7 +68,7 @@ public GraphStageLogic createLogic(final Attributes inheritedAttributes) throws
               public void onPush() throws Exception {
                 final HttpRequest request = grab(requestInlet);
                 final ContextScope scope = DatadogWrapperHelper.createSpan(request);
-                final AgentSpan span = AgentSpan.fromContext(scope.context());
+                final AgentSpan span = fromContext(scope.context());
                 RequestContext requestContext = span.getRequestContext();
                 if (requestContext != null) {
                   HttpResponse response =
@@ -133,7 +134,7 @@ public void onPush() throws Exception {
                 HttpResponse response = grab(responseInlet);
                 final ContextScope scope = scopes.poll();
                 if (scope != null) {
-                  AgentSpan span = AgentSpan.fromContext(scope.context());
+                  AgentSpan span = fromContext(scope.context());
                   HttpResponse newResponse =
                       BlockingResponseHelper.handleFinishForWaf(span, response);
                   if (newResponse != response) {
@@ -159,7 +160,7 @@ public void onUpstreamFinish() throws Exception {
                 // remaining spans
                 ContextScope scope = scopes.poll();
                 while (scope != null) {
-                  AgentSpan.fromContext(scope.context()).finish();
+                  fromContext(scope.context()).finish();
                   scope = scopes.poll();
                 }
                 completeStage();
@@ -170,14 +171,14 @@ public void onUpstreamFailure(final Throwable ex) throws Exception {
                 ContextScope scope = scopes.poll();
                 if (scope != null) {
                   // Mark the span as failed
-                  AgentSpan span = AgentSpan.fromContext(scope.context());
+                  AgentSpan span = fromContext(scope.context());
                   DatadogWrapperHelper.finishSpan(span, ex);
                 }
                 // We will not receive any more responses from the user code, so clean up any
                 // remaining spans
                 scope = scopes.poll();
                 while (scope != null) {
-                  AgentSpan.fromContext(scope.context()).finish();
+                  fromContext(scope.context()).finish();
                   scope = scopes.poll();
                 }
                 fail(responseOutlet, ex);

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

@@ -10,12 +10,12 @@
 
 public class DatadogWrapperHelper {
   public static ContextScope createSpan(final HttpRequest request) {
-    final Context extractedContext = DECORATE.extractContext(request);
-    final AgentSpan span = DECORATE.startSpan(request, extractedContext);
+    final Context context = DECORATE.extractContext(request);
+    final AgentSpan span = DECORATE.startSpan(request, context);
     DECORATE.afterStart(span);
-    DECORATE.onRequest(span, request, request, extractedContext);
+    DECORATE.onRequest(span, request, request, context);
 
-    return extractedContext.with(span).attach();
+    return context.with(span).attach();
   }
 
   public static void finishSpan(final AgentSpan span, final HttpResponse response) {

dd-java-agent/instrumentation/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java

@@ -1,15 +1,14 @@
 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.startSpan;
 import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
+import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
+import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.DECORATE;
 
 import akka.http.scaladsl.HttpExt;
 import akka.http.scaladsl.model.HttpRequest;
 import akka.http.scaladsl.model.HttpResponse;
-import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
@@ -25,13 +24,11 @@ public static AgentScope methodEnter(
       return null;
     }
 
-    final AgentSpan span = startSpan(AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST);
-    AkkaHttpClientDecorator.DECORATE.afterStart(span);
-    AkkaHttpClientDecorator.DECORATE.onRequest(span, request);
-
+    final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST);
+    DECORATE.afterStart(span);
+    DECORATE.onRequest(span, request);
     if (request != null) {
-      DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-      defaultPropagator().inject(getCurrentContext().with(span).with(dsmContext), request, headers);
+      DECORATE.injectContext(getCurrentContext().with(span), request, headers);
       // Request is immutable, so we have to assign new value once we update headers
       request = headers.getRequest();
     }
@@ -54,8 +51,8 @@ public static void methodExit(
       responseFuture.onComplete(
           new AkkaHttpClientHelpers.OnCompleteHandler(span), thiz.system().dispatcher());
     } else {
-      AkkaHttpClientDecorator.DECORATE.onError(span, throwable);
-      AkkaHttpClientDecorator.DECORATE.beforeFinish(span);
+      DECORATE.onError(span, throwable);
+      DECORATE.beforeFinish(span);
       span.finish();
     }
     scope.close();

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java

@@ -1,12 +1,9 @@
 package datadog.trace.instrumentation.apachehttpasyncclient;
 
-import static datadog.context.propagation.Propagators.defaultPropagator;
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
+import static datadog.context.Context.current;
 import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.context.Context;
-import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.io.IOException;
 import org.apache.http.HttpException;
@@ -35,10 +32,7 @@ public HttpHost getTarget() {
   public HttpRequest generateRequest() throws IOException, HttpException {
     final HttpRequest request = delegate.generateRequest();
     DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
-
-    DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-    defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER);
-
+    DECORATE.injectContext(current().with(span), request, SETTER);
     return request;
   }
 

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

@@ -1,15 +1,12 @@
 package datadog.trace.instrumentation.apachehttpclient;
 
-import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.context.Context.current;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
 import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.context.Context;
-import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -47,8 +44,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) {
-      DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-      defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER);
+      DECORATE.injectContext(current().with(span), request, SETTER);
     }
 
     return scope;

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

@@ -1,12 +1,9 @@
 package datadog.trace.instrumentation.apachehttpclient5;
 
-import static datadog.context.propagation.Propagators.defaultPropagator;
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
+import static datadog.context.Context.current;
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.context.Context;
-import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.io.IOException;
 import org.apache.hc.core5.http.EntityDetails;
@@ -28,9 +25,7 @@ public DelegatingRequestChannel(RequestChannel requestChannel, AgentSpan span) {
   public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpContext context)
       throws HttpException, IOException {
     DECORATE.onRequest(span, request);
-
-    DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-    defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER);
+    DECORATE.injectContext(current().with(span), request, SETTER);
     delegate.sendRequest(request, entityDetails, context);
   }
 }

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

@@ -1,15 +1,12 @@
 package datadog.trace.instrumentation.apachehttpclient5;
 
-import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.context.Context.current;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
-import datadog.context.Context;
-import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -47,8 +44,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) {
-      DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-      defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER);
+      DECORATE.injectContext(current().with(span), request, SETTER);
     }
 
     return scope;