Supporting Baggage for Instrumentations used in Weblog Tests

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

@@ -1,5 +1,7 @@
 package datadog.trace.bootstrap.instrumentation.decorator;
 
+import static datadog.context.Context.root;
+import static datadog.context.propagation.Propagators.defaultPropagator;
 import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS;
 import static datadog.trace.api.datastreams.DataStreamsContext.fromTags;
 import static datadog.trace.api.gateway.Events.EVENTS;
@@ -8,6 +10,7 @@
 import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
 
 import datadog.appsec.api.blocking.BlockingException;
+import datadog.context.Context;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.function.TriConsumer;
@@ -124,6 +127,7 @@ protected AgentTracer.TracerAPI tracer() {
     return AgentTracer.get();
   }
 
+  /** Deprecated. Use {@link #extractContext(REQUEST_CARRIER)} instead. */
   public AgentSpanContext.Extracted extract(REQUEST_CARRIER carrier) {
     AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
     if (null == carrier || null == getter) {
@@ -132,7 +136,18 @@ public AgentSpanContext.Extracted extract(REQUEST_CARRIER carrier) {
     return extractContextAndGetSpanContext(carrier, getter);
   }
 
-  /** Deprecated. Use {@link #startSpan(String, Object, AgentSpanContext.Extracted)} instead. */
+  /**
+   * Will be renamed to #extract(REQUEST_CARRIER) when refactoring of instrumentation's is complete
+   */
+  public Context extractContext(REQUEST_CARRIER carrier) {
+    AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
+    if (null == carrier || null == getter) {
+      return root();
+    }
+    return defaultPropagator().extract(root(), carrier, getter);
+  }
+
+  /** Deprecated. Use {@link #startSpanFromContext(String, Object, Context)} instead. */
   @Deprecated
   public AgentSpan startSpan(REQUEST_CARRIER carrier, AgentSpanContext.Extracted context) {
     return startSpan("http-server", carrier, context);
@@ -155,6 +170,20 @@ public AgentSpan startSpan(
     return span;
   }
 
+  /**
+   * Will be renamed to #startSpan(String, REQUEST_CARRIER, Context) when refactoring of
+   * instrumentation's is complete
+   */
+  public AgentSpan startSpanFromContext(
+      String instrumentationName, REQUEST_CARRIER carrier, Context context) {
+    return startSpan(instrumentationName, carrier, getSpanContext(context));
+  }
+
+  public AgentSpanContext.Extracted getSpanContext(Context context) {
+    AgentSpan extractedSpan = AgentSpan.fromContext(context);
+    return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
+  }
+
   public AgentSpan onRequest(
       final AgentSpan span,
       final CONNECTION connection,

dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/TracingInterceptor.java

@@ -8,6 +8,7 @@
 import static datadog.trace.instrumentation.okhttp3.OkHttpClientDecorator.OKHTTP_REQUEST;
 import static datadog.trace.instrumentation.okhttp3.RequestBuilderInjectAdapter.SETTER;
 
+import datadog.context.Context;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -31,7 +32,8 @@ public Response intercept(final Chain chain) throws IOException {
 
       final Request.Builder requestBuilder = chain.request().newBuilder();
       DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
-      defaultPropagator().inject(span.with(dsmContext), requestBuilder, SETTER);
+      defaultPropagator()
+          .inject(Context.current().with(span).with(dsmContext), requestBuilder, SETTER);
 
       final Response response;
       try {

dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java

@@ -12,14 +12,15 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
+import datadog.context.Context;
+import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.ExcludeFilterProvider;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.muzzle.Reference;
 import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.gateway.Flow;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -114,19 +115,22 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class ServiceAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static AgentScope onService(@Advice.Argument(0) org.apache.coyote.Request req) {
+    public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Request req) {
 
       Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE);
       if (existingSpan instanceof AgentSpan) {
         // Request already gone through initial processing, so just activate the span.
         return activateSpan((AgentSpan) existingSpan);
       }
 
-      final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req);
-      req.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
+      final Context extractedContext = DECORATE.extractContext(req);
+      // TODO: Migrate setting DD_EXTRACTED_CONTEXT_ATTRIBUTE from AgentSpanContext.Extracted to
+      // Context
+      req.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, DECORATE.getSpanContext(extractedContext));
+
+      final AgentSpan span = DECORATE.startSpanFromContext("http-server", req, extractedContext);
+      final ContextScope scope = extractedContext.with(span).attach();
 
-      final AgentSpan span = DECORATE.startSpan(req, extractedContext);
-      final AgentScope scope = activateSpan(span);
       // This span is finished when Request.recycle() is called by RequestInstrumentation.
       DECORATE.afterStart(span);
 
@@ -137,7 +141,7 @@ public static AgentScope onService(@Advice.Argument(0) org.apache.coyote.Request
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void closeScope(@Advice.Enter final AgentScope scope) {
+    public static void closeScope(@Advice.Enter final ContextScope scope) {
       scope.close();
     }