Support Context.with(span) for convenience

Author: mcculls

internal-api/build.gradle

@@ -207,6 +207,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.AgentTracer.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/AgentTracer.java

@@ -3,6 +3,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_ASYNC_PROPAGATING;
 import static java.util.Collections.emptyList;
 
+import datadog.context.ContextKey;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.EndpointCheckpointer;
@@ -33,6 +34,8 @@
 public class AgentTracer {
   private static final String DEFAULT_INSTRUMENTATION_NAME = "datadog";
 
+  static final ContextKey<AgentSpan> SPAN_KEY = ContextKey.named("dd-span-key");
+
   // Implicit parent
   /** Deprecated. Use {@link #startSpan(String, CharSequence)} instead. */
   @Deprecated