indy ready migration - part1 (#13759)

Co-authored-by: Jonas Kunz <[email protected]>
Co-authored-by: Jonas Kunz <[email protected]>

Author: SylvainJuge

instrumentation/activej-http-6.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/activejhttp/ActivejAsyncServletInstrumentation.java

@@ -16,15 +16,16 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
-import io.activej.http.AsyncServlet;
 import io.activej.http.HttpRequest;
 import io.activej.http.HttpResponse;
 import io.activej.promise.Promise;
 import io.opentelemetry.context.Context;
 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.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -52,39 +53,57 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class ServeAdvice {
 
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
-        @Advice.This AsyncServlet asyncServlet,
-        @Advice.Argument(0) HttpRequest request,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope,
-        @Advice.Local("httpRequest") HttpRequest httpRequest) {
-      Context parentContext = currentContext();
-      httpRequest = request;
-      if (!instrumenter().shouldStart(parentContext, request)) {
-        return;
+    public static class AdviceScope {
+      private final HttpRequest httpRequest;
+      private final Context context;
+      private final Scope scope;
+
+      private AdviceScope(Context context, Scope scope, HttpRequest httpRequest) {
+        this.context = context;
+        this.scope = scope;
+        this.httpRequest = httpRequest;
+      }
+
+      @Nullable
+      public static AdviceScope start(HttpRequest request) {
+        Context parentContext = currentContext();
+        if (!instrumenter().shouldStart(parentContext, request)) {
+          return null;
+        }
+        Context context = instrumenter().start(parentContext, request);
+        return new AdviceScope(context, context.makeCurrent(), request);
+      }
+
+      public Promise<HttpResponse> end(Promise<HttpResponse> responsePromise, Throwable throwable) {
+        scope.close();
+        Promise<HttpResponse> returnValue = responsePromise;
+        if (throwable != null) {
+          instrumenter().end(context, httpRequest, null, throwable);
+          return responsePromise;
+        } else {
+          return PromiseWrapper.wrap(responsePromise, httpRequest, context);
+        }
       }
-      context = instrumenter().start(parentContext, request);
-      scope = context.makeCurrent();
     }
 
+    @Nullable
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AdviceScope methodEnter(@Advice.Argument(0) HttpRequest request) {
+      return AdviceScope.start(request);
+    }
+
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(
-        @Advice.This AsyncServlet asyncServlet,
-        @Advice.Return(readOnly = false) Promise<HttpResponse> responsePromise,
+    public static Promise<HttpResponse> methodExit(
+        @Advice.Return Promise<HttpResponse> responsePromise,
         @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope,
-        @Advice.Local("httpRequest") HttpRequest httpRequest) {
-      if (scope == null) {
-        return;
-      }
-      scope.close();
-      if (throwable != null) {
-        instrumenter().end(context, httpRequest, null, throwable);
-      } else {
-        responsePromise = PromiseWrapper.wrap(responsePromise, httpRequest, context);
+        @Advice.Enter @Nullable AdviceScope adviceScope) {
+
+      if (adviceScope == null) {
+        return responsePromise;
       }
+
+      return adviceScope.end(responsePromise, throwable);
     }
   }
 }

instrumentation/activej-http-6.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/activejhttp/ActivejHttpServerInstrumentationModule.java

@@ -11,16 +11,23 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
-public class ActivejHttpServerInstrumentationModule extends InstrumentationModule {
+public class ActivejHttpServerInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
 
   public ActivejHttpServerInstrumentationModule() {
     super("activej-http", "activej-http-6.0");
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return singletonList(new ActivejAsyncServletInstrumentation());

instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaActorInstrumentationModule.java

@@ -10,14 +10,21 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class AkkaActorInstrumentationModule extends InstrumentationModule {
+public class AkkaActorInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public AkkaActorInstrumentationModule() {
     super("akka-actor", "akka-actor-2.3");
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(

instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaScheduleInstrumentation.java

@@ -11,6 +11,8 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -41,20 +43,20 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class ScheduleAdvice {
 
+    @AssignReturned.ToArguments(@ToArgument(2))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void enterSchedule(
-        @Advice.Argument(value = 2, readOnly = false) Runnable runnable) {
-      runnable = AkkaSchedulerTaskWrapper.wrap(runnable);
+    public static Runnable enterSchedule(@Advice.Argument(2) Runnable runnable) {
+      return AkkaSchedulerTaskWrapper.wrap(runnable);
     }
   }
 
   @SuppressWarnings("unused")
   public static class ScheduleOnceAdvice {
 
+    @AssignReturned.ToArguments(@ToArgument(1))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void enterScheduleOnce(
-        @Advice.Argument(value = 1, readOnly = false) Runnable runnable) {
-      runnable = AkkaSchedulerTaskWrapper.wrap(runnable);
+    public static Runnable enterScheduleOnce(@Advice.Argument(1) Runnable runnable) {
+      return AkkaSchedulerTaskWrapper.wrap(runnable);
     }
   }
 }

instrumentation/akka/akka-actor-fork-join-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaActorForkJoinInstrumentationModule.java

@@ -10,14 +10,21 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class AkkaActorForkJoinInstrumentationModule extends InstrumentationModule {
+public class AkkaActorForkJoinInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public AkkaActorForkJoinInstrumentationModule() {
     super("akka-actor-fork-join", "akka-actor-fork-join-2.5", "akka-actor");
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new AkkaForkJoinPoolInstrumentation(), new AkkaForkJoinTaskInstrumentation());

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/client/AkkaHttpClientInstrumentationModule.java

@@ -10,14 +10,21 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class AkkaHttpClientInstrumentationModule extends InstrumentationModule {
+public class AkkaHttpClientInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public AkkaHttpClientInstrumentationModule() {
     super("akka-http", "akka-http-10.0", "akka-http-client");
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new HttpExtClientInstrumentation(), new PoolMasterActorInstrumentation());

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/client/HttpExtClientInstrumentation.java

@@ -20,7 +20,10 @@
 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.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import scala.concurrent.Future;
@@ -43,55 +46,82 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class SingleRequestAdvice {
 
+    public static class AdviceScope {
+      private final Context context;
+      private final Scope scope;
+
+      private AdviceScope(Context context, Scope scope) {
+        this.context = context;
+        this.scope = scope;
+      }
+
+      public static AdviceScope start(HttpRequest request) {
+        Context parentContext = currentContext();
+        if (!instrumenter().shouldStart(parentContext, request)) {
+          return null;
+        }
+        Context context = instrumenter().start(parentContext, request);
+        // Making context current is required for header context propagation to work as expected
+        // because it implicitly relies on the current context when injecting headers.
+        Scope scope = context.makeCurrent();
+        return new AdviceScope(context, scope);
+      }
+
+      public HttpRequest injectHeaders(HttpRequest request) {
+        // Request is immutable, so we have to assign a new value once we update headers
+        return setter().inject(request);
+      }
+
+      public Future<HttpResponse> end(
+          @Nullable ActorSystem actorSystem,
+          HttpRequest request,
+          @Nullable Future<HttpResponse> responseFuture,
+          @Nullable Throwable throwable) {
+
+        scope.close();
+        if (actorSystem != null) {
+          if (throwable == null) {
+            responseFuture.onComplete(
+                new OnCompleteHandler(context, request), actorSystem.dispatcher());
+            return FutureWrapper.wrap(responseFuture, actorSystem.dispatcher(), currentContext());
+          } else {
+            instrumenter().end(context, request, null, throwable);
+          }
+        }
+        return responseFuture;
+      }
+    }
+
+    @AssignReturned.ToArguments(@ToArgument(value = 0, index = 1))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
-        @Advice.Argument(value = 0, readOnly = false) HttpRequest request,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
+    public static Object[] methodEnter(@Advice.Argument(0) HttpRequest request) {
+
       /*
       Versions 10.0 and 10.1 have slightly different structure that is hard to distinguish so here
       we cast 'wider net' and avoid instrumenting twice.
       In the future we may want to separate these, but since lots of code is reused we would need to come up
       with way of continuing to reusing it.
        */
-      Context parentContext = currentContext();
-      if (!instrumenter().shouldStart(parentContext, request)) {
-        return;
-      }
-
-      context = instrumenter().start(parentContext, request);
-      scope = context.makeCurrent();
-      // Request is immutable, so we have to assign new value once we update headers
-      request = setter().inject(request);
+      AdviceScope adviceScope = AdviceScope.start(request);
+      return new Object[] {
+        adviceScope, adviceScope == null ? request : adviceScope.injectHeaders(request)
+      };
     }
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(
-        @Advice.Argument(0) HttpRequest request,
+    public static Future<HttpResponse> methodExit(
         @Advice.This HttpExt thiz,
-        @Advice.Return(readOnly = false) Future<HttpResponse> responseFuture,
-        @Advice.Thrown Throwable throwable,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (scope == null) {
-        return;
-      }
+        @Advice.Return @Nullable Future<HttpResponse> responseFuture,
+        @Advice.Thrown @Nullable Throwable throwable,
+        @Advice.Enter Object[] enterResult) {
 
-      scope.close();
-      ActorSystem actorSystem = AkkaHttpClientUtil.getActorSystem(thiz);
-      if (actorSystem == null) {
-        return;
-      }
-      if (throwable == null) {
-        responseFuture.onComplete(
-            new OnCompleteHandler(context, request), actorSystem.dispatcher());
-      } else {
-        instrumenter().end(context, request, null, throwable);
-      }
-      if (responseFuture != null) {
-        responseFuture =
-            FutureWrapper.wrap(responseFuture, actorSystem.dispatcher(), currentContext());
+      AdviceScope adviceScope = (AdviceScope) enterResult[0];
+      if (adviceScope == null) {
+        return responseFuture;
       }
+      ActorSystem actorSystem = AkkaHttpClientUtil.getActorSystem(thiz);
+      return adviceScope.end(actorSystem, (HttpRequest) enterResult[1], responseFuture, throwable);
     }
   }
 }

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/AkkaHttpServerInstrumentationModule.java

@@ -22,6 +22,11 @@ public AkkaHttpServerInstrumentationModule() {
     super("akka-http", "akka-http-10.0", "akka-http-server");
   }
 
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+
   @Override
   public String getModuleGroup() {
     return "akka-http";

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/AkkaHttpServerSourceInstrumentation.java

@@ -14,6 +14,8 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -33,10 +35,11 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class AkkaBindAndHandleAdvice {
 
+    @AssignReturned.ToArguments(@ToArgument(0))
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void wrapHandler(
-        @Advice.Argument(value = 0, readOnly = false) Flow<HttpRequest, HttpResponse, ?> handler) {
-      handler = AkkaFlowWrapper.wrap(handler);
+    public static Flow<HttpRequest, HttpResponse, ?> wrapHandler(
+        @Advice.Argument(0) Flow<HttpRequest, HttpResponse, ?> handler) {
+      return AkkaFlowWrapper.wrap(handler);
     }
   }
 }