play-ws-2.x

Author: SylvainJuge

instrumentation/play/play-ws/play-ws-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/playws/v2_0/PlayWsInstrumentationModule.java

@@ -13,18 +13,22 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.javaagent.instrumentation.playws.AbstractBootstrapInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.AsyncHttpClientInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.HandlerPublisherInstrumentation;
 import java.util.List;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
 import play.shaded.ahc.org.asynchttpclient.AsyncHandler;
 import play.shaded.ahc.org.asynchttpclient.Request;
 import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler;
 import play.shaded.ahc.org.asynchttpclient.ws.WebSocketUpgradeHandler;
 
 @AutoService(InstrumentationModule.class)
-public class PlayWsInstrumentationModule extends InstrumentationModule {
+public class PlayWsInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public PlayWsInstrumentationModule() {
     super("play-ws", "play-ws-2.0");
   }
@@ -37,20 +41,26 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new AbstractBootstrapInstrumentation());
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @SuppressWarnings("unused")
   public static class ClientAdvice {
 
+    @AssignReturned.ToArguments(@ToArgument(value = 1, index = 1))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
+    public static Object[] methodEnter(
         @Advice.Argument(0) Request request,
-        @Advice.Argument(value = 1, readOnly = false) AsyncHandler<?> asyncHandler,
-        @Advice.Local("otelContext") Context context) {
+        @Advice.Argument(1) AsyncHandler<?> originalAsyncHandler) {
+      AsyncHandler<?> asyncHandler = originalAsyncHandler;
       Context parentContext = currentContext();
       if (!instrumenter().shouldStart(parentContext, request)) {
-        return;
+        return new Object[] {null, asyncHandler};
       }
 
-      context = instrumenter().start(parentContext, request);
+      Context context = instrumenter().start(parentContext, request);
 
       if (asyncHandler instanceof StreamedAsyncHandler) {
         asyncHandler =
@@ -60,13 +70,15 @@ public static void methodEnter(
         // websocket upgrade handlers aren't supported
         asyncHandler = new AsyncHandlerWrapper<>(asyncHandler, request, context, parentContext);
       }
+      return new Object[] {context, asyncHandler};
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Argument(0) Request request,
         @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelContext") Context context) {
+        @Advice.Enter Object[] enterResult) {
+      Context context = (Context) enterResult[0];
       if (context != null && throwable != null) {
         instrumenter().end(context, request, null, throwable);
       }

instrumentation/play/play-ws/play-ws-2.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/playws/v2_1/PlayWsInstrumentationModule.java

@@ -13,18 +13,23 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.javaagent.instrumentation.playws.AbstractBootstrapInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.AsyncHttpClientInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.playws.HandlerPublisherInstrumentation;
 import java.util.List;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
 import play.shaded.ahc.org.asynchttpclient.AsyncHandler;
 import play.shaded.ahc.org.asynchttpclient.Request;
 import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler;
 import play.shaded.ahc.org.asynchttpclient.ws.WebSocketUpgradeHandler;
 
 @AutoService(InstrumentationModule.class)
-public class PlayWsInstrumentationModule extends InstrumentationModule {
+public class PlayWsInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public PlayWsInstrumentationModule() {
     super("play-ws", "play-ws-2.1");
   }
@@ -37,20 +42,25 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new AbstractBootstrapInstrumentation());
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @SuppressWarnings("unused")
   public static class ClientAdvice {
 
+    @AssignReturned.ToArguments(@ToArgument(value = 1, index = 1))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
+    public static Object[] methodEnter(
         @Advice.Argument(0) Request request,
-        @Advice.Argument(value = 1, readOnly = false) AsyncHandler<?> asyncHandler,
-        @Advice.Local("otelContext") Context context) {
+        @Advice.Argument(1) AsyncHandler<?> originalAsyncHandler) {
+      AsyncHandler<?> asyncHandler = originalAsyncHandler;
       Context parentContext = currentContext();
       if (!instrumenter().shouldStart(parentContext, request)) {
-        return;
+        return new Object[] {null, asyncHandler};
       }
-
-      context = instrumenter().start(parentContext, request);
+      Context context = instrumenter().start(parentContext, request);
 
       if (asyncHandler instanceof StreamedAsyncHandler) {
         asyncHandler =
@@ -60,13 +70,15 @@ public static void methodEnter(
         // websocket upgrade handlers aren't supported
         asyncHandler = new AsyncHandlerWrapper<>(asyncHandler, request, context, parentContext);
       }
+      return new Object[] {context, asyncHandler};
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Argument(0) Request request,
-        @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelContext") Context context) {
+        @Advice.Thrown @Nullable Throwable throwable,
+        @Advice.Enter Object[] enterResult) {
+      Context context = (Context) enterResult[0];
       if (context != null && throwable != null) {
         instrumenter().end(context, request, null, throwable);
       }