Add instrumentation for vertx-sql-client

docs/supported-libraries.md

@@ -131,6 +131,7 @@ These are the supported libraries and frameworks:
 | [Vert.x HttpClient](https://vertx.io/docs/apidocs/io/vertx/core/http/HttpClient.html)                                                       | 3.0+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | [HTTP Client Spans], [HTTP Client Metrics]                                             |
 | [Vert.x Kafka Client](https://vertx.io/docs/vertx-kafka-client/java/)                                                                       | 3.6+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | [Messaging Spans]                                                                      |
 | [Vert.x RxJava2](https://vertx.io/docs/vertx-rx/java2/)                                                                                     | 3.5+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | context propagation only                                                               |
+| [Vert.x SQL Client](https://github.com/eclipse-vertx/vertx-sql-client/)                                                                     | 4.0+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | [Database Client Spans]                                                                |
 | [Vibur DBCP](https://www.vibur.org/)                                                                                                        | 11.0+                         | [opentelemetry-vibur-dbcp-11.0](../instrumentation/vibur-dbcp-11.0/library)                                                                                                                                                                                                                                                                                                             | [Database Pool Metrics]                                                                |
 | [ZIO](https://zio.dev/)                                                                                                                     | 2.0.0+                        | N/A                                                                                                                                                                                                                                                                                                                                                                                     | Context propagation                                                                    |
 

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

@@ -0,0 +1,25 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.vertx")
+    module.set("vertx-sql-client")
+    versions.set("[4.0.0,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  library("io.vertx:vertx-sql-client:4.1.0")
+  compileOnly("io.vertx:vertx-codegen:4.1.0")
+
+  testLibrary("io.vertx:vertx-pg-client:4.1.0")
+  testLibrary("io.vertx:vertx-codegen:4.1.0")
+  testLibrary("io.vertx:vertx-opentelemetry:4.1.0")
+}
+
+tasks.withType<Test>().configureEach {
+  usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
+}

instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/PoolInstrumentation.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql;
+
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.setSqlConnectOptions;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.vertx.sqlclient.Pool;
+import io.vertx.sqlclient.SqlConnectOptions;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class PoolInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.vertx.sqlclient.Pool");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("pool")
+            .and(takesArguments(3))
+            .and(takesArgument(1, named("io.vertx.sqlclient.SqlConnectOptions")))
+            .and(returns(named("io.vertx.sqlclient.Pool"))),
+        PoolInstrumentation.class.getName() + "$PoolAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class PoolAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(1) SqlConnectOptions sqlConnectOptions,
+        @Advice.Local("otelCallDepth") CallDepth callDepth) {
+      callDepth = CallDepth.forClass(Pool.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
+      // set connection options to ThreadLocal, they will be read in SqlClientBase constructor
+      setSqlConnectOptions(sqlConnectOptions);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Local("otelCallDepth") CallDepth callDepth) {
+      if (callDepth.decrementAndGet() > 0) {
+        return;
+      }
+
+      setSqlConnectOptions(null);
+    }
+  }
+}

instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/QueryExecutorInstrumentation.java

@@ -0,0 +1,127 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql;
+
+import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.OTEL_CONTEXT_KEY;
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.OTEL_PARENT_CONTEXT_KEY;
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.OTEL_REQUEST_KEY;
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.getSqlConnectOptions;
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.instrumenter;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.vertx.core.impl.future.PromiseInternal;
+import io.vertx.sqlclient.impl.PreparedStatement;
+import io.vertx.sqlclient.impl.QueryExecutorUtil;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class QueryExecutorInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.vertx.sqlclient.impl.QueryExecutor");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isConstructor(), QueryExecutorInstrumentation.class.getName() + "$ConstructorAdvice");
+    transformer.applyAdviceToMethod(
+        namedOneOf("executeSimpleQuery", "executeExtendedQuery", "executeBatchQuery"),
+        QueryExecutorInstrumentation.class.getName() + "$QueryAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class ConstructorAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.This Object queryExecutor) {
+      // copy connection options from ThreadLocal to VirtualField
+      QueryExecutorUtil.setConnectOptions(queryExecutor, getSqlConnectOptions());
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class QueryAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This Object queryExecutor,
+        @Advice.AllArguments Object[] arguments,
+        @Advice.Local("otelCallDepth") CallDepth callDepth,
+        @Advice.Local("otelRequest") VertxSqlClientRequest otelRequest,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      callDepth = CallDepth.forClass(queryExecutor.getClass());
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
+      // The parameter we need are in different positions, we are not going to have separate
+      // advices for all of them. The method gets the statement either as String or
+      // PreparedStatement, use the first argument that is either of these. PromiseInternal is
+      // always at the end of the argument list.
+      String sql = null;
+      PromiseInternal<?> promiseInternal = null;
+      for (Object argument : arguments) {
+        if (sql == null) {
+          if (argument instanceof String) {
+            sql = (String) argument;
+          } else if (argument instanceof PreparedStatement) {
+            sql = ((PreparedStatement) argument).sql();
+          }
+        } else if (argument instanceof PromiseInternal) {
+          promiseInternal = (PromiseInternal) argument;
+        }
+      }
+      if (sql == null || promiseInternal == null) {
+        return;
+      }
+
+      otelRequest =
+          new VertxSqlClientRequest(sql, QueryExecutorUtil.getConnectOptions(queryExecutor));
+      Context parentContext = currentContext();
+      if (!instrumenter().shouldStart(parentContext, otelRequest)) {
+        return;
+      }
+
+      context = instrumenter().start(parentContext, otelRequest);
+      scope = context.makeCurrent();
+      promiseInternal.context().localContextData().put(OTEL_REQUEST_KEY, otelRequest);
+      promiseInternal.context().localContextData().put(OTEL_CONTEXT_KEY, context);
+      promiseInternal.context().localContextData().put(OTEL_PARENT_CONTEXT_KEY, parentContext);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelCallDepth") CallDepth callDepth,
+        @Advice.Local("otelRequest") VertxSqlClientRequest otelRequest,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      if (callDepth.decrementAndGet() > 0) {
+        return;
+      }
+      if (scope == null) {
+        return;
+      }
+
+      scope.close();
+      if (throwable != null) {
+        instrumenter().end(context, otelRequest, null, throwable);
+      }
+      // span will be ended in QueryResultBuilderInstrumentation
+    }
+  }
+}

instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/QueryResultBuilderInstrumentation.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql;
+
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.endQuerySpan;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.vertx.core.Context;
+import io.vertx.core.impl.ContextInternal;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class QueryResultBuilderInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.vertx.sqlclient.impl.QueryResultBuilder");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("tryComplete"),
+        QueryResultBuilderInstrumentation.class.getName() + "$CompleteAdvice");
+    transformer.applyAdviceToMethod(
+        named("tryFail").and(takesArguments(Throwable.class)),
+        QueryResultBuilderInstrumentation.class.getName() + "$FailAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class CompleteAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope onEnter(@Advice.FieldValue("context") Context vertxContext) {
+      ContextInternal contextInternal = (ContextInternal) vertxContext;
+      return endQuerySpan(contextInternal.localContextData(), null);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class FailAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope onEnter(
+        @Advice.Argument(0) Throwable throwable,
+        @Advice.FieldValue("context") Context vertxContext) {
+      ContextInternal contextInternal = (ContextInternal) vertxContext;
+      return endQuerySpan(contextInternal.localContextData(), throwable);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}

instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/SqlClientBaseInstrumentation.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql;
+
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.getSqlConnectOptions;
+import static io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql.VertxSqlClientSingletons.setSqlConnectOptions;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.vertx.sqlclient.SqlConnectOptions;
+import io.vertx.sqlclient.impl.SqlClientBase;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class SqlClientBaseInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.vertx.sqlclient.impl.SqlClientBase");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isConstructor(), SqlClientBaseInstrumentation.class.getName() + "$ConstructorAdvice");
+    transformer.applyAdviceToMethod(
+        namedOneOf("query", "preparedQuery"),
+        SqlClientBaseInstrumentation.class.getName() + "$QueryAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class ConstructorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.This SqlClientBase<?> sqlClientBase) {
+      // copy connection options from ThreadLocal to VirtualField
+      VirtualField<SqlClientBase<?>, SqlConnectOptions> virtualField =
+          VirtualField.find(SqlClientBase.class, SqlConnectOptions.class);
+      virtualField.set(sqlClientBase, getSqlConnectOptions());
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class QueryAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This SqlClientBase<?> sqlClientBase,
+        @Advice.Local("otelCallDepth") CallDepth callDepth) {
+      callDepth = CallDepth.forClass(SqlClientBase.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
+      // set connection options to ThreadLocal, they will be read in QueryExecutor constructor
+      VirtualField<SqlClientBase<?>, SqlConnectOptions> virtualField =
+          VirtualField.find(SqlClientBase.class, SqlConnectOptions.class);
+      SqlConnectOptions sqlConnectOptions = virtualField.get(sqlClientBase);
+      setSqlConnectOptions(sqlConnectOptions);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable, @Advice.Local("otelCallDepth") CallDepth callDepth) {
+      if (callDepth.decrementAndGet() > 0) {
+        return;
+      }
+
+      setSqlConnectOptions(null);
+    }
+  }
+}