diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/build.gradle.kts b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/build.gradle.kts index 25525d1484..9831c0b4d0 100644 --- a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/build.gradle.kts +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/build.gradle.kts @@ -52,6 +52,7 @@ testing { implementation("org.hibernate.reactive:hibernate-reactive-core:2.0.0.Final") implementation("io.vertx:vertx-pg-client:4.4.2") } + compileOnly("io.vertx:vertx-codegen:4.4.2") } } } diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive1Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveTest.java b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive1Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveTest.java index d52500d76b..c5acf36cdd 100644 --- a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive1Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveTest.java +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive1Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveTest.java @@ -123,28 +123,156 @@ class HibernateReactiveTest { testing.runWithSpan( "parent", () -> - Vertx.vertx() - .getOrCreateContext() - .runOnContext( - event -> - stageSessionFactory - .withSession( - session -> { - if (!Span.current().getSpanContext().isValid()) { - throw new IllegalStateException("missing parent span"); - } + runWithVertx( + () -> + stageSessionFactory + .withSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } - return session - .find(Value.class, 1L) - .thenAccept( - value -> testing.runWithSpan("callback", () -> {})); - }) - .thenAccept(unused -> latch.countDown()))); + return session + .find(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); latch.await(30, TimeUnit.SECONDS); assertTrace(); } + @Test + void testStageWithStatelessSession() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .withStatelessSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .get(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageSessionWithTransaction() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .withSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .withTransaction(transaction -> session.find(Value.class, 1L)) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageStatelessSessionWithTransaction() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .withStatelessSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .withTransaction(transaction -> session.get(Value.class, 1L)) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageOpenSession() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .openSession() + .thenApply( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .find(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageOpenStatelessSession() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .openStatelessSession() + .thenApply( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .get(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + private static void runWithVertx(Runnable runnable) { + Vertx.vertx().getOrCreateContext().runOnContext(event -> runnable.run()); + } + @SuppressWarnings("deprecation") // until old http semconv are dropped in 2.0 private static void assertTrace() { testing.waitAndAssertTraces( diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive2Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v2_0/HibernateReactiveTest.java b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive2Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v2_0/HibernateReactiveTest.java index 091691d081..dc506f3bf1 100644 --- a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive2Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v2_0/HibernateReactiveTest.java +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/hibernateReactive2Test/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v2_0/HibernateReactiveTest.java @@ -18,9 +18,11 @@ import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension; import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension; +import io.vertx.core.Vertx; import jakarta.persistence.EntityManagerFactory; import jakarta.persistence.Persistence; import java.time.Duration; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import org.hibernate.reactive.mutiny.Mutiny; import org.hibernate.reactive.stage.Stage; @@ -138,6 +140,131 @@ class HibernateReactiveTest { assertTrace(); } + @Test + void testStageWithStatelessSession() throws Exception { + testing + .runWithSpan( + "parent", + () -> + stageSessionFactory + .withStatelessSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .get(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .toCompletableFuture()) + .get(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageSessionWithTransaction() throws Exception { + testing + .runWithSpan( + "parent", + () -> + stageSessionFactory + .withSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .withTransaction(transaction -> session.find(Value.class, 1L)) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .toCompletableFuture()) + .get(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageStatelessSessionWithTransaction() throws Exception { + testing + .runWithSpan( + "parent", + () -> + stageSessionFactory + .withStatelessSession( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .withTransaction(transaction -> session.get(Value.class, 1L)) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .toCompletableFuture()) + .get(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageOpenSession() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .openSession() + .thenApply( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .find(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + @Test + void testStageOpenStatelessSession() throws Exception { + CountDownLatch latch = new CountDownLatch(1); + testing.runWithSpan( + "parent", + () -> + runWithVertx( + () -> + stageSessionFactory + .openStatelessSession() + .thenApply( + session -> { + if (!Span.current().getSpanContext().isValid()) { + throw new IllegalStateException("missing parent span"); + } + + return session + .get(Value.class, 1L) + .thenAccept(value -> testing.runWithSpan("callback", () -> {})); + }) + .thenAccept(unused -> latch.countDown()))); + latch.await(30, TimeUnit.SECONDS); + + assertTrace(); + } + + private static void runWithVertx(Runnable runnable) { + Vertx.vertx().getOrCreateContext().runOnContext(event -> runnable.run()); + } + @SuppressWarnings("deprecation") // until old http semconv are dropped in 2.0 private static void assertTrace() { testing.waitAndAssertTraces( diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/CompletionStageWrapper.java b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/CompletionStageWrapper.java new file mode 100644 index 0000000000..d9ab3cb667 --- /dev/null +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/CompletionStageWrapper.java @@ -0,0 +1,40 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.hibernate.reactive.v1_0; + +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; + +public final class CompletionStageWrapper { + + private CompletionStageWrapper() {} + + public static CompletionStage wrap(CompletionStage future) { + Context context = Context.current(); + if (context != Context.root()) { + return wrap(future, context); + } + return future; + } + + private static CompletionStage wrap(CompletionStage completionStage, Context context) { + CompletableFuture result = new CompletableFuture<>(); + completionStage.whenComplete( + (T value, Throwable throwable) -> { + try (Scope ignored = context.makeCurrent()) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + result.complete(value); + } + } + }); + + return result; + } +} diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveInstrumentationModule.java b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveInstrumentationModule.java index 3e032e91e1..ae3d953c5a 100644 --- a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveInstrumentationModule.java +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/HibernateReactiveInstrumentationModule.java @@ -22,6 +22,8 @@ public class HibernateReactiveInstrumentationModule extends InstrumentationModul @Override public List typeInstrumentations() { return asList( - new StageSessionFactoryInstrumentation(), new MutinySessionFactoryInstrumentation()); + new StageSessionFactoryInstrumentation(), + new StageSessionImplInstrumentation(), + new MutinySessionFactoryInstrumentation()); } } diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionFactoryInstrumentation.java b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionFactoryInstrumentation.java index e578c8e939..96e55ce81e 100644 --- a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionFactoryInstrumentation.java +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionFactoryInstrumentation.java @@ -7,10 +7,12 @@ package io.opentelemetry.javaagent.instrumentation.hibernate.reactive.v1_0; import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; +import static net.bytebuddy.matcher.ElementMatchers.returns; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import java.util.concurrent.CompletionStage; import java.util.function.Function; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; @@ -30,6 +32,9 @@ public class StageSessionFactoryInstrumentation implements TypeInstrumentation { transformer.applyAdviceToMethod( namedOneOf("withSession", "withStatelessSession").and(takesArgument(1, Function.class)), this.getClass().getName() + "$Function1Advice"); + transformer.applyAdviceToMethod( + namedOneOf("openSession", "openStatelessSession").and(returns(CompletionStage.class)), + this.getClass().getName() + "$OpenSessionAdvice"); } @SuppressWarnings("unused") @@ -49,4 +54,12 @@ public class StageSessionFactoryInstrumentation implements TypeInstrumentation { function = FunctionWrapper.wrap(function); } } + + @SuppressWarnings("unused") + public static class OpenSessionAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit(@Advice.Return(readOnly = false) CompletionStage completionStage) { + completionStage = CompletionStageWrapper.wrap(completionStage); + } + } } diff --git a/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionImplInstrumentation.java b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionImplInstrumentation.java new file mode 100644 index 0000000000..26d7d1cd1c --- /dev/null +++ b/instrumentation/hibernate/hibernate-reactive-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/hibernate/reactive/v1_0/StageSessionImplInstrumentation.java @@ -0,0 +1,50 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.hibernate.reactive.v1_0; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; +import static net.bytebuddy.matcher.ElementMatchers.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; + +import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; +import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import java.util.concurrent.CompletionStage; +import java.util.function.Function; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class StageSessionImplInstrumentation implements TypeInstrumentation { + @Override + public ElementMatcher typeMatcher() { + return namedOneOf( + "org.hibernate.reactive.stage.impl.StageSessionImpl", + "org.hibernate.reactive.stage.impl.StageStatelessSessionImpl"); + } + + @Override + public void transform(TypeTransformer transformer) { + transformer.applyAdviceToMethod( + named("withTransaction") + .and(takesArgument(0, Function.class).and(returns(CompletionStage.class))), + this.getClass().getName() + "$WithTransactionAdvice"); + } + + @SuppressWarnings("unused") + public static class WithTransactionAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(value = 0, readOnly = false) Function function) { + function = FunctionWrapper.wrap(function); + } + + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit(@Advice.Return(readOnly = false) CompletionStage completionStage) { + completionStage = CompletionStageWrapper.wrap(completionStage); + } + } +}