vertx-redis-client-4.0

Author: SylvainJuge

instrumentation/vertx/vertx-redis-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/redis/RedisStandaloneConnectionInstrumentation.java

@@ -21,7 +21,9 @@
 import io.vertx.redis.client.impl.RedisStandaloneConnection;
 import io.vertx.redis.client.impl.RedisURI;
 import io.vertx.redis.client.impl.RequestUtil;
+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;
 
@@ -40,54 +42,80 @@ public void transform(TypeTransformer transformer) {
 
   @SuppressWarnings("unused")
   public static class SendAdvice {
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(
-        @Advice.This RedisStandaloneConnection connection,
-        @Advice.Argument(0) Request request,
-        @Advice.FieldValue("netSocket") NetSocket netSocket,
-        @Advice.Local("otelRequest") VertxRedisClientRequest otelRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (request == null) {
-        return;
-      }
+    public static class AdviceScope {
+      public VertxRedisClientRequest otelRequest;
+      public Context context;
+      public Scope scope;
+
+      @Nullable
+      public static AdviceScope start(
+          RedisStandaloneConnection connection, @Nullable Request request, NetSocket netSocket) {
 
-      String commandName = VertxRedisClientSingletons.getCommandName(request.command());
-      RedisURI redisUri = VertxRedisClientSingletons.getRedisUri(connection);
-      if (commandName == null || redisUri == null) {
-        return;
+        if (request == null) {
+          return null;
+        }
+
+        String commandName = VertxRedisClientSingletons.getCommandName(request.command());
+        RedisURI redisUri = VertxRedisClientSingletons.getRedisUri(connection);
+        if (commandName == null || redisUri == null) {
+          return null;
+        }
+
+        VertxRedisClientRequest otelRequest =
+            new VertxRedisClientRequest(
+                commandName, RequestUtil.getArgs(request), redisUri, netSocket);
+        Context parentContext = currentContext();
+        if (!instrumenter().shouldStart(parentContext, otelRequest)) {
+          return null;
+        }
+        AdviceScope locals = new AdviceScope();
+        locals.otelRequest = otelRequest;
+        locals.context = instrumenter().start(parentContext, locals.otelRequest);
+        locals.scope = locals.context.makeCurrent();
+        return locals;
       }
 
-      otelRequest =
-          new VertxRedisClientRequest(
-              commandName, RequestUtil.getArgs(request), redisUri, netSocket);
-      Context parentContext = currentContext();
-      if (!instrumenter().shouldStart(parentContext, otelRequest)) {
-        return;
+      @Nullable
+      public Future<Response> end(
+          @Nullable Future<Response> responseFuture, @Nullable Throwable throwable) {
+        if (scope == null) {
+          return responseFuture;
+        }
+
+        scope.close();
+        if (throwable != null) {
+          instrumenter().end(context, otelRequest, null, throwable);
+        } else {
+          responseFuture =
+              VertxRedisClientSingletons.wrapEndSpan(responseFuture, context, otelRequest);
+        }
+        return responseFuture;
       }
+    }
+
+    @Nullable
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AdviceScope onEnter(
+        @Advice.This RedisStandaloneConnection connection,
+        @Advice.Argument(0) @Nullable Request request,
+        @Advice.FieldValue("netSocket") NetSocket netSocket) {
 
-      context = instrumenter().start(parentContext, otelRequest);
-      scope = context.makeCurrent();
+      return AdviceScope.start(connection, request, netSocket);
     }
 
+    @Nullable
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(
+    public static Future<Response> onExit(
         @Advice.Thrown Throwable throwable,
-        @Advice.Return(readOnly = false) Future<Response> responseFuture,
-        @Advice.Local("otelRequest") VertxRedisClientRequest otelRequest,
-        @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
-      if (scope == null) {
-        return;
-      }
+        @Advice.Return @Nullable Future<Response> responseFuture,
+        @Advice.Enter @Nullable AdviceScope adviceScope) {
 
-      scope.close();
-      if (throwable != null) {
-        instrumenter().end(context, otelRequest, null, throwable);
-      } else {
-        responseFuture =
-            VertxRedisClientSingletons.wrapEndSpan(responseFuture, context, otelRequest);
+      if (adviceScope != null) {
+        return adviceScope.end(responseFuture, throwable);
       }
+
+      return responseFuture;
     }
   }
 

instrumentation/vertx/vertx-redis-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/redis/VertxRedisClientInstrumentationModule.java

@@ -39,4 +39,9 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new RedisConnectionProviderInstrumentation(),
         new CommandImplInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }