Fix mongo latest dep tests (#12331)

open-telemetry · Sep 25, 2024 · f1faa01 · f1faa01
1 parent e4752ab
commit f1faa01
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
diff --git a/...emetry/javaagent/instrumentation/mongo/v4_0/DefaultConnectionPoolTaskInstrumentation.java b/...emetry/javaagent/instrumentation/mongo/v4_0/DefaultConnectionPoolTaskInstrumentation.java
@@ -26,18 +26,33 @@ public ElementMatcher<TypeDescription> 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(
         @Advice.Argument(value = 2, readOnly = false) Consumer<Object> action) {
       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<Object> action) {
+      action = new TaskWrapper(Java8BytecodeBridge.currentContext(), action);
+    }
+  }
 }
diff --git a/...lemetry/javaagent/instrumentation/mongo/v4_0/InternalStreamConnectionInstrumentation.java b/...lemetry/javaagent/instrumentation/mongo/v4_0/InternalStreamConnectionInstrumentation.java
@@ -26,21 +26,42 @@ public ElementMatcher<TypeDescription> 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<Object> 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<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
 }