Support Context.with(span) for convenience

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Constants.java

@@ -15,6 +15,7 @@ public final class Constants {
    */
   public static final String[] BOOTSTRAP_PACKAGE_PREFIXES = {
     "datadog.slf4j",
+    "datadog.context",
     "datadog.appsec.api",
     "datadog.trace.api",
     "datadog.trace.bootstrap",

dd-java-agent/build.gradle

@@ -135,6 +135,7 @@ def sharedShadowJar = tasks.register('sharedShadowJar', ShadowJar) {
     exclude(project(':dd-java-agent:agent-logging'))
     exclude(project(':dd-trace-api'))
     exclude(project(':internal-api'))
+    exclude(project(':components:context'))
     exclude(project(':utils:time-utils'))
     exclude(dependency('org.slf4j::'))
   }

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

@@ -38,6 +38,7 @@ public class SpockRunner extends JUnitPlatform {
    */
   public static final String[] BOOTSTRAP_PACKAGE_PREFIXES_COPY = {
     "datadog.slf4j",
+    "datadog.context",
     "datadog.appsec.api",
     "datadog.trace.api",
     "datadog.trace.bootstrap",

gradle/dependencies.gradle

@@ -16,6 +16,7 @@ final class CachedData {
       exclude(project(':internal-api'))
       exclude(project(':internal-api:internal-api-9'))
       exclude(project(':communication'))
+      exclude(project(':components:context'))
       exclude(project(':components:json'))
       exclude(project(':remote-config:remote-config-api'))
       exclude(project(':remote-config:remote-config-core'))

internal-api/build.gradle

@@ -210,6 +210,7 @@ dependencies {
   // references TraceScope and Continuation from public api
   api project(':dd-trace-api')
   api libs.slf4j
+  api project(':components:context')
   api project(":utils:time-utils")
 
   // has to be loaded by system classloader:

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -1,12 +1,16 @@
 package datadog.trace.bootstrap.instrumentation.api;
 
+import static datadog.trace.bootstrap.instrumentation.api.InternalContextKeys.SPAN_KEY;
+
+import datadog.context.Context;
+import datadog.context.ImplicitContextKeyed;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.gateway.IGSpanInfo;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.interceptor.MutableSpan;
 
-public interface AgentSpan extends MutableSpan, IGSpanInfo, WithAgentSpan {
+public interface AgentSpan extends MutableSpan, ImplicitContextKeyed, IGSpanInfo, WithAgentSpan {
 
   DDTraceId getTraceId();
 
@@ -145,4 +149,9 @@ public interface AgentSpan extends MutableSpan, IGSpanInfo, WithAgentSpan {
   default AgentSpan asAgentSpan() {
     return this;
   }
+
+  @Override
+  default Context storeInto(Context context) {
+    return context.with(SPAN_KEY, this);
+  }
 }

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/InternalContextKeys.java

@@ -0,0 +1,7 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+import datadog.context.ContextKey;
+
+final class InternalContextKeys {
+  static final ContextKey<AgentSpan> SPAN_KEY = ContextKey.named("dd-span-key");
+}