feat(dsm): Clean up context injection

Author: PerfectSlayer

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

@@ -1,7 +1,7 @@
 package datadog.trace.bootstrap.instrumentation.decorator;
 
 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;
@@ -35,7 +35,6 @@
 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;
@@ -59,15 +58,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");
   private static final boolean SHOULD_SET_404_RESOURCE_NAME =
@@ -150,7 +140,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;
   }

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy

@@ -36,7 +36,7 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
   protected static final int READ_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(5) as int
   protected static final BASIC_AUTH_KEY = "custom_authorization_header"
   protected static final BASIC_AUTH_VAL = "plain text auth token"
-  protected static final DSM_EDGE_TAGS = DataStreamsContext.client().sortedTags().collect { key, value ->
+  protected static final DSM_EDGE_TAGS = DataStreamsContext.forHttpClient().sortedTags().collect { key, value ->
     return key + ":" + value
   }
 

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy

@@ -14,6 +14,7 @@ import datadog.trace.api.DDTags
 import datadog.trace.api.ProductActivation
 import datadog.trace.api.config.GeneralConfig
 import datadog.trace.api.config.TracerConfig
+import datadog.trace.api.datastreams.DataStreamsContext
 import datadog.trace.api.env.CapturedEnvironment
 import datadog.trace.api.function.TriConsumer
 import datadog.trace.api.function.TriFunction
@@ -98,14 +99,13 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.get
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.SERVER_PATHWAY_EDGE_TAGS
 import static java.nio.charset.StandardCharsets.UTF_8
 import static org.junit.Assume.assumeTrue
 
 abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
 
   public static final Logger SERVER_LOGGER = LoggerFactory.getLogger("http-server")
-  protected static final DSM_EDGE_TAGS = SERVER_PATHWAY_EDGE_TAGS.collect {
+  protected static final DSM_EDGE_TAGS = DataStreamsContext.forHttpServer().sortedTags().collect {
     key, value ->
     return key + ":" + value
   }

dd-trace-core/src/main/java/datadog/trace/core/datastreams/DataStreamsPropagator.java

@@ -104,7 +104,7 @@ private static DataStreamsContext createDataStreamsContext(AgentSpan span) {
         && span.getTag(COMPONENT) instanceof CharSequence
         && "java-web-servlet-dispatcher"
             .equals(((CharSequence) span.getTag(COMPONENT)).toString())) {
-      return DataStreamsContext.client();
+      return DataStreamsContext.forHttpClient();
     }
     return null;
   }

internal-api/src/main/java/datadog/trace/api/datastreams/DataStreamsContext.java

@@ -9,6 +9,7 @@ public class DataStreamsContext implements ImplicitContextKeyed {
   private static final ContextKey<DataStreamsContext> CONTEXT_KEY =
       ContextKey.named("dsm-context-key");
   private static final LinkedHashMap<String, String> CLIENT_PATHWAY_EDGE_TAGS;
+  private static final LinkedHashMap<String, String> SERVER_PATHWAY_EDGE_TAGS;
 
   final LinkedHashMap<String, String> sortedTags;
   final long defaultTimestamp;
@@ -20,21 +21,34 @@ public class DataStreamsContext implements ImplicitContextKeyed {
     // 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");
+    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");
   }
 
   public static DataStreamsContext fromContext(Context context) {
     return context.get(CONTEXT_KEY);
   }
 
   /**
-   * Return default DSM context for HTTP clients.
+   * Gets default DSM context for HTTP clients.
    *
    * @return The default DSM context for HTTP clients.
    */
-  public static DataStreamsContext client() {
+  public static DataStreamsContext forHttpClient() {
     return fromTags(CLIENT_PATHWAY_EDGE_TAGS);
   }
 
+  /**
+   * Gets default DSM context for HTTP servers.
+   *
+   * @return The default DSM context for HTTP servers.
+   */
+  public static DataStreamsContext forHttpServer() {
+    return fromTags(SERVER_PATHWAY_EDGE_TAGS);
+  }
+
   /**
    * Creates a DSM context.
    *