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 index a87070f926eb..2d142c66a7db 100644 --- 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 @@ -26,13 +26,18 @@ public ElementMatcher typeMatcher() { @Override public void transform(TypeTransformer transformer) { // outer class this is passed as arg 0 to constructor + // before 5.2.0 transformer.applyAdviceToMethod( isConstructor().and(takesArgument(2, Consumer.class)), - this.getClass().getName() + "$TaskAdvice"); + this.getClass().getName() + "$TaskArg2Advice"); + // since 5.2.0 + transformer.applyAdviceToMethod( + isConstructor().and(takesArgument(3, Consumer.class)), + this.getClass().getName() + "$TaskArg3Advice"); } @SuppressWarnings("unused") - public static class TaskAdvice { + public static class TaskArg2Advice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void wrapCallback( @@ -40,4 +45,14 @@ public static void wrapCallback( action = new TaskWrapper(Java8BytecodeBridge.currentContext(), action); } } + + @SuppressWarnings("unused") + public static class TaskArg3Advice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void wrapCallback( + @Advice.Argument(value = 3, 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/InternalStreamConnectionInstrumentation.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/InternalStreamConnectionInstrumentation.java index c22896ffcc44..3316353583df 100644 --- a/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/InternalStreamConnectionInstrumentation.java +++ b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/InternalStreamConnectionInstrumentation.java @@ -26,21 +26,42 @@ public ElementMatcher typeMatcher() { @Override public void transform(TypeTransformer transformer) { + // before 5.2.0 transformer.applyAdviceToMethod( isMethod() .and(named("openAsync")) .and(takesArgument(0, named("com.mongodb.internal.async.SingleResultCallback"))), this.getClass().getName() + "$SingleResultCallbackArg0Advice"); + // since 5.2.0 + transformer.applyAdviceToMethod( + isMethod() + .and(named("openAsync")) + .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))), + this.getClass().getName() + "$SingleResultCallbackArg1Advice"); + // before 5.2.0 transformer.applyAdviceToMethod( isMethod() .and(named("readAsync")) .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))), this.getClass().getName() + "$SingleResultCallbackArg1Advice"); + // since 5.2.0 + transformer.applyAdviceToMethod( + isMethod() + .and(named("readAsync")) + .and(takesArgument(2, named("com.mongodb.internal.async.SingleResultCallback"))), + this.getClass().getName() + "$SingleResultCallbackArg2Advice"); + // removed in 5.2.0 transformer.applyAdviceToMethod( isMethod() .and(named("writeAsync")) .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))), this.getClass().getName() + "$SingleResultCallbackArg1Advice"); + // since 5.2.0, earlier versions instrument writeAsync instead + transformer.applyAdviceToMethod( + isMethod() + .and(named("sendMessageAsync")) + .and(takesArgument(3, named("com.mongodb.internal.async.SingleResultCallback"))), + this.getClass().getName() + "$SingleResultCallbackArg3Advice"); } @SuppressWarnings("unused") @@ -62,4 +83,24 @@ 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); + } + } + + @SuppressWarnings("unused") + public static class SingleResultCallbackArg3Advice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void wrapCallback( + @Advice.Argument(value = 3, readOnly = false) SingleResultCallback callback) { + callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback); + } + } }