diff --git a/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/HandlerWrapper.java b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/HandlerWrapper.java new file mode 100644 index 0000000000..53e4f302c1 --- /dev/null +++ b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/HandlerWrapper.java @@ -0,0 +1,35 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql; + +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.vertx.core.Handler; + +public class HandlerWrapper implements Handler { + private final Handler delegate; + private final Context context; + + private HandlerWrapper(Handler delegate, Context context) { + this.delegate = delegate; + this.context = context; + } + + public static Handler wrap(Handler handler) { + Context current = Context.current(); + if (handler != null && !(handler instanceof HandlerWrapper) && current != Context.root()) { + handler = new HandlerWrapper<>(handler, current); + } + return handler; + } + + @Override + public void handle(T t) { + try (Scope ignore = context.makeCurrent()) { + delegate.handle(t); + } + } +} diff --git a/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/TransactionImplInstrumentation.java b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/TransactionImplInstrumentation.java new file mode 100644 index 0000000000..464479dd63 --- /dev/null +++ b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/TransactionImplInstrumentation.java @@ -0,0 +1,39 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.vertx.v4_0.sql; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.returns; + +import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; +import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import io.vertx.core.Handler; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class TransactionImplInstrumentation implements TypeInstrumentation { + + @Override + public ElementMatcher typeMatcher() { + return named("io.vertx.sqlclient.impl.TransactionImpl"); + } + + @Override + public void transform(TypeTransformer transformer) { + transformer.applyAdviceToMethod( + named("wrap").and(returns(named("io.vertx.core.Handler"))), + TransactionImplInstrumentation.class.getName() + "$WrapHandlerAdvice"); + } + + @SuppressWarnings("unused") + public static class WrapHandlerAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void wrapHandler(@Advice.Return(readOnly = false) Handler handler) { + handler = HandlerWrapper.wrap(handler); + } + } +} diff --git a/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientInstrumentationModule.java b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientInstrumentationModule.java index 7f71cf586b..b4ccf21011 100644 --- a/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientInstrumentationModule.java +++ b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientInstrumentationModule.java @@ -30,6 +30,7 @@ public class VertxSqlClientInstrumentationModule extends InstrumentationModule { new PoolInstrumentation(), new SqlClientBaseInstrumentation(), new QueryExecutorInstrumentation(), - new QueryResultBuilderInstrumentation()); + new QueryResultBuilderInstrumentation(), + new TransactionImplInstrumentation()); } } diff --git a/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientTest.java b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientTest.java index 088bb42c77..80a94702f5 100644 --- a/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientTest.java +++ b/instrumentation/vertx/vertx-sql-client-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/vertx/v4_0/sql/VertxSqlClientTest.java @@ -203,6 +203,10 @@ class VertxSqlClientTest { .toCompletableFuture() .get(30, TimeUnit.SECONDS); + assertPreparedSelect(); + } + + private static void assertPreparedSelect() { testing.waitAndAssertTraces( trace -> trace.hasSpansSatisfyingExactly( @@ -250,4 +254,38 @@ class VertxSqlClientTest { equalTo(NET_PEER_NAME, "localhost"), equalTo(NET_PEER_PORT, port)))); } + + @Test + void testWithTransaction() throws Exception { + testing + .runWithSpan( + "parent", + () -> + pool.withTransaction( + conn -> + conn.preparedQuery("select * from test where id = $1") + .execute(Tuple.of(1)))) + .toCompletionStage() + .toCompletableFuture() + .get(30, TimeUnit.SECONDS); + + assertPreparedSelect(); + } + + @Test + void testWithConnection() throws Exception { + testing + .runWithSpan( + "parent", + () -> + pool.withConnection( + conn -> + conn.preparedQuery("select * from test where id = $1") + .execute(Tuple.of(1)))) + .toCompletionStage() + .toCompletableFuture() + .get(30, TimeUnit.SECONDS); + + assertPreparedSelect(); + } }