Mongo reactive context propagation (#2758)

instrumentation/mongo/mongo-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_1/MongoClientInstrumentationModule.java

@@ -11,6 +11,7 @@ import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
@@ -19,8 +20,6 @@ import com.mongodb.event.CommandListener;
 import io.opentelemetry.javaagent.instrumentation.mongo.TracingCommandListener;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.lang.reflect.Modifier;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -48,14 +47,10 @@ public class MongoClientInstrumentationModule extends InstrumentationModule {
           .and(
               declaresMethod(
                   named("addCommandListener")
+                      .and(isPublic())
                       .and(
-                          takesArguments(
+                          takesArguments(1)
-                              new TypeDescription.Latent(
+                              .and(takesArgument(0, named("com.mongodb.event.CommandListener"))))));
-                                  "com.mongodb.event.CommandListener",
-                                  Modifier.PUBLIC,
-                                  null,
-                                  Collections.<TypeDescription.Generic>emptyList())))
-                      .and(isPublic())));
     }
 
     @Override

instrumentation/mongo/mongo-3.1/javaagent/src/test/groovy/MongoClientTest.groovy

@@ -17,7 +17,7 @@ import org.bson.BsonString
 import org.bson.Document
 import spock.lang.Shared
 
-class MongoClientTest extends AbstractMongoClientTest {
+class MongoClientTest extends AbstractMongoClientTest<MongoCollection<Document>> {
 
   @Shared
   MongoClient client
@@ -61,19 +61,24 @@ class MongoClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int insert(String dbName, String collectionName) {
+  MongoCollection<Document> setupInsert(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
       return db.getCollection(collectionName)
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int insert(MongoCollection<Document> collection) {
     collection.insertOne(new Document("password", "SECRET"))
     return collection.count()
   }
 
   @Override
-  int update(String dbName, String collectionName) {
+  MongoCollection<Document> setupUpdate(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
@@ -82,6 +87,11 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int update(MongoCollection<Document> collection) {
     def result = collection.updateOne(
       new BsonDocument("password", new BsonString("OLDPW")),
       new BsonDocument('$set', new BsonDocument("password", new BsonString("NEWPW"))))
@@ -90,7 +100,7 @@ class MongoClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int delete(String dbName, String collectionName) {
+  MongoCollection<Document> setupDelete(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
@@ -99,13 +109,18 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int delete(MongoCollection<Document> collection) {
     def result = collection.deleteOne(new BsonDocument("password", new BsonString("SECRET")))
     collection.count()
     return result.deletedCount
   }
 
   @Override
-  void getMore(String dbName, String collectionName) {
+  MongoCollection<Document> setupGetMore(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def coll = db.getCollection(collectionName)
@@ -113,6 +128,11 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  void getMore(MongoCollection<Document> collection) {
     collection.find().filter(new Document("_id", new Document('$gte', 0)))
       .batchSize(2).into(new ArrayList())
   }

instrumentation/mongo/mongo-3.7/javaagent/mongo-3.7-javaagent.gradle

@@ -4,7 +4,7 @@ muzzle {
   pass {
     group = "org.mongodb"
     module = "mongo-java-driver"
-    versions = "[3.7,)"
+    versions = "[3.7, 4.0)"
     assertInverse = true
   }
   pass {
@@ -13,8 +13,7 @@ muzzle {
     // this instrumentation is backwards compatible with early versions of the new API that shipped in 3.7
     // the legacy API instrumented in mongo-3.1 continues to be shipped in 4.x, but doesn't conflict here
     // because they are triggered by different types: MongoClientSettings(new) vs MongoClientOptions(legacy)
-    versions = "[3.7,)"
+    versions = "[3.7, 4.0)"
-    extraDependency "org.mongodb:bson"
     assertInverse = true
   }
 }

instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/MongoClientInstrumentationModule.java

@@ -5,22 +5,24 @@
 
 package io.opentelemetry.javaagent.instrumentation.mongo.v3_7;
 
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import com.mongodb.MongoClientSettings;
+import com.mongodb.async.SingleResultCallback;
 import com.mongodb.event.CommandListener;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.mongo.TracingCommandListener;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.lang.reflect.Modifier;
+import java.util.HashMap;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -37,7 +39,10 @@ public class MongoClientInstrumentationModule extends InstrumentationModule {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new MongoClientSettingsBuilderInstrumentation());
+    return asList(
+        new MongoClientSettingsBuilderInstrumentation(),
+        new InternalStreamConnectionInstrumentation(),
+        new BaseClusterInstrumentation());
   }
 
   private static final class MongoClientSettingsBuilderInstrumentation
@@ -48,14 +53,10 @@ public class MongoClientInstrumentationModule extends InstrumentationModule {
           .and(
               declaresMethod(
                   named("addCommandListener")
+                      .and(isPublic())
                       .and(
-                          takesArguments(
+                          takesArguments(1)
-                              new TypeDescription.Latent(
+                              .and(takesArgument(0, named("com.mongodb.event.CommandListener"))))));
-                                  "com.mongodb.event.CommandListener",
-                                  Modifier.PUBLIC,
-                                  null,
-                                  Collections.<TypeDescription.Generic>emptyList())))
-                      .and(isPublic())));
     }
 
     @Override
@@ -80,4 +81,72 @@ public class MongoClientInstrumentationModule extends InstrumentationModule {
       builder.addCommandListener(new TracingCommandListener());
     }
   }
+
+  private static final class InternalStreamConnectionInstrumentation
+      implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.internal.connection.InternalStreamConnection");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>();
+      transformers.put(
+          isMethod()
+              .and(named("openAsync"))
+              .and(takesArgument(0, named("com.mongodb.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg0Advice");
+      transformers.put(
+          isMethod()
+              .and(named("readAsync"))
+              .and(takesArgument(1, named("com.mongodb.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg1Advice");
+      transformers.put(
+          isMethod()
+              .and(named("writeAsync"))
+              .and(takesArgument(1, named("com.mongodb.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg1Advice");
+      return transformers;
+    }
+  }
+
+  private static final class BaseClusterInstrumentation implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.connection.BaseCluster")
+          .or(named("com.mongodb.internal.connection.BaseCluster"));
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      return singletonMap(
+          isMethod()
+              .and(isPublic())
+              .and(named("selectServerAsync"))
+              .and(takesArgument(0, named("com.mongodb.selector.ServerSelector")))
+              .and(takesArgument(1, named("com.mongodb.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg1Advice");
+    }
+  }
+
+  public static class SingleResultCallbackArg0Advice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrapCallback(
+        @Advice.Argument(value = 0, readOnly = false) SingleResultCallback<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
+
+  public static class SingleResultCallbackArg1Advice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrapCallback(
+        @Advice.Argument(value = 1, readOnly = false) SingleResultCallback<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
 }

instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/SingleResultCallbackWrapper.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.mongo.v3_7;
+
+import com.mongodb.async.SingleResultCallback;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+public class SingleResultCallbackWrapper implements SingleResultCallback<Object> {
+  private final Context context;
+  private final SingleResultCallback<Object> delegate;
+
+  public SingleResultCallbackWrapper(Context context, SingleResultCallback<Object> delegate) {
+    this.context = context;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void onResult(Object server, Throwable throwable) {
+    try (Scope ignored = context.makeCurrent()) {
+      delegate.onResult(server, throwable);
+    }
+  }
+}

instrumentation/mongo/mongo-3.7/javaagent/src/test/groovy/MongoClientTest.groovy

@@ -18,7 +18,7 @@ import org.bson.BsonString
 import org.bson.Document
 import spock.lang.Shared
 
-class MongoClientTest extends AbstractMongoClientTest {
+class MongoClientTest extends AbstractMongoClientTest<MongoCollection<Document>> {
 
   @Shared
   MongoClient client
@@ -71,19 +71,24 @@ class MongoClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int insert(String dbName, String collectionName) {
+  MongoCollection<Document> setupInsert(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
       return db.getCollection(collectionName)
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int insert(MongoCollection<Document> collection) {
     collection.insertOne(new Document("password", "SECRET"))
     return collection.count()
   }
 
   @Override
-  int update(String dbName, String collectionName) {
+  MongoCollection<Document> setupUpdate(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
@@ -92,6 +97,11 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int update(MongoCollection<Document> collection) {
     def result = collection.updateOne(
       new BsonDocument("password", new BsonString("OLDPW")),
       new BsonDocument('$set', new BsonDocument("password", new BsonString("NEWPW"))))
@@ -100,7 +110,7 @@ class MongoClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int delete(String dbName, String collectionName) {
+  MongoCollection<Document> setupDelete(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
@@ -109,13 +119,18 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int delete(MongoCollection<Document> collection) {
     def result = collection.deleteOne(new BsonDocument("password", new BsonString("SECRET")))
     collection.count()
     return result.deletedCount
   }
 
   @Override
-  void getMore(String dbName, String collectionName) {
+  MongoCollection<Document> setupGetMore(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def coll = db.getCollection(collectionName)
@@ -123,6 +138,11 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  void getMore(MongoCollection<Document> collection) {
     collection.find().filter(new Document("_id", new Document('$gte', 0)))
       .batchSize(2).into(new ArrayList())
   }

instrumentation/mongo/mongo-4.0-testing/mongo-4.0-testing.gradle

@@ -1,15 +0,0 @@
-ext {
-  skipPublish = true
-}
-apply from: "$rootDir/gradle/instrumentation.gradle"
-
-dependencies {
-  testInstrumentation(project(':instrumentation:mongo:mongo-3.7:javaagent')) {
-    exclude group: 'org.mongodb', module: 'mongo-java-driver'
-  }
-  testImplementation project(':instrumentation:mongo:mongo-testing')
-  testImplementation group: 'org.mongodb', name: 'mongodb-driver-core', version: '4.0.0'
-  testImplementation group: 'org.mongodb', name: 'mongodb-driver-sync', version: '4.0.0'
-  testImplementation group: 'org.mongodb', name: 'mongodb-driver-reactivestreams', version: '4.0.0'
-  testImplementation group: 'de.flapdoodle.embed', name: 'de.flapdoodle.embed.mongo', version: '1.50.5'
-}

instrumentation/mongo/mongo-4.0/javaagent/mongo-4.0-javaagent.gradle

@@ -0,0 +1,22 @@
+apply from: "$rootDir/gradle/instrumentation.gradle"
+
+muzzle {
+  pass {
+    group = "org.mongodb"
+    module = "mongodb-driver-core"
+    versions = "[4.0,)"
+    assertInverse = true
+  }
+}
+
+dependencies {
+  implementation(project(':instrumentation:mongo:mongo-common:javaagent'))
+
+  library group: 'org.mongodb', name: 'mongodb-driver-core', version: '4.0.0'
+
+  testLibrary group: 'org.mongodb', name: 'mongodb-driver-sync', version: '4.0.0'
+  testLibrary group: 'org.mongodb', name: 'mongodb-driver-reactivestreams', version: '4.0.0'
+
+  testImplementation project(':instrumentation:mongo:mongo-testing')
+  testImplementation group: 'de.flapdoodle.embed', name: 'de.flapdoodle.embed.mongo', version: '1.50.5'
+}

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

@@ -0,0 +1,151 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.mongo.v4_0;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import com.mongodb.MongoClientSettings;
+import com.mongodb.event.CommandListener;
+import com.mongodb.internal.async.SingleResultCallback;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.mongo.TracingCommandListener;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+public class MongoClientInstrumentationModule extends InstrumentationModule {
+
+  public MongoClientInstrumentationModule() {
+    super("mongo", "mongo-4.0");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return asList(
+        new MongoClientSettingsBuilderInstrumentation(),
+        new InternalStreamConnectionInstrumentation(),
+        new BaseClusterInstrumentation());
+  }
+
+  private static final class MongoClientSettingsBuilderInstrumentation
+      implements TypeInstrumentation {
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.MongoClientSettings$Builder")
+          .and(
+              declaresMethod(
+                  named("addCommandListener")
+                      .and(isPublic())
+                      .and(
+                          takesArguments(1)
+                              .and(takesArgument(0, named("com.mongodb.event.CommandListener"))))));
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      return singletonMap(
+          isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)),
+          MongoClientInstrumentationModule.class.getName() + "$MongoClientAdvice");
+    }
+  }
+
+  public static class MongoClientAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void injectTraceListener(
+        @Advice.This MongoClientSettings.Builder builder,
+        @Advice.FieldValue("commandListeners") List<CommandListener> commandListeners) {
+      for (CommandListener commandListener : commandListeners) {
+        if (commandListener instanceof TracingCommandListener) {
+          return;
+        }
+      }
+      builder.addCommandListener(new TracingCommandListener());
+    }
+  }
+
+  private static final class InternalStreamConnectionInstrumentation
+      implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.internal.connection.InternalStreamConnection");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>();
+      transformers.put(
+          isMethod()
+              .and(named("openAsync"))
+              .and(takesArgument(0, named("com.mongodb.internal.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg0Advice");
+      transformers.put(
+          isMethod()
+              .and(named("readAsync"))
+              .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg1Advice");
+      transformers.put(
+          isMethod()
+              .and(named("writeAsync"))
+              .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg1Advice");
+      return transformers;
+    }
+  }
+
+  private static final class BaseClusterInstrumentation implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.internal.connection.BaseCluster");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      return singletonMap(
+          isMethod()
+              .and(isPublic())
+              .and(named("selectServerAsync"))
+              .and(takesArgument(0, named("com.mongodb.selector.ServerSelector")))
+              .and(takesArgument(1, named("com.mongodb.internal.async.SingleResultCallback"))),
+          MongoClientInstrumentationModule.class.getName() + "$SingleResultCallbackArg1Advice");
+    }
+  }
+
+  public static class SingleResultCallbackArg0Advice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrapCallback(
+        @Advice.Argument(value = 0, readOnly = false) SingleResultCallback<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
+
+  public static class SingleResultCallbackArg1Advice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrapCallback(
+        @Advice.Argument(value = 1, 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/SingleResultCallbackWrapper.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.mongo.v4_0;
+
+import com.mongodb.internal.async.SingleResultCallback;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+public class SingleResultCallbackWrapper implements SingleResultCallback<Object> {
+  private final Context context;
+  private final SingleResultCallback<Object> delegate;
+
+  public SingleResultCallbackWrapper(Context context, SingleResultCallback<Object> delegate) {
+    this.context = context;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void onResult(Object server, Throwable throwable) {
+    try (Scope ignored = context.makeCurrent()) {
+      delegate.onResult(server, throwable);
+    }
+  }
+}

instrumentation/mongo/mongo-4.0/javaagent/src/test/groovy/Mongo4ReactiveClientTest.groovy

@@ -21,7 +21,7 @@ import org.reactivestreams.Subscriber
 import org.reactivestreams.Subscription
 import spock.lang.Shared
 
-class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
+class Mongo4ReactiveClientTest extends AbstractMongoClientTest<MongoCollection<Document>> {
 
   @Shared
   MongoClient client
@@ -61,16 +61,20 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int insert(String dbName, String collectionName) {
+  MongoCollection<Document> setupInsert(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def latch1 = new CountDownLatch(1)
-      // This creates a trace that isn't linked to the parent... using NIO internally that we don't handle.
       db.createCollection(collectionName).subscribe(toSubscriber { latch1.countDown() })
       latch1.await()
       return db.getCollection(collectionName)
     }
-    ignoreTracesAndClear(2)
+    ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int insert(MongoCollection<Document> collection) {
     def count = new CompletableFuture<Integer>()
     collection.insertOne(new Document("password", "SECRET")).subscribe(toSubscriber {
       collection.estimatedDocumentCount().subscribe(toSubscriber { count.complete(it) })
@@ -79,7 +83,7 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int update(String dbName, String collectionName) {
+  MongoCollection<Document> setupUpdate(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def latch1 = new CountDownLatch(1)
@@ -92,6 +96,11 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int update(MongoCollection<Document> collection) {
     def result = new CompletableFuture<UpdateResult>()
     def count = new CompletableFuture()
     collection.updateOne(
@@ -104,7 +113,7 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int delete(String dbName, String collectionName) {
+  MongoCollection<Document> setupDelete(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def latch1 = new CountDownLatch(1)
@@ -117,6 +126,11 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int delete(MongoCollection<Document> collection) {
     def result = new CompletableFuture<DeleteResult>()
     def count = new CompletableFuture()
     collection.deleteOne(new BsonDocument("password", new BsonString("SECRET"))).subscribe(toSubscriber {
@@ -127,7 +141,12 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  void getMore(String dbName, String collectionName) {
+  MongoCollection<Document> setupGetMore(String dbName, String collectionName) {
+    throw new AssumptionViolatedException("not tested on reactive")
+  }
+
+  @Override
+  void getMore(MongoCollection<Document> collection) {
     throw new AssumptionViolatedException("not tested on reactive")
   }
 

instrumentation/mongo/mongo-4.0/javaagent/src/test/groovy/MongoClientTest.groovy

@@ -15,7 +15,7 @@ import org.bson.Document
 import org.junit.AssumptionViolatedException
 import spock.lang.Shared
 
-class MongoClientTest extends AbstractMongoClientTest {
+class MongoClientTest extends AbstractMongoClientTest<MongoCollection<Document>> {
 
   @Shared
   MongoClient client
@@ -53,19 +53,24 @@ class MongoClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int insert(String dbName, String collectionName) {
+  MongoCollection<Document> setupInsert(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
       return db.getCollection(collectionName)
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int insert(MongoCollection<Document> collection) {
     collection.insertOne(new Document("password", "SECRET"))
     return collection.estimatedDocumentCount()
   }
 
   @Override
-  int update(String dbName, String collectionName) {
+  MongoCollection<Document> setupUpdate(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
@@ -74,6 +79,11 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int update(MongoCollection<Document> collection) {
     def result = collection.updateOne(
       new BsonDocument("password", new BsonString("OLDPW")),
       new BsonDocument('$set', new BsonDocument("password", new BsonString("NEWPW"))))
@@ -82,7 +92,7 @@ class MongoClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int delete(String dbName, String collectionName) {
+  MongoCollection<Document> setupDelete(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       db.createCollection(collectionName)
@@ -91,13 +101,18 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int delete(MongoCollection<Document> collection) {
     def result = collection.deleteOne(new BsonDocument("password", new BsonString("SECRET")))
     collection.estimatedDocumentCount()
     return result.deletedCount
   }
 
   @Override
-  void getMore(String dbName, String collectionName) {
+  MongoCollection<Document> setupGetMore(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def coll = db.getCollection(collectionName)
@@ -105,6 +120,11 @@ class MongoClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  void getMore(MongoCollection<Document> collection) {
     collection.find().filter(new Document("_id", new Document('$gte', 0)))
       .batchSize(2).into(new ArrayList())
   }

instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/MongoAsyncClientInstrumentationModule.java

@@ -5,23 +5,25 @@
 
 package io.opentelemetry.javaagent.instrumentation.mongoasync.v3_3;
 
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.declaresField;
 import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
+import com.mongodb.async.SingleResultCallback;
 import com.mongodb.async.client.MongoClientSettings;
 import com.mongodb.event.CommandListener;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.mongo.TracingCommandListener;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.lang.reflect.Modifier;
+import java.util.HashMap;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -38,7 +40,10 @@ public class MongoAsyncClientInstrumentationModule extends InstrumentationModule
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new MongoClientSettingsBuildersInstrumentation());
+    return asList(
+        new MongoClientSettingsBuildersInstrumentation(),
+        new InternalStreamConnectionInstrumentation(),
+        new BaseClusterInstrumentation());
   }
 
   private static final class MongoClientSettingsBuildersInstrumentation
@@ -49,14 +54,10 @@ public class MongoAsyncClientInstrumentationModule extends InstrumentationModule
           .and(
               declaresMethod(
                   named("addCommandListener")
+                      .and(isPublic())
                       .and(
-                          takesArguments(
+                          takesArguments(1)
-                              new TypeDescription.Latent(
+                              .and(takesArgument(0, named("com.mongodb.event.CommandListener"))))))
-                                  "com.mongodb.event.CommandListener",
-                                  Modifier.PUBLIC,
-                                  null,
-                                  Collections.<TypeDescription.Generic>emptyList())))
-                      .and(isPublic())))
           .and(declaresField(named("commandListeners")));
     }
 
@@ -82,4 +83,84 @@ public class MongoAsyncClientInstrumentationModule extends InstrumentationModule
       builder.addCommandListener(new TracingCommandListener());
     }
   }
+
+  private static final class InternalStreamConnectionInstrumentation
+      implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.connection.InternalStreamConnection");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>();
+      transformers.put(
+          isMethod()
+              .and(named("openAsync"))
+              .and(takesArgument(0, named("com.mongodb.async.SingleResultCallback"))),
+          MongoAsyncClientInstrumentationModule.class.getName()
+              + "$SingleResultCallbackArg0Advice");
+      transformers.put(
+          isMethod()
+              .and(named("readAsync"))
+              .and(takesArgument(1, named("com.mongodb.async.SingleResultCallback"))),
+          MongoAsyncClientInstrumentationModule.class.getName()
+              + "$SingleResultCallbackArg1Advice");
+      transformers.put(
+          isMethod()
+              .and(named("sendMessageAsync"))
+              .and(takesArgument(2, named("com.mongodb.async.SingleResultCallback"))),
+          MongoAsyncClientInstrumentationModule.class.getName()
+              + "$SingleResultCallbackArg2Advice");
+      return transformers;
+    }
+  }
+
+  private static final class BaseClusterInstrumentation implements TypeInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.mongodb.connection.BaseCluster");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      return singletonMap(
+          isMethod()
+              .and(isPublic())
+              .and(named("selectServerAsync"))
+              .and(takesArgument(0, named("com.mongodb.selector.ServerSelector")))
+              .and(takesArgument(1, named("com.mongodb.async.SingleResultCallback"))),
+          MongoAsyncClientInstrumentationModule.class.getName()
+              + "$SingleResultCallbackArg1Advice");
+    }
+  }
+
+  public static class SingleResultCallbackArg0Advice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrapCallback(
+        @Advice.Argument(value = 0, readOnly = false) SingleResultCallback<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
+
+  public static class SingleResultCallbackArg1Advice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void wrapCallback(
+        @Advice.Argument(value = 1, readOnly = false) SingleResultCallback<Object> callback) {
+      callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback);
+    }
+  }
+
+  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-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/SingleResultCallbackWrapper.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.mongoasync.v3_3;
+
+import com.mongodb.async.SingleResultCallback;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+public class SingleResultCallbackWrapper implements SingleResultCallback<Object> {
+  private final Context context;
+  private final SingleResultCallback<Object> delegate;
+
+  public SingleResultCallbackWrapper(Context context, SingleResultCallback<Object> delegate) {
+    this.context = context;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void onResult(Object server, Throwable throwable) {
+    try (Scope ignored = context.makeCurrent()) {
+      delegate.onResult(server, throwable);
+    }
+  }
+}

instrumentation/mongo/mongo-async-3.3/javaagent/src/test/groovy/MongoAsyncClientTest.groovy

@@ -23,7 +23,7 @@ import org.bson.Document
 import org.junit.AssumptionViolatedException
 import spock.lang.Shared
 
-class MongoAsyncClientTest extends AbstractMongoClientTest {
+class MongoAsyncClientTest extends AbstractMongoClientTest<MongoCollection<Document>> {
 
   @Shared
   MongoClient client
@@ -73,7 +73,7 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int insert(String dbName, String collectionName) {
+  MongoCollection<Document> setupInsert(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def latch1 = new CountDownLatch(1)
@@ -82,6 +82,11 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
       return db.getCollection(collectionName)
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int insert(MongoCollection<Document> collection) {
     def count = new CompletableFuture<Integer>()
     collection.insertOne(new Document("password", "SECRET"), toCallback {
       collection.count toCallback { count.complete(it) }
@@ -90,7 +95,7 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int update(String dbName, String collectionName) {
+  MongoCollection<Document> setupUpdate(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def latch1 = new CountDownLatch(1)
@@ -103,6 +108,11 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int update(MongoCollection<Document> collection) {
     def result = new CompletableFuture<UpdateResult>()
     def count = new CompletableFuture()
     collection.updateOne(
@@ -115,7 +125,7 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  int delete(String dbName, String collectionName) {
+  MongoCollection<Document> setupDelete(String dbName, String collectionName) {
     MongoCollection<Document> collection = runUnderTrace("setup") {
       MongoDatabase db = client.getDatabase(dbName)
       def latch1 = new CountDownLatch(1)
@@ -128,6 +138,11 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
       return coll
     }
     ignoreTracesAndClear(1)
+    return collection
+  }
+
+  @Override
+  int delete(MongoCollection<Document> collection) {
     def result = new CompletableFuture<DeleteResult>()
     def count = new CompletableFuture()
     collection.deleteOne(new BsonDocument("password", new BsonString("SECRET")), toCallback {
@@ -138,7 +153,12 @@ class MongoAsyncClientTest extends AbstractMongoClientTest {
   }
 
   @Override
-  void getMore(String dbName, String collectionName) {
+  MongoCollection<Document> setupGetMore(String dbName, String collectionName) {
+    throw new AssumptionViolatedException("not tested on async")
+  }
+
+  @Override
+  void getMore(MongoCollection<Document> collection) {
     throw new AssumptionViolatedException("not tested on async")
   }
 

instrumentation/mongo/mongo-testing/src/main/groovy/AbstractMongoClientTest.groovy

@@ -4,6 +4,8 @@
  */
 
 import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.instrumentation.test.utils.TraceUtils.basicSpan
+import static io.opentelemetry.instrumentation.test.utils.TraceUtils.runUnderTrace
 
 import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.instrumentation.test.asserts.TraceAssert
@@ -15,7 +17,7 @@ import org.testcontainers.containers.GenericContainer
 import org.testcontainers.containers.output.Slf4jLogConsumer
 import spock.lang.Shared
 
-abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification {
+abstract class AbstractMongoClientTest<T> extends AgentInstrumentationSpecification {
 
   @Shared
   GenericContainer mongodb
@@ -50,13 +52,17 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   abstract int getCollection(String dbName, String collectionName)
 
-  abstract int insert(String dbName, String collectionName)
+  abstract T setupInsert(String dbName, String collectionName)
+  abstract int insert(T collection)
 
-  abstract int update(String dbName, String collectionName)
+  abstract T setupUpdate(String dbName, String collectionName)
+  abstract int update(T collection)
 
-  abstract int delete(String dbName, String collectionName)
+  abstract T setupDelete(String dbName, String collectionName)
+  abstract int delete(T collection)
 
-  abstract void getMore(String dbName, String collectionName)
+  abstract T setupGetMore(String dbName, String collectionName)
+  abstract void getMore(T collection)
 
   abstract void error(String dbName, String collectionName)
 
@@ -70,12 +76,15 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test create collection"() {
     when:
-    createCollection(dbName, collectionName)
+    runUnderTrace("parent") {
+      createCollection(dbName, collectionName)
+    }
 
     then:
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 2) {
-        mongoSpan(it, 0, "create", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "create", collectionName, dbName, span(0)) {
           assert it == "{\"create\":\"$collectionName\",\"capped\":\"?\"}" ||
             it == "{\"create\": \"$collectionName\", \"capped\": \"?\", \"\$db\": \"?\", \"\$readPreference\": {\"mode\": \"?\"}}"
           true
@@ -90,12 +99,15 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test create collection no description"() {
     when:
-    createCollectionNoDescription(dbName, collectionName)
+    runUnderTrace("parent") {
+      createCollectionNoDescription(dbName, collectionName)
+    }
 
     then:
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 2) {
-        mongoSpan(it, 0, "create", collectionName, dbName, {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "create", collectionName, dbName, span(0), {
           assert it == "{\"create\":\"$collectionName\",\"capped\":\"?\"}" ||
             it == "{\"create\": \"$collectionName\", \"capped\": \"?\", \"\$db\": \"?\", \"\$readPreference\": {\"mode\": \"?\"}}"
           true
@@ -110,13 +122,16 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test get collection"() {
     when:
-    def count = getCollection(dbName, collectionName)
+    def count = runUnderTrace("parent") {
+      getCollection(dbName, collectionName)
+    }
 
     then:
     count == 0
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 2) {
-        mongoSpan(it, 0, "count", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "count", collectionName, dbName, span(0)) {
           assert it == "{\"count\":\"$collectionName\",\"query\":{}}" ||
             it == "{\"count\": \"$collectionName\", \"query\": {}, \"\$db\": \"?\", \"\$readPreference\": {\"mode\": \"?\"}}"
           true
@@ -131,20 +146,22 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test insert"() {
     when:
-    def count = insert(dbName, collectionName)
+    def collection = setupInsert(dbName, collectionName)
+    def count = runUnderTrace("parent") {
+      insert(collection)
+    }
 
     then:
     count == 1
-    assertTraces(2) {
+    assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
-        mongoSpan(it, 0, "insert", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "insert", collectionName, dbName, span(0)) {
           assert it == "{\"insert\":\"$collectionName\",\"ordered\":\"?\",\"documents\":[{\"_id\":\"?\",\"password\":\"?\"}]}" ||
             it == "{\"insert\": \"$collectionName\", \"ordered\": \"?\", \"\$db\": \"?\", \"documents\": [{\"_id\": \"?\", \"password\": \"?\"}]}"
           true
         }
-      }
+        mongoSpan(it, 2, "count", collectionName, dbName, span(0)) {
-      trace(1, 1) {
-        mongoSpan(it, 0, "count", collectionName, dbName) {
           assert it == "{\"count\":\"$collectionName\",\"query\":{}}" ||
             it == "{\"count\": \"$collectionName\", \"query\": {}, \"\$db\": \"?\", \"\$readPreference\": {\"mode\": \"?\"}}"
           true
@@ -159,20 +176,22 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test update"() {
     when:
-    int modifiedCount = update(dbName, collectionName)
+    def collection = setupUpdate(dbName, collectionName)
+    int modifiedCount = runUnderTrace("parent") {
+      update(collection)
+    }
 
     then:
     modifiedCount == 1
-    assertTraces(2) {
+    assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
-        mongoSpan(it, 0, "update", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "update", collectionName, dbName, span(0)) {
           assert it == "{\"update\":\"$collectionName\",\"ordered\":\"?\",\"updates\":[{\"q\":{\"password\":\"?\"},\"u\":{\"\$set\":{\"password\":\"?\"}}}]}" ||
             it == "{\"update\": \"?\", \"ordered\": \"?\", \"\$db\": \"?\", \"updates\": [{\"q\": {\"password\": \"?\"}, \"u\": {\"\$set\": {\"password\": \"?\"}}}]}"
           true
         }
-      }
+        mongoSpan(it, 2, "count", collectionName, dbName, span(0)) {
-      trace(1, 1) {
-        mongoSpan(it, 0, "count", collectionName, dbName) {
           assert it == "{\"count\":\"$collectionName\",\"query\":{}}" ||
             it == "{\"count\": \"$collectionName\", \"query\": {}, \"\$db\": \"?\", \"\$readPreference\": {\"mode\": \"?\"}}"
           true
@@ -187,20 +206,22 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test delete"() {
     when:
-    int deletedCount = delete(dbName, collectionName)
+    def collection = setupDelete(dbName, collectionName)
+    int deletedCount = runUnderTrace("parent") {
+      delete(collection)
+    }
 
     then:
     deletedCount == 1
-    assertTraces(2) {
+    assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
-        mongoSpan(it, 0, "delete", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "delete", collectionName, dbName, span(0)) {
           assert it == "{\"delete\":\"$collectionName\",\"ordered\":\"?\",\"deletes\":[{\"q\":{\"password\":\"?\"},\"limit\":\"?\"}]}" ||
             it == "{\"delete\": \"?\", \"ordered\": \"?\", \"\$db\": \"?\", \"deletes\": [{\"q\": {\"password\": \"?\"}, \"limit\": \"?\"}]}"
           true
         }
-      }
+        mongoSpan(it, 2, "count", collectionName, dbName, span(0)) {
-      trace(1, 1) {
-        mongoSpan(it, 0, "count", collectionName, dbName) {
           assert it == "{\"count\":\"$collectionName\",\"query\":{}}" ||
             it == "{\"count\": \"$collectionName\", \"query\": {}, \"\$db\": \"?\", \"\$readPreference\": {\"mode\": \"?\"}}"
           true
@@ -215,18 +236,20 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test collection name for getMore command"() {
     when:
-    getMore(dbName, collectionName)
+    def collection = setupGetMore(dbName, collectionName)
+    runUnderTrace("parent") {
+      getMore(collection)
+    }
 
     then:
-    assertTraces(2) {
+    assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
-        mongoSpan(it, 0, "find", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "find", collectionName, dbName, span(0)) {
           assert it == '{"find":"' + collectionName + '","filter":{"_id":{"$gte":"?"}},"batchSize":"?"}'
           true
         }
-      }
+        mongoSpan(it, 2, "getMore", collectionName, dbName, span(0)) {
-      trace(1, 1) {
-        mongoSpan(it, 0, "getMore", collectionName, dbName) {
           assert it == '{"getMore":"?","collection":"?","batchSize":"?"}'
           true
         }
@@ -254,12 +277,15 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def "test create collection with already built ClientOptions"() {
     when:
-    createCollectionWithAlreadyBuiltClientOptions(dbName, collectionName)
+    runUnderTrace("parent") {
+      createCollectionWithAlreadyBuiltClientOptions(dbName, collectionName)
+    }
 
     then:
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 2) {
-        mongoSpan(it, 0, "create", collectionName, dbName) {
+        basicSpan(it, 0, "parent")
+        mongoSpan(it, 1, "create", collectionName, dbName, span(0)) {
           assert it == "{\"create\":\"$collectionName\",\"capped\":\"?\"}"
           true
         }
@@ -279,8 +305,8 @@ abstract class AbstractMongoClientTest extends AgentInstrumentationSpecification
 
   def mongoSpan(TraceAssert trace, int index,
                 String operation, String collection,
-                String dbName, Closure<Boolean> statementEval,
+                String dbName, Object parentSpan,
-                Object parentSpan = null, Throwable exception = null) {
+                Closure<Boolean> statementEval, Throwable exception = null) {
     trace.span(index) {
       name { operation + " " + dbName + "." + collection }
       kind CLIENT

settings.gradle

@@ -185,7 +185,7 @@ include ':instrumentation:logback:logback-1.0:testing'
 include ':instrumentation:methods:javaagent'
 include ':instrumentation:mongo:mongo-3.1:javaagent'
 include ':instrumentation:mongo:mongo-3.7:javaagent'
-include ':instrumentation:mongo:mongo-4.0-testing'
+include ':instrumentation:mongo:mongo-4.0:javaagent'
 include ':instrumentation:mongo:mongo-async-3.3:javaagent'
 include ':instrumentation:mongo:mongo-common:javaagent'
 include ':instrumentation:mongo:mongo-testing'