feat(play): Improve context usage in advice

Author: PerfectSlayer

dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java

@@ -9,6 +9,7 @@
 import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
 import net.bytebuddy.asm.Advice;
 import play.api.mvc.Action;
@@ -19,9 +20,8 @@
 
 public class PlayAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static ContextScope onEnter(
-      @Advice.Argument(value = 0, readOnly = false) Request<?> req,
-      @Advice.Local("extractedContext") Context extractedContext) {
+  public static ContextScope onEnter(@Advice.Argument(value = 0, readOnly = false) Request<?> req) {
+    final Context parentContext;
     final AgentSpan span;
     final ContextScope scope;
 
@@ -32,13 +32,14 @@ public static ContextScope onEnter(
 
     if (activeSpan() == null) {
       final Headers headers = req.headers();
-      extractedContext = DECORATE.extract(headers);
-      final Context context = DECORATE.startSpan(headers, extractedContext);
+      parentContext = DECORATE.extract(headers);
+      final Context context = DECORATE.startSpan(headers, parentContext);
       span = spanFromContext(context);
       scope = context.attach();
     } else {
       // An upstream framework (e.g. akka-http, netty) has already started the span.
       // Do not extract the context.
+      parentContext = Java8BytecodeBridge.getRootContext();
       span = startSpan("play", PLAY_REQUEST);
       scope = span.attach();
     }
@@ -50,15 +51,14 @@ public static ContextScope onEnter(
 
     // Moved from OnMethodExit
     // Call onRequest on return after tags are populated.
-    DECORATE.onRequest(span, req, req, extractedContext);
+    DECORATE.onRequest(span, req, req, parentContext);
 
     return scope;
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void stopTraceOnResponse(
       @Advice.Enter final ContextScope playControllerScope,
-      @Advice.Local("extractedContext") Context extractedContext,
       @Advice.This final Object thisAction,
       @Advice.Thrown final Throwable throwable,
       @Advice.Argument(0) final Request<?> req,

dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayHttpServerDecorator.java

@@ -3,6 +3,7 @@
 import static datadog.trace.api.gateway.Events.EVENTS;
 import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
 
+import datadog.context.Context;
 import datadog.trace.api.Config;
 import datadog.trace.api.cache.DDCache;
 import datadog.trace.api.cache.DDCaches;
@@ -11,7 +12,6 @@
 import datadog.trace.api.gateway.RequestContextSlot;
 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.ResourceNamePriorities;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
 import datadog.trace.bootstrap.instrumentation.api.URIUtils;
@@ -36,7 +36,7 @@
 import scala.Option;
 
 public class PlayHttpServerDecorator
-    extends HttpServerDecorator<Request, Request, Result, Headers> {
+    extends HttpServerDecorator<Request<?>, Request<?>, Result, Headers> {
   private static final Logger LOG = LoggerFactory.getLogger(PlayHttpServerDecorator.class);
   public static final boolean REPORT_HTTP_STATUS = Config.get().getPlayReportHttpStatus();
   public static final CharSequence PLAY_REQUEST = UTF8BytesString.create("play.request");
@@ -96,17 +96,17 @@ public CharSequence spanName() {
   }
 
   @Override
-  protected String method(final Request httpRequest) {
+  protected String method(final Request<?> httpRequest) {
     return httpRequest.method();
   }
 
   @Override
-  protected URIDataAdapter url(final Request request) {
+  protected URIDataAdapter url(final Request<?> request) {
     return new RequestURIDataAdapter(request);
   }
 
   @Override
-  protected String peerHostIP(final Request request) {
+  protected String peerHostIP(final Request<?> request) {
     RemoteConnection connection = request.connection();
     if (connection instanceof RemoteConnectionWithRawAddress) {
       return ((RemoteConnectionWithRawAddress) connection).rawRemoteAddressString();
@@ -116,7 +116,7 @@ protected String peerHostIP(final Request request) {
   }
 
   @Override
-  protected int peerPort(final Request request) {
+  protected int peerPort(final Request<?> request) {
     return 0;
   }
 
@@ -128,9 +128,9 @@ protected int status(final Result httpResponse) {
   @Override
   public AgentSpan onRequest(
       final AgentSpan span,
-      final Request connection,
-      final Request request,
-      AgentSpanContext.Extracted context) {
+      final Request<?> connection,
+      final Request<?> request,
+      final Context context) {
     super.onRequest(span, connection, request, context);
     if (request != null) {
       // more about routes here: