Improve Instrumenter API to use Context instead of Span

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

@@ -1,5 +1,6 @@
 package datadog.trace.bootstrap.instrumentation.decorator;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
@@ -58,7 +59,8 @@ public void setUp() {
             .build();
     GlobalTracer.forceRegister(tracer);
     decorator = new BenchmarkHttpServerDecorator();
-    span = decorator.startSpan(Collections.emptyMap(), (Context) null);
+    Context context = decorator.startSpan(Collections.emptyMap(), Context.root());
+    span = fromContext(context);
   }
 
   @Benchmark
@@ -99,7 +101,7 @@ public BenchmarkHttpServerDecorator() {
 
     @Override
     protected String[] instrumentationNames() {
-      return new String[0];
+      return new String[] {"benchmark"};
     }
 
     @Override

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

@@ -118,6 +118,12 @@ protected AgentTracer.TracerAPI tracer() {
     return AgentTracer.get();
   }
 
+  /**
+   * Extracts context from an upstream service.
+   *
+   * @param carrier The request carrier to get the context from.
+   * @return The extracted context, {@code Context#root()} if no valid context to extract.
+   */
   public Context extract(REQUEST_CARRIER carrier) {
     AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
     if (null == carrier || null == getter) {
@@ -126,12 +132,19 @@ public Context extract(REQUEST_CARRIER carrier) {
     return Propagators.defaultPropagator().extract(root(), carrier, getter);
   }
 
-  public AgentSpan startSpan(
-      String instrumentationName, REQUEST_CARRIER carrier, AgentSpanContext.Extracted context) {
+  /**
+   * Starts a span.
+   *
+   * @param carrier The request carrier.
+   * @param context The parent context of the span to create.
+   * @return A new context bundling the span, child of the given parent context.
+   */
+  public Context startSpan(REQUEST_CARRIER carrier, Context context) {
+    String instrumentationName =
+        instrumentationNames().length == 0 ? "http-server" : instrumentationNames()[0];
+    AgentSpanContext.Extracted extracted = callIGCallbackStart(context);
     AgentSpan span =
-        tracer()
-            .startSpan(instrumentationName, spanName(), callIGCallbackStart(context))
-            .setMeasured(true);
+        tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
     Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
     if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
       span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
@@ -140,16 +153,7 @@ public AgentSpan startSpan(
     if (null != carrier && null != getter) {
       tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
     }
-    return span;
-  }
-
-  public AgentSpan startSpan(REQUEST_CARRIER carrier, Context context) {
-    return startSpan("http-server", carrier, getExtractedSpanContext(context));
-  }
-
-  public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
-    AgentSpan extractedSpan = AgentSpan.fromContext(context);
-    return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
+    return context.with(span);
   }
 
   public AgentSpan onRequest(
@@ -160,6 +164,11 @@ public AgentSpan onRequest(
     return onRequest(span, connection, request, getExtractedSpanContext(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,
@@ -375,34 +384,23 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
     return span;
   }
 
-  //  @Override
-  //  public Span onError(final Span span, final Throwable throwable) {
-  //    assert span != null;
-  //    // FIXME
-  //    final Object status = span.getTag("http.status");
-  //    if (status == null || status.equals(200)) {
-  //      // Ensure status set correctly
-  //      span.setTag("http.status", 500);
-  //    }
-  //    return super.onError(span, throwable);
-  //  }
-
-  private AgentSpanContext.Extracted callIGCallbackStart(AgentSpanContext.Extracted context) {
+  private AgentSpanContext.Extracted callIGCallbackStart(final Context context) {
+    AgentSpanContext.Extracted extracted = getExtractedSpanContext(context);
     AgentTracer.TracerAPI tracer = tracer();
     Supplier<Flow<Object>> startedCbAppSec =
         tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());
     Supplier<Flow<Object>> startedCbIast =
         tracer.getCallbackProvider(RequestContextSlot.IAST).getCallback(EVENTS.requestStarted());
 
     if (startedCbAppSec == null && startedCbIast == null) {
-      return context;
+      return extracted;
     }
 
     TagContext tagContext = null;
-    if (context == null) {
+    if (extracted == null) {
       tagContext = new TagContext();
-    } else if (context instanceof TagContext) {
-      tagContext = (TagContext) context;
+    } else if (extracted instanceof TagContext) {
+      tagContext = (TagContext) extracted;
     }
 
     if (tagContext != null) {
@@ -415,7 +413,7 @@ private AgentSpanContext.Extracted callIGCallbackStart(AgentSpanContext.Extracte
       return tagContext;
     }
 
-    return context;
+    return extracted;
   }
 
   @Override

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

@@ -25,6 +25,7 @@ import datadog.trace.core.datastreams.DataStreamsMonitoring
 import java.util.function.Function
 import java.util.function.Supplier
 
+import static datadog.context.Context.root
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_DECODED_RESOURCE_PRESERVE_SPACES
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_QUERY_STRING
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_RESOURCE
@@ -497,7 +498,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
     def decorator = newDecorator(mTracer, null)
 
     when:
-    decorator.startSpan("test", headers, null)
+    decorator.startSpan(headers, root())
 
     then:
     1 * mSpan.setMeasured(true) >> mSpan

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

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.akkahttp;
 
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.instrumentation.akkahttp.AkkaHttpServerDecorator.DECORATE;
 
 import akka.http.scaladsl.model.HttpRequest;
@@ -10,12 +11,13 @@
 
 public class DatadogWrapperHelper {
   public static ContextScope createSpan(final HttpRequest request) {
-    final Context context = DECORATE.extract(request);
-    final AgentSpan span = DECORATE.startSpan(request, context);
+    final Context parentContext = DECORATE.extract(request);
+    final Context context = DECORATE.startSpan(request, parentContext);
+    final AgentSpan span = spanFromContext(context);
     DECORATE.afterStart(span);
-    DECORATE.onRequest(span, request, request, context);
+    DECORATE.onRequest(span, request, request, parentContext);
 
-    return context.with(span).attach();
+    return context.attach();
   }
 
   public static void finishSpan(final AgentSpan span, final HttpResponse response) {

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

@@ -1,4 +1,4 @@
-package datadog.trace.instrumentation.azurefunctions;
+package datadog.trace.instrumentation.azure.functions;
 
 import com.microsoft.azure.functions.HttpRequestMessage;
 import com.microsoft.azure.functions.HttpResponseMessage;

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

@@ -1,11 +1,11 @@
-package datadog.trace.instrumentation.azurefunctions;
+package datadog.trace.instrumentation.azure.functions;
 
 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.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
-import static datadog.trace.instrumentation.azurefunctions.AzureFunctionsDecorator.DECORATE;
+import static datadog.trace.instrumentation.azure.functions.AzureFunctionsDecorator.DECORATE;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -65,15 +65,16 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class AzureFunctionsAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static ContextScope methodEnter(
-        @Advice.Argument(0) final HttpRequestMessage request,
-        @Advice.Argument(1) final ExecutionContext context) {
-      final Context extractedContext = DECORATE.extract(request);
-      final AgentSpan span = DECORATE.startSpan(request, extractedContext);
-      DECORATE.afterStart(span, context.getFunctionName());
-      DECORATE.onRequest(span, request, request, extractedContext);
+        @Advice.Argument(0) final HttpRequestMessage<?> request,
+        @Advice.Argument(1) final ExecutionContext executionContext) {
+      final Context parentContext = DECORATE.extract(request);
+      final Context context = DECORATE.startSpan(request, parentContext);
+      final AgentSpan span = spanFromContext(context);
+      DECORATE.afterStart(span, executionContext.getFunctionName());
+      DECORATE.onRequest(span, request, request, parentContext);
       HTTP_RESOURCE_DECORATOR.withRoute(
           span, request.getHttpMethod().name(), request.getUri().getPath());
-      return extractedContext.with(span).attach();
+      return context.attach();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

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

@@ -1,4 +1,4 @@
-package datadog.trace.instrumentation.azurefunctions;
+package datadog.trace.instrumentation.azure.functions;
 
 import com.microsoft.azure.functions.HttpRequestMessage;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;

dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java

@@ -13,7 +13,6 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.CorrelationIdentifier;
-import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
@@ -74,15 +73,17 @@ public static class HandleAdvice {
       }
 
       final Context parentContext = DECORATE.extract(request);
-      final AgentSpan span = DECORATE.startSpan(request, parentContext);
+      final Context context = DECORATE.startSpan(request, parentContext);
+      final AgentSpan span = spanFromContext(context);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request, request, parentContext);
 
-      scope = parentContext.with(span).attach();
+      scope = context.attach();
 
       request.setAttribute(DD_SPAN_ATTRIBUTE, span);
-      request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
-      request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
+      request.setAttribute(
+          CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
+      request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
 
       Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
       if (rba != null) {

dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java

@@ -1,5 +1,7 @@
 package datadog.trace.instrumentation.grizzlyhttp232;
 
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
+
 import datadog.appsec.api.blocking.BlockingContentType;
 import datadog.context.Context;
 import datadog.context.ContextScope;
@@ -112,13 +114,14 @@ public static NextAction onHttpCodecFilterExit(
     }
     HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader;
     HttpResponsePacket httpResponse = httpRequest.getResponse();
-    Context context = DECORATE.extract(httpRequest);
-    AgentSpan span = DECORATE.startSpan(httpRequest, context);
-    ContextScope scope = context.with(span).attach();
+    Context parentContext = DECORATE.extract(httpRequest);
+    Context context = DECORATE.startSpan(httpRequest, parentContext);
+    ContextScope scope = context.attach();
+    AgentSpan span = spanFromContext(context);
     DECORATE.afterStart(span);
     ctx.getAttributes().setAttribute(DD_SPAN_ATTRIBUTE, span);
     ctx.getAttributes().setAttribute(DD_RESPONSE_ATTRIBUTE, httpResponse);
-    DECORATE.onRequest(span, httpRequest, httpRequest, context);
+    DECORATE.onRequest(span, httpRequest, httpRequest, parentContext);
 
     Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
     if (rba != null && thiz instanceof HttpServerFilter) {

dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java

@@ -1,13 +1,13 @@
 package datadog.trace.instrumentation.jetty11;
 
 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.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
 import static datadog.trace.instrumentation.jetty11.JettyDecorator.DECORATE;
 
 import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.api.CorrelationIdentifier;
-import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 import org.eclipse.jetty.server.HttpChannel;
@@ -26,16 +26,17 @@ public static ContextScope onEnter(
         return activateSpan((AgentSpan) existingSpan);
       }
 
-      final Context context = DECORATE.extract(req);
-      span = DECORATE.startSpan(req, context);
-      final ContextScope scope = context.with(span).attach();
+      final Context parentContext = DECORATE.extract(req);
+      final Context context = DECORATE.startSpan(req, parentContext);
+      final ContextScope scope = context.attach();
+      span = spanFromContext(context);
       span.setMeasured(true);
       DECORATE.afterStart(span);
-      DECORATE.onRequest(span, req, req, context);
+      DECORATE.onRequest(span, req, req, parentContext);
 
       req.setAttribute(DD_SPAN_ATTRIBUTE, span);
-      req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
-      req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
+      req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
+      req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
       return scope;
     }
 

dd-java-agent/instrumentation/jetty-11/src/test/groovy/HttpChannelAppSecTest.groovy

@@ -73,7 +73,7 @@ class HttpChannelAppSecTest extends AgentTestRunner {
 
     @Override
     void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
-      if ('datadog/trace/bootstrap/instrumentation/api/AgentSpan' == owner && 'getRequestBlockingAction' == name) {
+      if ('datadog/trace/instrumentation/jetty/JettyBlockingHelper' == owner && 'hasRequestBlockingAction' == name) {
         blockApplied = true
       }
       super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)

dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java

@@ -1,11 +1,12 @@
 package datadog.trace.instrumentation.jetty12;
 
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
+import static datadog.trace.instrumentation.jetty12.JettyDecorator.DECORATE;
 
 import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.api.CorrelationIdentifier;
-import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 import org.eclipse.jetty.server.Request;
@@ -27,16 +28,17 @@ public static void onExit(
         }
       }
 
-      final Context context = JettyDecorator.DECORATE.extract(req);
-      final AgentSpan span = JettyDecorator.DECORATE.startSpan(req, context);
-      try (final ContextScope scope = context.with(span).attach()) {
+      final Context parentContext = DECORATE.extract(req);
+      final Context context = DECORATE.startSpan(req, parentContext);
+      try (final ContextScope ignored = context.attach()) {
+        final AgentSpan span = spanFromContext(context);
         span.setMeasured(true);
-        JettyDecorator.DECORATE.afterStart(span);
-        JettyDecorator.DECORATE.onRequest(span, req, req, context);
+        DECORATE.afterStart(span);
+        DECORATE.onRequest(span, req, req, parentContext);
 
         req.setAttribute(DD_SPAN_ATTRIBUTE, span);
-        req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
-        req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
+        req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
+        req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
         ret = JettyRunnableWrapper.wrapIfNeeded(ret);
       }
     }
@@ -53,8 +55,8 @@ public static void stopSpan(@Advice.This final HttpChannelState channel) {
       Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE);
       if (spanObj instanceof AgentSpan) {
         final AgentSpan span = (AgentSpan) spanObj;
-        JettyDecorator.DECORATE.onResponse(span, channel);
-        JettyDecorator.DECORATE.beforeFinish(span);
+        DECORATE.onResponse(span, channel);
+        DECORATE.beforeFinish(span);
         span.finish();
       }
     }