Fix vertx sql client instumentation for 5.1

Author: laurit

.github/config/latest-dep-versions.json

@@ -212,6 +212,7 @@
     "io.vertx:vertx-codegen#+": "5.1.0",
     "io.vertx:vertx-codegen#3.+": "3.9.16",
     "io.vertx:vertx-codegen#4.+": "4.5.27",
+    "io.vertx:vertx-codegen#5.0.+": "5.0.12",
     "io.vertx:vertx-core#+": "5.1.0",
     "io.vertx:vertx-core#3.+": "3.9.16",
     "io.vertx:vertx-core#4.+": "4.5.27",
@@ -220,18 +221,21 @@
     "io.vertx:vertx-jdbc-client#+": "5.1.0",
     "io.vertx:vertx-jdbc-client#3.+": "3.9.16",
     "io.vertx:vertx-jdbc-client#4.+": "4.5.27",
+    "io.vertx:vertx-jdbc-client#5.0.+": "5.0.12",
     "io.vertx:vertx-kafka-client#+": "5.1.0",
     "io.vertx:vertx-kafka-client#3.+": "3.9.16",
     "io.vertx:vertx-kafka-client#4.+": "4.5.27",
     "io.vertx:vertx-lang-kotlin-coroutines#+": "5.1.0",
     "io.vertx:vertx-pg-client#+": "5.1.0",
     "io.vertx:vertx-pg-client#4.+": "4.5.27",
+    "io.vertx:vertx-pg-client#5.0.+": "5.0.12",
     "io.vertx:vertx-redis-client#+": "5.1.0",
     "io.vertx:vertx-rx-java2#+": "5.1.0",
     "io.vertx:vertx-rx-java2#3.+": "3.9.16",
     "io.vertx:vertx-rx-java2#4.+": "4.5.27",
     "io.vertx:vertx-sql-client#+": "5.1.0",
     "io.vertx:vertx-sql-client#4.+": "4.5.27",
+    "io.vertx:vertx-sql-client#5.0.+": "5.0.12",
     "io.vertx:vertx-web#+": "5.1.0",
     "io.vertx:vertx-web#3.+": "3.9.16",
     "io.vertx:vertx-web#4.+": "4.5.27",

instrumentation/vertx/vertx-sql-client/vertx-sql-client-5.0/javaagent/build.gradle.kts

@@ -20,14 +20,16 @@ dependencies {
   library("io.vertx:vertx-sql-client:$version")
   library("io.vertx:vertx-codegen:$version")
 
+  compileOnly(project(":muzzle")) // For @NoMuzzle
+
   implementation(project(":instrumentation:vertx:vertx-sql-client:vertx-sql-client-common-4.0:javaagent"))
 
   testInstrumentation(project(":instrumentation:jdbc:javaagent"))
   testInstrumentation(project(":instrumentation:netty:netty-4.1:javaagent"))
   testInstrumentation(project(":instrumentation:vertx:vertx-sql-client:vertx-sql-client-4.0:javaagent"))
 
   testLibrary("io.vertx:vertx-pg-client:$version")
-  testImplementation("io.vertx:vertx-jdbc-client:$version")
+  testLibrary("io.vertx:vertx-jdbc-client:$version")
   testImplementation("io.agroal:agroal-pool:2.5")
   testImplementation("org.hsqldb:hsqldb:2.3.4")
 }

instrumentation/vertx/vertx-sql-client/vertx-sql-client-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/sqlclient/v5_0/CommandSchedulerInstrumentation.java

@@ -8,13 +8,13 @@
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 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 io.vertx.sqlclient.internal.command.CommandBase;
 import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -49,7 +49,12 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         named("schedule")
-            .and(takesArgument(0, named("io.vertx.sqlclient.internal.command.CommandBase")))
+            .and(
+                takesArgument(
+                    0,
+                    namedOneOf(
+                        "io.vertx.sqlclient.internal.command.CommandBase",
+                        "io.vertx.sqlclient.spi.protocol.CommandBase")))
             .and(takesArgument(1, named("io.vertx.core.Completable"))),
         getClass().getName() + "$ScheduleAdvice");
   }
@@ -59,7 +64,7 @@ public static class ScheduleAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
     @Nullable
-    public static Scope onEnter(@Advice.Argument(0) CommandBase<?> command) {
+    public static Scope onEnter(@Advice.Argument(0) Object command) {
       Context stored = VertxSqlClientSingletons.getCommandContext(command);
       if (stored == null) {
         // First schedule call (query executor → pool or direct connection).

instrumentation/vertx/vertx-sql-client/vertx-sql-client-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/sqlclient/v5_0/VertxSqlClientInstrumentationModule.java

@@ -10,10 +10,12 @@
 import static java.util.Collections.singletonList;
 
 import com.google.auto.service.AutoService;
+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 java.util.List;
+import java.util.function.BiConsumer;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -51,4 +53,16 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new QueryResultBuilderInstrumentation(),
         new TransactionImplInstrumentation());
   }
+
+  @Override
+  public void registerVirtualFields(BiConsumer<String, String> virtualFieldRegistrar) {
+    // we add the virtual field to CommandBase manually because it is in different package in 5.0
+    // and 5.1
+    // used in 5.0
+    virtualFieldRegistrar.accept(
+        "io.vertx.sqlclient.internal.command.CommandBase", Context.class.getName());
+    // used in 5.1
+    virtualFieldRegistrar.accept(
+        "io.vertx.sqlclient.spi.protocol.CommandBase", Context.class.getName());
+  }
 }

instrumentation/vertx/vertx-sql-client/vertx-sql-client-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/sqlclient/v5_0/VertxSqlClientSingletons.java

@@ -12,12 +12,12 @@
 import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.instrumentation.vertx.sqlclient.common.v4_0.VertxSqlClientRequest;
 import io.opentelemetry.javaagent.instrumentation.vertx.sqlclient.common.v4_0.VertxSqlInstrumenterFactory;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
 import io.vertx.core.Future;
 import io.vertx.sqlclient.Pool;
 import io.vertx.sqlclient.SqlConnectOptions;
 import io.vertx.sqlclient.SqlConnection;
 import io.vertx.sqlclient.internal.SqlClientBase;
-import io.vertx.sqlclient.internal.command.CommandBase;
 import javax.annotation.Nullable;
 
 public class VertxSqlClientSingletons {
@@ -34,26 +34,47 @@ public class VertxSqlClientSingletons {
   private static final VirtualField<SqlClientBase, SqlConnectOptions> connectOptionsField =
       VirtualField.find(SqlClientBase.class, SqlConnectOptions.class);
 
-  // CommandBase is a generic type; VirtualField.find requires the raw type
-  @SuppressWarnings("rawtypes")
-  private static final VirtualField<CommandBase, Context> commandContextField =
-      VirtualField.find(CommandBase.class, Context.class);
+  @Nullable
+  private static final VirtualField<Object, Context> commandContextField =
+      getCommandContextVirtualField();
 
   public static Instrumenter<VertxSqlClientRequest, Void> instrumenter() {
     return instrumenter;
   }
 
-  // CommandBase is a generic type used as VirtualField key
-  @SuppressWarnings("rawtypes")
+  @NoMuzzle // to skip virtual field detection in this method
+  @SuppressWarnings("unchecked") // virtual field key type is not known at compile time
+  private static VirtualField<Object, Context> getCommandContextVirtualField() {
+    // CommandBase that we want to attach context to is in different packages in 5.0 and 5.1
+    Class<?> commandClass = null;
+    try {
+      // 5.0.0
+      commandClass = Class.forName("io.vertx.sqlclient.internal.command.CommandBase");
+    } catch (ClassNotFoundException ignored) {
+      // ignored
+    }
+    if (commandClass == null) {
+      try {
+        // 5.1.0
+        commandClass = Class.forName("io.vertx.sqlclient.spi.protocol.CommandBase");
+      } catch (ClassNotFoundException ignored) {
+        // ignored
+      }
+    }
+    return commandClass != null
+        ? (VirtualField<Object, Context>) VirtualField.find(commandClass, Context.class)
+        : null;
+  }
+
   @Nullable
-  public static Context getCommandContext(CommandBase<?> command) {
-    return commandContextField.get(command);
+  public static Context getCommandContext(Object command) {
+    return commandContextField != null ? commandContextField.get(command) : null;
   }
 
-  // CommandBase is a generic type used as VirtualField key
-  @SuppressWarnings("rawtypes")
-  public static void setCommandContext(CommandBase<?> command, Context context) {
-    commandContextField.set(command, context);
+  public static void setCommandContext(Object command, Context context) {
+    if (commandContextField != null) {
+      commandContextField.set(command, context);
+    }
   }
 
   public static void storePoolDbSystem(Pool pool, String dbSystem) {

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/internal/ExperimentalInstrumentationModule.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiConsumer;
 import net.bytebuddy.utility.JavaModule;
 
 /**
@@ -19,6 +20,12 @@
  */
 public interface ExperimentalInstrumentationModule {
 
+  /**
+   * Register virtual field. First argument for the consumer is dot class name of the type where the
+   * field is added and the second argument is the dot class name of the field type.
+   */
+  default void registerVirtualFields(BiConsumer<String, String> virtualFieldRegistrar) {}
+
   /**
    * Returns a list of helper classes that will be defined in the class loader of the instrumented
    * library.

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/field/VirtualFieldImplementationInstallerFactory.java

@@ -9,6 +9,7 @@
 
 import io.opentelemetry.instrumentation.api.internal.RuntimeVirtualFieldSupplier;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
 import io.opentelemetry.javaagent.tooling.muzzle.VirtualFieldMappings;
@@ -31,6 +32,10 @@ public static VirtualFieldImplementationInstallerFactory getInstance() {
 
   public VirtualFieldImplementationInstaller create(InstrumentationModule instrumentationModule) {
     VirtualFieldMappingsBuilderImpl builder = new VirtualFieldMappingsBuilderImpl();
+    if (instrumentationModule instanceof ExperimentalInstrumentationModule) {
+      ((ExperimentalInstrumentationModule) instrumentationModule)
+          .registerVirtualFields(builder::register);
+    }
     if (instrumentationModule instanceof InstrumentationModuleMuzzle) {
       ((InstrumentationModuleMuzzle) instrumentationModule).registerMuzzleVirtualFields(builder);
     } else {

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/ReferenceCollectingClassVisitor.java

@@ -275,16 +275,21 @@ public void visitEnd() {
                 }
               }
             }
-            MethodVisitor target =
-                skip ? methodVisitor : new AdviceReferenceMethodVisitor(methodVisitor);
+            MethodVisitor target;
+            if (skip) {
+              target = methodVisitor;
+            } else {
+              target = new AdviceReferenceMethodVisitor(methodVisitor);
+              target = new VirtualFieldCollectingMethodVisitor(target);
+            }
             if (target != null) {
               accept(target);
             }
           }
         };
     // Additional references we could check
     // - Classes in signature (return type, params) and visible from this package
-    return new VirtualFieldCollectingMethodVisitor(methodNode);
+    return methodNode;
   }
 
   private static VisibilityFlag computeVisibilityFlag(int access) {