Migrating all HttpServer Instrumentations to Extract full Context

dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorBenchmark.java

@@ -6,6 +6,7 @@
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.ContextVisitors;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
 import datadog.trace.bootstrap.instrumentation.api.URIDefaultDataAdapter;
@@ -61,7 +62,7 @@ public void setUp() {
 
   @Benchmark
   public AgentSpan onRequest() {
-    return decorator.onRequest(span, null, request, null);
+    return decorator.onRequest(span, null, request, (AgentSpanContext.Extracted) null);
   }
 
   public static class Request {

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge.java

@@ -29,4 +29,9 @@ public static Context getContextFrom(Object carrier) {
   public static Context detachContextFrom(Object carrier) {
     return Context.detachFrom(carrier);
   }
+
+  /** @see AgentSpan#fromContext(Context) */
+  public static AgentSpan spanFromContext(Context context) {
+    return AgentSpan.fromContext(context);
+  }
 }

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

@@ -3,6 +3,7 @@
 import static datadog.trace.api.cache.RadixTreeCache.PORTS;
 import static datadog.trace.api.cache.RadixTreeCache.UNSET_PORT;
 
+import datadog.context.ContextScope;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.Functions;
@@ -105,6 +106,11 @@ public AgentSpan onError(final AgentSpan span, final Throwable throwable, byte e
     return span;
   }
 
+  public ContextScope onError(final ContextScope scope, final Throwable throwable) {
+    onError(AgentSpan.fromContext(scope.context()), throwable);
+    return scope;
+  }
+
   public AgentSpan onPeerConnection(
       final AgentSpan span, final InetSocketAddress remoteConnection) {
     if (remoteConnection != null) {

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

@@ -1,7 +1,5 @@
 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;
@@ -11,6 +9,7 @@
 
 import datadog.appsec.api.blocking.BlockingException;
 import datadog.context.Context;
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.function.TriConsumer;
@@ -136,18 +135,15 @@ public AgentSpanContext.Extracted extract(REQUEST_CARRIER carrier) {
     return extractContextAndGetSpanContext(carrier, getter);
   }
 
-  /**
-   * 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 null;
     }
-    return defaultPropagator().extract(root(), carrier, getter);
+    return Propagators.defaultPropagator().extract(Context.root(), carrier, getter);
   }
 
-  /** Deprecated. Use {@link #startSpanFromContext(String, Object, Context)} instead. */
+  /** Deprecated. Use {@link #startSpanFromContext(Object, Context)} instead. */
   @Deprecated
   public AgentSpan startSpan(REQUEST_CARRIER carrier, AgentSpanContext.Extracted context) {
     return startSpan("http-server", carrier, context);
@@ -170,20 +166,23 @@ 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 AgentSpan startSpanFromContext(REQUEST_CARRIER carrier, Context context) {
+    return startSpan("http-server", carrier, getExtractedSpanContext(context));
   }
 
-  public AgentSpanContext.Extracted getSpanContext(Context context) {
+  public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
     AgentSpan extractedSpan = AgentSpan.fromContext(context);
     return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
   }
 
+  public AgentSpan onRequest(
+      final AgentSpan span,
+      final CONNECTION connection,
+      final REQUEST request,
+      final Context context) {
+    return onRequest(span, connection, request, getExtractedSpanContext(context));
+  }
+
   public AgentSpan onRequest(
       final AgentSpan span,
       final CONNECTION connection,

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy

@@ -51,7 +51,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
     def decorator = newDecorator()
 
     when:
-    decorator.onRequest(this.span, null, req, null)
+    decorator.onRequest(this.span, null, req, (AgentSpanContext.Extracted) null)
 
     then:
     if (req) {
@@ -84,7 +84,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
     def decorator = newDecorator()
 
     when:
-    decorator.onRequest(this.span, null, req, null)
+    decorator.onRequest(this.span, null, req, (AgentSpanContext.Extracted) null)
 
     then:
     if (expectedUrl) {
@@ -135,7 +135,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
     def decorator = newDecorator()
 
     when:
-    decorator.onRequest(this.span, null, req, null)
+    decorator.onRequest(this.span, null, req, (AgentSpanContext.Extracted) null)
 
     then:
     1 * this.span.setTag(Tags.HTTP_URL, {it.toString() == expectedUrl})
@@ -165,7 +165,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
     def decorator = newDecorator()
 
     when:
-    decorator.onRequest(this.span, null, [url: new URI('http://host/p%20ath')], null)
+    decorator.onRequest(this.span, null, [url: new URI('http://host/p%20ath')], (AgentSpanContext.Extracted) null)
 
     then:
     1 * this.span.setResourceName({ it as String == '/path' }, ResourceNamePriorities.HTTP_PATH_NORMALIZER)

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

@@ -4,8 +4,9 @@
 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.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.instrumentation.akkahttp.appsec.BlockingResponseHelper;
 import scala.Function1;
@@ -26,8 +27,9 @@ public DatadogAsyncHandlerWrapper(
 
   @Override
   public Future<HttpResponse> apply(final HttpRequest request) {
-    final AgentScope scope = DatadogWrapperHelper.createSpan(request);
-    AgentSpan span = scope.span();
+    final ContextScope scope = DatadogWrapperHelper.createSpan(request);
+    Context context = scope.context();
+    final AgentSpan span = AgentSpan.fromContext(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

@@ -13,8 +13,8 @@
 import akka.stream.stage.AbstractOutHandler;
 import akka.stream.stage.GraphStage;
 import akka.stream.stage.GraphStageLogic;
+import datadog.context.ContextScope;
 import datadog.trace.api.gateway.RequestContext;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.instrumentation.akkahttp.appsec.BlockingResponseHelper;
 import java.util.Queue;
@@ -56,7 +56,7 @@ public GraphStageLogic createLogic(final Attributes inheritedAttributes) throws
         // that this connection was created with. This means that we can safely
         // close the span at the front of the queue when we receive the response
         // from the user code, since it will match up to the request for that span.
-        final Queue<AgentScope> scopes = new ArrayBlockingQueue<>(pipeliningLimit);
+        final Queue<ContextScope> scopes = new ArrayBlockingQueue<>(pipeliningLimit);
         boolean[] skipNextPull = new boolean[] {false};
 
         // This is where the request comes in from the server and TCP layer
@@ -66,8 +66,8 @@ public GraphStageLogic createLogic(final Attributes inheritedAttributes) throws
               @Override
               public void onPush() throws Exception {
                 final HttpRequest request = grab(requestInlet);
-                final AgentScope scope = DatadogWrapperHelper.createSpan(request);
-                AgentSpan span = scope.span();
+                final ContextScope scope = DatadogWrapperHelper.createSpan(request);
+                final AgentSpan span = AgentSpan.fromContext(scope.context());
                 RequestContext requestContext = span.getRequestContext();
                 if (requestContext != null) {
                   HttpResponse response =
@@ -77,7 +77,7 @@ public void onPush() throws Exception {
                     skipNextPull[0] = true;
                     requestContext.getTraceSegment().effectivelyBlocked();
                     emit(responseOutlet, response);
-                    DatadogWrapperHelper.finishSpan(scope.span(), response);
+                    DatadogWrapperHelper.finishSpan(span, response);
                     pull(requestInlet);
                     scope.close();
                     return;
@@ -131,9 +131,9 @@ public void onDownstreamFinish() throws Exception {
               @Override
               public void onPush() throws Exception {
                 HttpResponse response = grab(responseInlet);
-                final AgentScope scope = scopes.poll();
+                final ContextScope scope = scopes.poll();
                 if (scope != null) {
-                  AgentSpan span = scope.span();
+                  AgentSpan span = AgentSpan.fromContext(scope.context());
                   HttpResponse newResponse =
                       BlockingResponseHelper.handleFinishForWaf(span, response);
                   if (newResponse != response) {
@@ -146,7 +146,7 @@ public void onPush() throws Exception {
                   // and close it. If it's not, then it will be cleaned up actor message
                   // processing instrumentation that drives this state machine
                   AgentSpan activeSpan = activeSpan();
-                  if (activeSpan == scope.span()) {
+                  if (activeSpan == span) {
                     scope.close();
                   }
                 }
@@ -157,26 +157,27 @@ public void onPush() throws Exception {
               public void onUpstreamFinish() throws Exception {
                 // We will not receive any more responses from the user code, so clean up any
                 // remaining spans
-                AgentScope scope = scopes.poll();
+                ContextScope scope = scopes.poll();
                 while (scope != null) {
-                  scope.span().finish();
+                  AgentSpan.fromContext(scope.context()).finish();
                   scope = scopes.poll();
                 }
                 completeStage();
               }
 
               @Override
               public void onUpstreamFailure(final Throwable ex) throws Exception {
-                AgentScope scope = scopes.poll();
+                ContextScope scope = scopes.poll();
+                AgentSpan span = AgentSpan.fromContext(scope.context());
                 if (scope != null) {
                   // Mark the span as failed
-                  DatadogWrapperHelper.finishSpan(scope.span(), ex);
+                  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) {
-                  scope.span().finish();
+                  span.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

@@ -1,22 +1,21 @@
 package datadog.trace.instrumentation.akkahttp;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.instrumentation.akkahttp.AkkaHttpServerDecorator.DECORATE;
 
 import akka.http.scaladsl.model.HttpRequest;
 import akka.http.scaladsl.model.HttpResponse;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.context.Context;
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 
 public class DatadogWrapperHelper {
-  public static AgentScope createSpan(final HttpRequest request) {
-    final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request);
-    final AgentSpan span = DECORATE.startSpan(request, extractedContext);
+  public static ContextScope createSpan(final HttpRequest request) {
+    final Context extractedContext = DECORATE.extractContext(request);
+    final AgentSpan span = DECORATE.startSpanFromContext(request, extractedContext);
     DECORATE.afterStart(span);
     DECORATE.onRequest(span, request, request, extractedContext);
 
-    return activateSpan(span);
+    return extractedContext.with(span).attach();
   }
 
   public static void finishSpan(final AgentSpan span, final HttpResponse response) {

dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java

@@ -8,6 +8,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import net.bytebuddy.asm.Advice;
 
 public class HTTPPluginAdvice {
@@ -17,7 +18,8 @@ public static AgentScope onEnter(@Advice.Argument(value = 2) final Object server
     final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true);
     DECORATE.afterStart(span);
     // serverTransaction is like request + connection in one object:
-    DECORATE.onRequest(span, serverTransaction, serverTransaction, null);
+    DECORATE.onRequest(
+        span, serverTransaction, serverTransaction, (AgentSpanContext.Extracted) null);
     final AgentScope scope = activateSpan(span);
     return scope;
   }

dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java

@@ -3,7 +3,7 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.isAnnotatedWith;
 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.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
 import static datadog.trace.instrumentation.azurefunctions.AzureFunctionsDecorator.DECORATE;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -14,11 +14,11 @@
 import com.microsoft.azure.functions.ExecutionContext;
 import com.microsoft.azure.functions.HttpRequestMessage;
 import com.microsoft.azure.functions.HttpResponseMessage;
+import datadog.context.Context;
+import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -64,24 +64,24 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static class AzureFunctionsAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static AgentScope methodEnter(
+    public static ContextScope methodEnter(
         @Advice.Argument(0) final HttpRequestMessage request,
         @Advice.Argument(1) final ExecutionContext context) {
-      final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request);
-      final AgentSpan span = DECORATE.startSpan(request, extractedContext);
+      final Context extractedContext = DECORATE.extractContext(request);
+      final AgentSpan span = DECORATE.startSpanFromContext(request, extractedContext);
       DECORATE.afterStart(span, context.getFunctionName());
       DECORATE.onRequest(span, request, request, extractedContext);
       HTTP_RESOURCE_DECORATOR.withRoute(
           span, request.getHttpMethod().name(), request.getUri().getPath());
-      return activateSpan(span);
+      return extractedContext.with(span).attach();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
-        @Advice.Enter final AgentScope scope,
+        @Advice.Enter final ContextScope scope,
         @Advice.Return final HttpResponseMessage response,
         @Advice.Thrown final Throwable throwable) {
-      final AgentSpan span = scope.span();
+      final AgentSpan span = spanFromContext(scope.context());
       DECORATE.onError(span, throwable);
       DECORATE.onResponse(span, response);
       DECORATE.beforeFinish(span);