Fix mongo latest dep tests (#8785)

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<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
 }

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<TypeDescription> 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<Object> action) {
+      action = new TaskWrapper(Java8BytecodeBridge.currentContext(), action);
+    }
+  }
+}

instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java

@@ -33,6 +33,7 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new InternalStreamConnectionInstrumentation(),
         new BaseClusterInstrumentation(),
         new DefaultConnectionPoolInstrumentation(),
+        new DefaultConnectionPoolTaskInstrumentation(),
         new AsyncWorkManagerInstrumentation());
   }
 }

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<Object> {
+  private final Context context;
+  private final Consumer<Object> delegate;
+
+  public TaskWrapper(Context context, Consumer<Object> delegate) {
+    this.context = context;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void accept(Object value) {
+    try (Scope ignored = context.makeCurrent()) {
+      delegate.accept(value);
+    }
+  }
+}