Add vertx postgresql client instrumentation

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/build.gradle

@@ -0,0 +1,37 @@
+
+apply from: "$rootDir/gradle/java.gradle"
+
+muzzle {
+  pass {
+    group = 'io.vertx'
+    module = 'vertx-pg-client'
+    versions = '[4.2.0)'
+    assertInverse = true
+  }
+}
+
+addTestSuiteForDir('latestDepTest', 'test')
+addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test')
+
+apply from: "$rootDir/gradle/configure_tests.gradle"
+
+latestDepTest {
+  finalizedBy 'latestDepForkedTest'
+}
+
+dependencies {
+  compileOnly group: 'io.vertx', name: 'vertx-pg-client', version: '4.2.0'
+
+  testImplementation group: 'io.vertx', name: 'vertx-pg-client', version: '4.2.0'
+
+  // This is needed for the test container to start
+  testImplementation group: 'com.ongres.scram', name: 'client', version: '2.1'
+  testImplementation group: 'org.postgresql', name: 'postgresql', version: '42.7.4'
+  testImplementation group: 'org.testcontainers', name: 'postgresql', version: libs.versions.testcontainers.get()
+
+  latestDepTestImplementation group: 'io.vertx', name: 'vertx-pg-client', version: '4.2.+'
+}
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/CursorImplInstrumentation.java

@@ -0,0 +1,45 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import java.util.Map;
+
+@AutoService(InstrumenterModule.class)
+public class CursorImplInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public CursorImplInstrumentation() {
+    super("vertx", "vertx-sql-client");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("io.vertx.sqlclient.PreparedStatement", "datadog.trace.api.Pair");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".QueryResultHandlerWrapper", packageName + ".VertxSqlClientDecorator",
+    };
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "io.vertx.sqlclient.impl.CursorImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("read"))
+            .and(takesArgument(1, named("io.vertx.core.Handler"))),
+        packageName + ".CursorReadAdvice");
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/CursorReadAdvice.java

@@ -0,0 +1,48 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.*;
+import static datadog.trace.instrumentation.vertx_pg_client_4_2_0.VertxSqlClientDecorator.DECORATE;
+
+import datadog.trace.api.Pair;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.sqlclient.PreparedStatement;
+import io.vertx.sqlclient.Row;
+import io.vertx.sqlclient.RowSet;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+
+public class CursorReadAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope beforeRead(
+      @Advice.Argument(value = 1, readOnly = false) Handler<AsyncResult<RowSet<Row>>> handler,
+      @Advice.FieldValue(value = "ps", typing = Assigner.Typing.DYNAMIC)
+          final PreparedStatement ps) {
+    if (handler instanceof QueryResultHandlerWrapper) {
+      return null;
+    }
+    final AgentSpan parentSpan = activeSpan();
+    final AgentScope.Continuation parentContinuation =
+        null == parentSpan ? null : captureSpan(parentSpan);
+    final AgentSpan clientSpan =
+        DECORATE.startAndDecorateSpanForStatement(
+            ps, InstrumentationContext.get(PreparedStatement.class, Pair.class), true);
+    if (null == clientSpan) {
+      return null;
+    }
+    handler = new QueryResultHandlerWrapper<>(handler, clientSpan, parentContinuation);
+
+    return activateSpan(clientSpan, true);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void afterRead(
+      @Advice.Thrown final Throwable throwable, @Advice.Enter final AgentScope clientScope) {
+    if (null != clientScope) {
+      clientScope.close();
+    }
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PgConnectionFactoryConstructorAdvice.java

@@ -0,0 +1,25 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import io.vertx.pgclient.PgConnectOptions;
+import io.vertx.pgclient.impl.PgConnectionFactory;
+import net.bytebuddy.asm.Advice;
+
+public class PgConnectionFactoryConstructorAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void afterConstructor(
+      @Advice.This final PgConnectionFactory zis,
+      @Advice.Argument(1) final PgConnectOptions options) {
+    DBInfo.Builder builder = DBInfo.DEFAULT.toBuilder();
+    DBInfo info =
+        builder
+            .host(options.getHost())
+            .port(options.getPort())
+            .db(options.getDatabase())
+            .user(options.getUser())
+            .type("postgresql")
+            .build();
+    InstrumentationContext.get(PgConnectionFactory.class, DBInfo.class).put(zis, info);
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PgConnectionFactoryInstrumentation.java

@@ -0,0 +1,38 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import java.util.Map;
+
+@AutoService(InstrumenterModule.class)
+public class PgConnectionFactoryInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public PgConnectionFactoryInstrumentation() {
+    super("vertx", "vertx-sql-client");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("io.vertx.pgclient.impl.PgConnectionFactory", DBInfo.class.getName());
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "io.vertx.pgclient.impl.PgConnectionFactory";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor()
+            .and(takesArguments(2))
+            .and(takesArgument(1, named("io.vertx.pgclient.PgConnectOptions"))),
+        packageName + ".PgConnectionFactoryConstructorAdvice");
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PgConnectionImplConstructorAdvice.java

@@ -0,0 +1,16 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import io.vertx.pgclient.impl.PgConnectionFactory;
+import io.vertx.sqlclient.SqlClient;
+import net.bytebuddy.asm.Advice;
+
+public class PgConnectionImplConstructorAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void afterConstructor(
+      @Advice.This final SqlClient zis, @Advice.Argument(0) final PgConnectionFactory factory) {
+    InstrumentationContext.get(SqlClient.class, DBInfo.class)
+        .put(zis, InstrumentationContext.get(PgConnectionFactory.class, DBInfo.class).get(factory));
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PgConnectionImplInstrumentation.java

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import java.util.HashMap;
+import java.util.Map;
+
+@AutoService(InstrumenterModule.class)
+public class PgConnectionImplInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public PgConnectionImplInstrumentation() {
+    super("vertx", "vertx-sql-client");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> contextStores = new HashMap<>();
+    contextStores.put("io.vertx.pgclient.impl.PgConnectionFactory", DBInfo.class.getName());
+    contextStores.put("io.vertx.sqlclient.SqlClient", DBInfo.class.getName());
+    return contextStores;
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "io.vertx.pgclient.impl.PgConnectionImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor().and(takesArgument(0, named("io.vertx.pgclient.impl.PgConnectionFactory"))),
+        packageName + ".PgConnectionImplConstructorAdvice");
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PgPoolImplAdvice.java

@@ -0,0 +1,25 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import io.vertx.pgclient.PgConnectOptions;
+import io.vertx.sqlclient.SqlClient;
+import net.bytebuddy.asm.Advice;
+
+public class PgPoolImplAdvice {
+
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void afterCreate(
+      @Advice.Return final SqlClient zis, @Advice.Argument(2) PgConnectOptions options) {
+    DBInfo.Builder builder = DBInfo.DEFAULT.toBuilder();
+    DBInfo info =
+        builder
+            .host(options.getHost())
+            .port(options.getPort())
+            .db(options.getDatabase())
+            .user(options.getUser())
+            .type("postgresql")
+            .build();
+    InstrumentationContext.get(SqlClient.class, DBInfo.class).put(zis, info);
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PgPoolImplInstrumentation.java

@@ -0,0 +1,41 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import java.util.Map;
+
+@AutoService(InstrumenterModule.class)
+public class PgPoolImplInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public PgPoolImplInstrumentation() {
+    super("vertx", "vertx-sql-client");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("io.vertx.sqlclient.SqlClient", DBInfo.class.getName());
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "io.vertx.pgclient.impl.PgPoolImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isStatic()
+            .and(isPublic())
+            .and(isMethod())
+            .and(named("create"))
+            .and(takesArguments(4))
+            .and(takesArgument(2, named("io.vertx.pgclient.PgConnectOptions"))),
+        packageName + ".PgPoolImplAdvice");
+  }
+}

dd-java-agent/instrumentation/vertx-pg-client-4.2.0/src/main/java/datadog/trace/instrumentation/vertx_pg_client_4_2_0/PrepareHandlerWrapper.java

@@ -0,0 +1,32 @@
+package datadog.trace.instrumentation.vertx_pg_client_4_2_0;
+
+import datadog.trace.api.Pair;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.sqlclient.PreparedStatement;
+
+public class PrepareHandlerWrapper implements Handler<AsyncResult<PreparedStatement>> {
+  private final Handler<AsyncResult<PreparedStatement>> handler;
+  private final ContextStore<PreparedStatement, Pair> contextStore;
+  private final Pair<DBInfo, DBQueryInfo> queryInfo;
+
+  public PrepareHandlerWrapper(
+      Handler<AsyncResult<PreparedStatement>> handler,
+      ContextStore<PreparedStatement, Pair> contextStore,
+      Pair<DBInfo, DBQueryInfo> queryInfo) {
+    this.handler = handler;
+    this.contextStore = contextStore;
+    this.queryInfo = queryInfo;
+  }
+
+  @Override
+  public void handle(AsyncResult<PreparedStatement> event) {
+    if (event.succeeded()) {
+      contextStore.put(event.result(), queryInfo);
+    }
+    handler.handle(event);
+  }
+}