From acd1c92173ebcd1489ee25f490bc107c0a3aa6bb Mon Sep 17 00:00:00 2001 From: Lauri Tulmin Date: Thu, 22 Jun 2023 12:44:30 +0300 Subject: [PATCH] Fix mongo latest dep tests (#8785) --- .../v4_0/BaseClusterInstrumentation.java | 18 ++++++++ ...aultConnectionPoolTaskInstrumentation.java | 43 +++++++++++++++++++ .../MongoClientInstrumentationModule.java | 1 + .../mongo/v4_0/TaskWrapper.java | 27 ++++++++++++ 4 files changed, 89 insertions(+) create mode 100644 instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/DefaultConnectionPoolTaskInstrumentation.java create mode 100644 instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/TaskWrapper.java diff --git a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/BaseClusterInstrumentation.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/BaseClusterInstrumentation.java index 0eff79c30dd2..24994d4d1972 100644 --- a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/BaseClusterInstrumentation.java +++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/BaseClusterInstrumentation.java @@ -34,6 +34,14 @@ public void transform(TypeTransformer transformer) { .and(takesArgument(0, named("com.mongodb.selector.ServerSelector"))) .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))), this.getClass().getName() + "$SingleResultCallbackArg1Advice"); + + transformer.applyAdviceToMethod( + isMethod() + .and(isPublic()) + .and(named("selectServerAsync")) + .and(takesArgument(0, named("com.mongodb.selector.ServerSelector"))) + .and(takesArgument(2, named("com.mongodb.internal.async.SingleResultCallback"))), + this.getClass().getName() + "$SingleResultCallbackArg2Advice"); } @SuppressWarnings("unused") @@ -45,4 +53,14 @@ public static void wrapCallback( callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback); } } + + @SuppressWarnings("unused") + public static class SingleResultCallbackArg2Advice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void wrapCallback( + @Advice.Argument(value = 2, readOnly = false) SingleResultCallback callback) { + callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback); + } + } } diff --git a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/DefaultConnectionPoolTaskInstrumentation.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/DefaultConnectionPoolTaskInstrumentation.java new file mode 100644 index 000000000000..a87070f926eb --- /dev/null +++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/DefaultConnectionPoolTaskInstrumentation.java @@ -0,0 +1,43 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.mongo.v4_0; + +import static net.bytebuddy.matcher.ElementMatchers.isConstructor; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; + +import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge; +import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; +import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import java.util.function.Consumer; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class DefaultConnectionPoolTaskInstrumentation implements TypeInstrumentation { + @Override + public ElementMatcher typeMatcher() { + return named("com.mongodb.internal.connection.DefaultConnectionPool$Task"); + } + + @Override + public void transform(TypeTransformer transformer) { + // outer class this is passed as arg 0 to constructor + transformer.applyAdviceToMethod( + isConstructor().and(takesArgument(2, Consumer.class)), + this.getClass().getName() + "$TaskAdvice"); + } + + @SuppressWarnings("unused") + public static class TaskAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void wrapCallback( + @Advice.Argument(value = 2, readOnly = false) Consumer action) { + action = new TaskWrapper(Java8BytecodeBridge.currentContext(), action); + } + } +} diff --git a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java index ecc578f7bced..25b8abe5c0c2 100644 --- a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java +++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java @@ -33,6 +33,7 @@ public List typeInstrumentations() { new InternalStreamConnectionInstrumentation(), new BaseClusterInstrumentation(), new DefaultConnectionPoolInstrumentation(), + new DefaultConnectionPoolTaskInstrumentation(), new AsyncWorkManagerInstrumentation()); } } diff --git a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/TaskWrapper.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/TaskWrapper.java new file mode 100644 index 000000000000..ea4f72229143 --- /dev/null +++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/TaskWrapper.java @@ -0,0 +1,27 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.mongo.v4_0; + +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import java.util.function.Consumer; + +public class TaskWrapper implements Consumer { + private final Context context; + private final Consumer delegate; + + public TaskWrapper(Context context, Consumer delegate) { + this.context = context; + this.delegate = delegate; + } + + @Override + public void accept(Object value) { + try (Scope ignored = context.makeCurrent()) { + delegate.accept(value); + } + } +}