Context API beforeFinish Migration

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

@@ -4,6 +4,7 @@
 import static datadog.trace.api.cache.RadixTreeCache.UNSET_PORT;
 import static datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver.hostName;
 
+import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
@@ -76,15 +77,19 @@ public AgentSpan afterStart(final AgentSpan span) {
     return span;
   }
 
-  public AgentScope beforeFinish(final AgentScope scope) {
-    beforeFinish(scope.span());
+  public ContextScope beforeFinish(final ContextScope scope) {
+    beforeFinish(scope.context());
     return scope;
   }
 
   public AgentSpan beforeFinish(final AgentSpan span) {
     return span;
   }
 
+  public Context beforeFinish(final Context context) {
+    return context;
+  }
+
   public AgentScope onError(final AgentScope scope, final Throwable throwable) {
     onError(scope.span(), throwable);
     return scope;

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

@@ -55,7 +55,7 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
   private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class);
   private static final int UNSET_PORT = 0;
 
-  public static final String DD_SPAN_ATTRIBUTE = "datadog.span";
+  public static final String DD_CONTEXT_ATTRIBUTE = "datadog.context";
   public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch";
   public static final String DD_RUM_INJECTED = "datadog.rum.injected";
   public static final String DD_FIN_DISP_LIST_SPAN_ATTRIBUTE =
@@ -147,17 +147,18 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) {
         instrumentationNames != null && instrumentationNames.length > 0
             ? instrumentationNames[0]
             : DEFAULT_INSTRUMENTATION_NAME;
-    AgentSpanContext.Extracted extracted = callIGCallbackStart(getExtractedSpanContext(context));
+    AgentSpanContext extracted = getExtractedSpanContext(context);
+    // Call IG callbacks
+    extracted = callIGCallbackStart(extracted);
     AgentSpan span =
         tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
+    // Apply RequestBlockingAction if any
     Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
     if (flow.getAction() instanceof RequestBlockingAction) {
       span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
     }
-    AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
-    if (null != carrier && null != getter) {
-      tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
-    }
+    // DSM Checkpoint
+    tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
     return context.with(span);
   }
 
@@ -381,8 +382,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
     return span;
   }
 
-  private AgentSpanContext.Extracted callIGCallbackStart(
-      @Nullable final AgentSpanContext.Extracted extracted) {
+  private AgentSpanContext callIGCallbackStart(@Nullable final AgentSpanContext extracted) {
     AgentTracer.TracerAPI tracer = tracer();
     Supplier<Flow<Object>> startedCbAppSec =
         tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());
@@ -517,9 +517,11 @@ private Flow<Void> callIGCallbackURI(
   }
 
   @Override
-  public AgentSpan beforeFinish(AgentSpan span) {
+  public Context beforeFinish(Context context) {
+    AgentSpan span = AgentSpan.fromContext(context);
     onRequestEndForInstrumentationGateway(span);
-    return super.beforeFinish(span);
+
+    return super.beforeFinish(context);
   }
 
   private void onRequestEndForInstrumentationGateway(@Nonnull final AgentSpan span) {

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

@@ -38,15 +38,15 @@ public Future<HttpResponse> apply(final HttpRequest request) {
       request.discardEntityBytes(materializer);
       HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(rba, request);
       span.getRequestContext().getTraceSegment().effectivelyBlocked();
-      DatadogWrapperHelper.finishSpan(span, response);
+      DatadogWrapperHelper.finishSpan(scope.context(), response);
       return FastFuture$.MODULE$.<HttpResponse>successful().apply(response);
     }
 
     try {
       futureResponse = userHandler.apply(request);
     } catch (final Throwable t) {
       scope.close();
-      DatadogWrapperHelper.finishSpan(span, t);
+      DatadogWrapperHelper.finishSpan(scope.context(), t);
       throw t;
     }
 
@@ -67,14 +67,14 @@ public HttpResponse apply(HttpResponse response) {
                       response = newResponse;
                     }
 
-                    DatadogWrapperHelper.finishSpan(span, response);
+                    DatadogWrapperHelper.finishSpan(scope.context(), response);
                     return response;
                   }
                 },
                 new AbstractFunction1<Throwable, Throwable>() {
                   @Override
                   public Throwable apply(final Throwable t) {
-                    DatadogWrapperHelper.finishSpan(span, t);
+                    DatadogWrapperHelper.finishSpan(scope.context(), t);
                     return t;
                   }
                 },

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

@@ -78,7 +78,7 @@ public void onPush() throws Exception {
                     skipNextPull[0] = true;
                     requestContext.getTraceSegment().effectivelyBlocked();
                     emit(responseOutlet, response);
-                    DatadogWrapperHelper.finishSpan(span, response);
+                    DatadogWrapperHelper.finishSpan(scope.context(), response);
                     pull(requestInlet);
                     scope.close();
                     return;
@@ -142,7 +142,7 @@ public void onPush() throws Exception {
                     response.discardEntityBytes(materializer());
                     response = newResponse;
                   }
-                  DatadogWrapperHelper.finishSpan(span, response);
+                  DatadogWrapperHelper.finishSpan(scope.context(), response);
                   // Check if the active span matches the scope from when the request came in,
                   // and close it. If it's not, then it will be cleaned up actor message
                   // processing instrumentation that drives this state machine
@@ -172,7 +172,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception {
                 if (scope != null) {
                   // Mark the span as failed
                   AgentSpan span = fromContext(scope.context());
-                  DatadogWrapperHelper.finishSpan(span, ex);
+                  DatadogWrapperHelper.finishSpan(scope.context(), ex);
                 }
                 // We will not receive any more responses from the user code, so clean up any
                 // remaining spans

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

@@ -20,17 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) {
     return context.attach();
   }
 
-  public static void finishSpan(final AgentSpan span, final HttpResponse response) {
+  public static void finishSpan(final Context context, final HttpResponse response) {
+    final AgentSpan span = fromContext(context);
     DECORATE.onResponse(span, response);
-    DECORATE.beforeFinish(span);
+    DECORATE.beforeFinish(context);
 
     span.finish();
   }
 
-  public static void finishSpan(final AgentSpan span, final Throwable t) {
+  public static void finishSpan(final Context context, final Throwable t) {
+    final AgentSpan span = fromContext(context);
     DECORATE.onError(span, t);
     span.setHttpStatusCode(500);
-    DECORATE.beforeFinish(span);
+    DECORATE.beforeFinish(context);
 
     span.finish();
   }

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

@@ -45,7 +45,7 @@ public static void onExit(
       if (throwable != null) {
         DECORATE.onError(span, throwable);
       }
-      DECORATE.beforeFinish(span);
+      DECORATE.beforeFinish(scope.context());
     } finally {
       scope.close();
       span.finish();

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

@@ -38,7 +38,7 @@ public static void onExit(
       if (throwable != null) {
         DECORATE.onError(span, throwable);
       }
-      DECORATE.beforeFinish(span);
+      DECORATE.beforeFinish(scope.context());
     } finally {
       scope.close();
       span.finish();

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

@@ -85,7 +85,7 @@ public static void methodExit(
       final AgentSpan span = fromContext(scope.context());
       DECORATE.onError(span, throwable);
       DECORATE.onResponse(span, response);
-      DECORATE.beforeFinish(span);
+      DECORATE.beforeFinish(scope.context());
       scope.close();
       span.finish();
     }

dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.cxf;
 
 import com.google.auto.service.AutoService;
+import datadog.context.Context;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
@@ -9,6 +10,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -59,11 +61,16 @@ public static AgentScope beforeInvoke(@Advice.Argument(0) final Exchange exchang
       if (exchange == null || exchange.getInMessage() == null || AgentTracer.activeSpan() != null) {
         return null;
       }
-      final Object span =
+      final Object contextObj =
           ServletHelper.getServletRequestAttribute(
-              exchange.getInMessage().get("HTTP.REQUEST"), HttpServerDecorator.DD_SPAN_ATTRIBUTE);
-      if (span instanceof AgentSpan) {
-        return AgentTracer.activateSpan((AgentSpan) span);
+              exchange.getInMessage().get("HTTP.REQUEST"),
+              HttpServerDecorator.DD_CONTEXT_ATTRIBUTE);
+      if (contextObj instanceof Context) {
+        Context context = (Context) contextObj;
+        AgentSpan span = Java8BytecodeBridge.spanFromContext(context);
+        if (span != null) {
+          return AgentTracer.activateSpan(span);
+        }
       }
       return null;
     }

dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java

@@ -92,7 +92,7 @@ public static void setupCallback(
       final AgentSpan span = scope.span();
       if (throwable != null) {
         DECORATE.onError(span, throwable);
-        DECORATE.beforeFinish(span);
+        DECORATE.beforeFinish(scope.context());
         span.finish();
         scope.close();
         return;

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

@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.grizzly;
 
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE;
 
 import datadog.appsec.api.blocking.BlockingContentType;
+import datadog.context.Context;
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.bootstrap.blocking.BlockingActionHelper;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -35,14 +37,14 @@ public class GrizzlyBlockingHelper {
   private GrizzlyBlockingHelper() {}
 
   public static boolean block(
-      Request request, Response response, Flow.Action.RequestBlockingAction rba, AgentSpan span) {
+      Request request, Response response, Flow.Action.RequestBlockingAction rba, Context context) {
     return block(
         request,
         response,
         rba.getStatusCode(),
         rba.getBlockingContentType(),
         rba.getExtraHeaders(),
-        span);
+        context);
   }
 
   public static boolean block(
@@ -51,11 +53,12 @@ public static boolean block(
       int statusCode,
       BlockingContentType bct,
       Map<String, String> extraHeaders,
-      AgentSpan span) {
+      Context context) {
     if (GET_OUTPUT_STREAM == null) {
       return false;
     }
 
+    AgentSpan span = spanFromContext(context);
     try {
       OutputStream os = (OutputStream) GET_OUTPUT_STREAM.invoke(response);
       response.setStatus(BlockingActionHelper.getHttpCode(statusCode));
@@ -76,13 +79,17 @@ public static boolean block(
       os.close();
       response.finish();
 
-      span.getRequestContext().getTraceSegment().effectivelyBlocked();
+      if (span != null) {
+        span.getRequestContext().getTraceSegment().effectivelyBlocked();
+      }
       SpanClosingListener.LISTENER.onAfterService(request);
     } catch (Throwable e) {
       log.info("Error committing blocking response", e);
-      DECORATE.onError(span, e);
-      DECORATE.beforeFinish(span);
-      span.finish();
+      if (span != null) {
+        DECORATE.onError(span, e);
+        DECORATE.beforeFinish(context);
+        span.finish();
+      }
     }
 
     return true;

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

@@ -2,7 +2,7 @@
 
 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.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
 import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -68,7 +68,7 @@ public static class HandleAdvice {
         @Advice.Local("contextScope") ContextScope scope,
         @Advice.Argument(0) final Request request,
         @Advice.Argument(1) final Response response) {
-      if (request.getAttribute(DD_SPAN_ATTRIBUTE) != null) {
+      if (request.getAttribute(DD_CONTEXT_ATTRIBUTE) != null) {
         return false;
       }
 
@@ -80,14 +80,14 @@ public static class HandleAdvice {
 
       scope = context.attach();
 
-      request.setAttribute(DD_SPAN_ATTRIBUTE, span);
+      request.setAttribute(DD_CONTEXT_ATTRIBUTE, context);
       request.setAttribute(
           CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
       request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
 
       Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
       if (rba != null) {
-        boolean success = GrizzlyBlockingHelper.block(request, response, rba, span);
+        boolean success = GrizzlyBlockingHelper.block(request, response, rba, context);
         if (success) {
           return true; /* skip body */
         }

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

@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.grizzly;
 
-import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
 import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import org.glassfish.grizzly.http.server.AfterServiceListener;
 import org.glassfish.grizzly.http.server.Request;
@@ -12,13 +14,17 @@ public class SpanClosingListener implements AfterServiceListener {
 
   @Override
   public void onAfterService(final Request request) {
-    final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE);
-    if (spanAttr instanceof AgentSpan) {
-      request.removeAttribute(DD_SPAN_ATTRIBUTE);
-      final AgentSpan span = (AgentSpan) spanAttr;
-      DECORATE.onResponse(span, request.getResponse());
-      DECORATE.beforeFinish(span);
-      span.finish();
+    final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE);
+    if (contextAttr instanceof Context) {
+      request.removeAttribute(DD_CONTEXT_ATTRIBUTE);
+
+      final Context context = (Context) contextAttr;
+      final AgentSpan span = spanFromContext(context);
+      if (span != null) {
+        DECORATE.onResponse(span, request.getResponse());
+        DECORATE.beforeFinish(context);
+        span.finish();
+      }
     }
   }
 }