Update play mvc to use AdviceScope (#13140)

Co-authored-by: Lauri Tulmin <[email protected]>

Author: jaydeluca

instrumentation/play/play-mvc/play-mvc-2.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/play/v2_4/ActionInstrumentation.java

@@ -18,6 +18,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -49,41 +50,60 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class ApplyAdvice {
 
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(0) Request<?> req,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      Context parentContext = currentContext();
-      if (!instrumenter().shouldStart(parentContext, null)) {
-        return;
+    public static class AdviceScope {
+      private final Context context;
+      private final Scope scope;
+
+      public AdviceScope(Context context, Scope scope) {
+        this.context = context;
+        this.scope = scope;
+      }
+
+      @Nullable
+      public static AdviceScope start(Context parentContext) {
+        if (!instrumenter().shouldStart(parentContext, null)) {
+          return null;
+        }
+
+        Context context = instrumenter().start(parentContext, null);
+        return new AdviceScope(context, context.makeCurrent());
       }
 
-      context = instrumenter().start(parentContext, null);
-      scope = context.makeCurrent();
+      public void end(
+          @Nullable Throwable throwable,
+          Future<Result> responseFuture,
+          Action<?> thisAction,
+          Request<?> req) {
+        scope.close();
+        updateSpan(context, req);
+
+        if (throwable == null) {
+          // span finished in RequestCompleteCallback
+          responseFuture.onComplete(
+              new RequestCompleteCallback(context), thisAction.executionContext());
+        } else {
+          instrumenter().end(context, null, null, throwable);
+        }
+      }
+    }
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AdviceScope onEnter(@Advice.Argument(0) Request<?> req) {
+      return AdviceScope.start(currentContext());
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopTraceOnResponse(
-        @Advice.This Object thisAction,
+        @Advice.This Action<?> thisAction,
         @Advice.Thrown Throwable throwable,
         @Advice.Argument(0) Request<?> req,
-        @Advice.Return(readOnly = false) Future<Result> responseFuture,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (scope == null) {
+        @Advice.Return Future<Result> responseFuture,
+        @Advice.Enter @Nullable AdviceScope actionScope) {
+      if (actionScope == null) {
         return;
       }
-      scope.close();
-
-      updateSpan(context, req);
-      // span finished in RequestCompleteCallback
-      if (throwable == null) {
-        responseFuture.onComplete(
-            new RequestCompleteCallback(context), ((Action<?>) thisAction).executionContext());
-      } else {
-        instrumenter().end(context, null, null, throwable);
-      }
+
+      actionScope.end(throwable, responseFuture, thisAction, req);
     }
   }
 }

instrumentation/play/play-mvc/play-mvc-2.6/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/play/v2_6/ActionInstrumentation.java

@@ -18,6 +18,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -49,44 +50,65 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class ApplyAdvice {
 
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.Argument(0) Request<?> req,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      Context parentContext = currentContext();
-      if (!instrumenter().shouldStart(parentContext, null)) {
-        return;
+    public static class AdviceScope {
+      private final Context context;
+      private final Scope scope;
+
+      public AdviceScope(Context context, Scope scope) {
+        this.context = context;
+        this.scope = scope;
+      }
+
+      @Nullable
+      public static AdviceScope start(Context parentContext) {
+        if (!instrumenter().shouldStart(parentContext, null)) {
+          return null;
+        }
+
+        Context context = instrumenter().start(parentContext, null);
+        return new AdviceScope(context, context.makeCurrent());
+      }
+
+      public Future<Result> end(
+          @Nullable Throwable throwable,
+          Future<Result> responseFuture,
+          Action<?> thisAction,
+          Request<?> req) {
+        scope.close();
+        updateSpan(context, req);
+
+        if (throwable == null) {
+          // span is finished when future completes
+          // not using responseFuture.onComplete() because that doesn't guarantee this handler span
+          // will be completed before the server span completes
+          responseFuture =
+              ResponseFutureWrapper.wrap(responseFuture, context, thisAction.executionContext());
+        } else {
+          instrumenter().end(context, null, null, throwable);
+        }
+
+        return responseFuture;
       }
+    }
 
-      context = instrumenter().start(parentContext, null);
-      scope = context.makeCurrent();
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AdviceScope onEnter(@Advice.Argument(0) Request<?> req) {
+      return AdviceScope.start(currentContext());
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void stopTraceOnResponse(
-        @Advice.This Object thisAction,
+    @Advice.AssignReturned.ToReturned
+    public static Future<Result> stopTraceOnResponse(
+        @Advice.This Action<?> thisAction,
         @Advice.Thrown Throwable throwable,
         @Advice.Argument(0) Request<?> req,
-        @Advice.Return(readOnly = false) Future<Result> responseFuture,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (scope == null) {
-        return;
-      }
-      scope.close();
-
-      updateSpan(context, req);
-      if (throwable == null) {
-        // span is finished when future completes
-        // not using responseFuture.onComplete() because that doesn't guarantee this handler span
-        // will be completed before the server span completes
-        responseFuture =
-            ResponseFutureWrapper.wrap(
-                responseFuture, context, ((Action<?>) thisAction).executionContext());
-      } else {
-        instrumenter().end(context, null, null, throwable);
+        @Advice.Return Future<Result> responseFuture,
+        @Advice.Enter @Nullable AdviceScope actionScope) {
+      if (actionScope == null) {
+        return responseFuture;
       }
+
+      return actionScope.end(throwable, responseFuture, thisAction, req);
     }
   }
 }