From 718a5a17a3bc2efa54ae514a1874a8c3e71084bc Mon Sep 17 00:00:00 2001 From: Lauri Tulmin Date: Tue, 13 Apr 2021 01:58:25 +0300 Subject: [PATCH] Mongo reactive context propagation (#2758) --- .../MongoClientInstrumentationModule.java | 13 +- .../src/test/groovy/MongoClientTest.groovy | 30 +++- .../javaagent/mongo-3.7-javaagent.gradle | 5 +- .../MongoClientInstrumentationModule.java | 91 +++++++++-- .../v3_7/SingleResultCallbackWrapper.java | 27 ++++ .../src/test/groovy/MongoClientTest.groovy | 30 +++- .../mongo-4.0-testing.gradle | 15 -- .../javaagent/mongo-4.0-javaagent.gradle | 22 +++ .../MongoClientInstrumentationModule.java | 151 ++++++++++++++++++ .../v4_0/SingleResultCallbackWrapper.java | 27 ++++ .../groovy/Mongo4ReactiveClientTest.groovy | 33 +++- .../src/test/groovy/MongoClientTest.groovy | 30 +++- ...MongoAsyncClientInstrumentationModule.java | 103 ++++++++++-- .../v3_3/SingleResultCallbackWrapper.java | 27 ++++ .../test/groovy/MongoAsyncClientTest.groovy | 30 +++- .../groovy/AbstractMongoClientTest.groovy | 120 ++++++++------ settings.gradle | 2 +- 17 files changed, 632 insertions(+), 124 deletions(-) create mode 100644 instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/SingleResultCallbackWrapper.java delete mode 100644 instrumentation/mongo/mongo-4.0-testing/mongo-4.0-testing.gradle create mode 100644 instrumentation/mongo/mongo-4.0/javaagent/mongo-4.0-javaagent.gradle create mode 100644 instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/MongoClientInstrumentationModule.java create mode 100644 instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/SingleResultCallbackWrapper.java rename instrumentation/mongo/{mongo-4.0-testing => mongo-4.0/javaagent}/src/test/groovy/Mongo4ReactiveClientTest.groovy (87%) rename instrumentation/mongo/{mongo-4.0-testing => mongo-4.0/javaagent}/src/test/groovy/MongoClientTest.groovy (82%) create mode 100644 instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/SingleResultCallbackWrapper.java diff --git a/instrumentation/mongo/mongo-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_1/MongoClientInstrumentationModule.java b/instrumentation/mongo/mongo-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_1/MongoClientInstrumentationModule.java index 9eb5b6d731..00f0f255d7 100644 --- a/instrumentation/mongo/mongo-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_1/MongoClientInstrumentationModule.java +++ b/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( - new TypeDescription.Latent( - "com.mongodb.event.CommandListener", - Modifier.PUBLIC, - null, - Collections.emptyList()))) - .and(isPublic()))); + takesArguments(1) + .and(takesArgument(0, named("com.mongodb.event.CommandListener")))))); } @Override diff --git a/instrumentation/mongo/mongo-3.1/javaagent/src/test/groovy/MongoClientTest.groovy b/instrumentation/mongo/mongo-3.1/javaagent/src/test/groovy/MongoClientTest.groovy index 956f0b7616..8d1bd6d876 100644 --- a/instrumentation/mongo/mongo-3.1/javaagent/src/test/groovy/MongoClientTest.groovy +++ b/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> { @Shared MongoClient client @@ -61,19 +61,24 @@ class MongoClientTest extends AbstractMongoClientTest { } @Override - int insert(String dbName, String collectionName) { + MongoCollection setupInsert(String dbName, String collectionName) { MongoCollection collection = runUnderTrace("setup") { MongoDatabase db = client.getDatabase(dbName) db.createCollection(collectionName) return db.getCollection(collectionName) } ignoreTracesAndClear(1) + return collection + } + + @Override + int insert(MongoCollection collection) { collection.insertOne(new Document("password", "SECRET")) return collection.count() } @Override - int update(String dbName, String collectionName) { + MongoCollection setupUpdate(String dbName, String collectionName) { MongoCollection 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 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 setupDelete(String dbName, String collectionName) { MongoCollection 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 collection) { def result = collection.deleteOne(new BsonDocument("password", new BsonString("SECRET"))) collection.count() return result.deletedCount } @Override - void getMore(String dbName, String collectionName) { + MongoCollection setupGetMore(String dbName, String collectionName) { MongoCollection 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 collection) { collection.find().filter(new Document("_id", new Document('$gte', 0))) .batchSize(2).into(new ArrayList()) } diff --git a/instrumentation/mongo/mongo-3.7/javaagent/mongo-3.7-javaagent.gradle b/instrumentation/mongo/mongo-3.7/javaagent/mongo-3.7-javaagent.gradle index 2bf3b72f8d..7ac2eab4f1 100644 --- a/instrumentation/mongo/mongo-3.7/javaagent/mongo-3.7-javaagent.gradle +++ b/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,)" - extraDependency "org.mongodb:bson" + versions = "[3.7, 4.0)" assertInverse = true } } diff --git a/instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/MongoClientInstrumentationModule.java b/instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/MongoClientInstrumentationModule.java index 656a523a77..171748e07e 100644 --- a/instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/MongoClientInstrumentationModule.java +++ b/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.Collections; +import java.util.HashMap; 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 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( - new TypeDescription.Latent( - "com.mongodb.event.CommandListener", - Modifier.PUBLIC, - null, - Collections.emptyList()))) - .and(isPublic()))); + takesArguments(1) + .and(takesArgument(0, named("com.mongodb.event.CommandListener")))))); } @Override @@ -80,4 +81,72 @@ public class MongoClientInstrumentationModule extends InstrumentationModule { builder.addCommandListener(new TracingCommandListener()); } } + + private static final class InternalStreamConnectionInstrumentation + implements TypeInstrumentation { + + @Override + public ElementMatcher typeMatcher() { + return named("com.mongodb.internal.connection.InternalStreamConnection"); + } + + @Override + public Map, String> transformers() { + Map, 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 typeMatcher() { + return named("com.mongodb.connection.BaseCluster") + .or(named("com.mongodb.internal.connection.BaseCluster")); + } + + @Override + public Map, 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 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 callback) { + callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback); + } + } } diff --git a/instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/SingleResultCallbackWrapper.java b/instrumentation/mongo/mongo-3.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v3_7/SingleResultCallbackWrapper.java new file mode 100644 index 0000000000..b5fbf7dc82 --- /dev/null +++ b/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 { + private final Context context; + private final SingleResultCallback delegate; + + public SingleResultCallbackWrapper(Context context, SingleResultCallback delegate) { + this.context = context; + this.delegate = delegate; + } + + @Override + public void onResult(Object server, Throwable throwable) { + try (Scope ignored = context.makeCurrent()) { + delegate.onResult(server, throwable); + } + } +} diff --git a/instrumentation/mongo/mongo-3.7/javaagent/src/test/groovy/MongoClientTest.groovy b/instrumentation/mongo/mongo-3.7/javaagent/src/test/groovy/MongoClientTest.groovy index d8d0c61378..89779e68d5 100644 --- a/instrumentation/mongo/mongo-3.7/javaagent/src/test/groovy/MongoClientTest.groovy +++ b/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> { @Shared MongoClient client @@ -71,19 +71,24 @@ class MongoClientTest extends AbstractMongoClientTest { } @Override - int insert(String dbName, String collectionName) { + MongoCollection setupInsert(String dbName, String collectionName) { MongoCollection collection = runUnderTrace("setup") { MongoDatabase db = client.getDatabase(dbName) db.createCollection(collectionName) return db.getCollection(collectionName) } ignoreTracesAndClear(1) + return collection + } + + @Override + int insert(MongoCollection collection) { collection.insertOne(new Document("password", "SECRET")) return collection.count() } @Override - int update(String dbName, String collectionName) { + MongoCollection setupUpdate(String dbName, String collectionName) { MongoCollection 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 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 setupDelete(String dbName, String collectionName) { MongoCollection 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 collection) { def result = collection.deleteOne(new BsonDocument("password", new BsonString("SECRET"))) collection.count() return result.deletedCount } @Override - void getMore(String dbName, String collectionName) { + MongoCollection setupGetMore(String dbName, String collectionName) { MongoCollection 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 collection) { collection.find().filter(new Document("_id", new Document('$gte', 0))) .batchSize(2).into(new ArrayList()) } diff --git a/instrumentation/mongo/mongo-4.0-testing/mongo-4.0-testing.gradle b/instrumentation/mongo/mongo-4.0-testing/mongo-4.0-testing.gradle deleted file mode 100644 index d5f197d996..0000000000 --- a/instrumentation/mongo/mongo-4.0-testing/mongo-4.0-testing.gradle +++ /dev/null @@ -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' -} diff --git a/instrumentation/mongo/mongo-4.0/javaagent/mongo-4.0-javaagent.gradle b/instrumentation/mongo/mongo-4.0/javaagent/mongo-4.0-javaagent.gradle new file mode 100644 index 0000000000..a45bf03ed0 --- /dev/null +++ b/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' +} 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 new file mode 100644 index 0000000000..9a5ae650d0 --- /dev/null +++ b/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 typeInstrumentations() { + return asList( + new MongoClientSettingsBuilderInstrumentation(), + new InternalStreamConnectionInstrumentation(), + new BaseClusterInstrumentation()); + } + + private static final class MongoClientSettingsBuilderInstrumentation + implements TypeInstrumentation { + @Override + public ElementMatcher 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, 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 commandListeners) { + for (CommandListener commandListener : commandListeners) { + if (commandListener instanceof TracingCommandListener) { + return; + } + } + builder.addCommandListener(new TracingCommandListener()); + } + } + + private static final class InternalStreamConnectionInstrumentation + implements TypeInstrumentation { + + @Override + public ElementMatcher typeMatcher() { + return named("com.mongodb.internal.connection.InternalStreamConnection"); + } + + @Override + public Map, String> transformers() { + Map, 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 typeMatcher() { + return named("com.mongodb.internal.connection.BaseCluster"); + } + + @Override + public Map, 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 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 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/SingleResultCallbackWrapper.java b/instrumentation/mongo/mongo-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongo/v4_0/SingleResultCallbackWrapper.java new file mode 100644 index 0000000000..730ffe4c60 --- /dev/null +++ b/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 { + private final Context context; + private final SingleResultCallback delegate; + + public SingleResultCallbackWrapper(Context context, SingleResultCallback delegate) { + this.context = context; + this.delegate = delegate; + } + + @Override + public void onResult(Object server, Throwable throwable) { + try (Scope ignored = context.makeCurrent()) { + delegate.onResult(server, throwable); + } + } +} diff --git a/instrumentation/mongo/mongo-4.0-testing/src/test/groovy/Mongo4ReactiveClientTest.groovy b/instrumentation/mongo/mongo-4.0/javaagent/src/test/groovy/Mongo4ReactiveClientTest.groovy similarity index 87% rename from instrumentation/mongo/mongo-4.0-testing/src/test/groovy/Mongo4ReactiveClientTest.groovy rename to instrumentation/mongo/mongo-4.0/javaagent/src/test/groovy/Mongo4ReactiveClientTest.groovy index de53e2d945..7c09fc462b 100644 --- a/instrumentation/mongo/mongo-4.0-testing/src/test/groovy/Mongo4ReactiveClientTest.groovy +++ b/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> { @Shared MongoClient client @@ -61,16 +61,20 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest { } @Override - int insert(String dbName, String collectionName) { + MongoCollection setupInsert(String dbName, String collectionName) { MongoCollection 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 collection) { def count = new CompletableFuture() 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 setupUpdate(String dbName, String collectionName) { MongoCollection 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 collection) { def result = new CompletableFuture() def count = new CompletableFuture() collection.updateOne( @@ -104,7 +113,7 @@ class Mongo4ReactiveClientTest extends AbstractMongoClientTest { } @Override - int delete(String dbName, String collectionName) { + MongoCollection setupDelete(String dbName, String collectionName) { MongoCollection 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 collection) { def result = new CompletableFuture() 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 setupGetMore(String dbName, String collectionName) { + throw new AssumptionViolatedException("not tested on reactive") + } + + @Override + void getMore(MongoCollection collection) { throw new AssumptionViolatedException("not tested on reactive") } diff --git a/instrumentation/mongo/mongo-4.0-testing/src/test/groovy/MongoClientTest.groovy b/instrumentation/mongo/mongo-4.0/javaagent/src/test/groovy/MongoClientTest.groovy similarity index 82% rename from instrumentation/mongo/mongo-4.0-testing/src/test/groovy/MongoClientTest.groovy rename to instrumentation/mongo/mongo-4.0/javaagent/src/test/groovy/MongoClientTest.groovy index 1142791c43..76c9340bee 100644 --- a/instrumentation/mongo/mongo-4.0-testing/src/test/groovy/MongoClientTest.groovy +++ b/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> { @Shared MongoClient client @@ -53,19 +53,24 @@ class MongoClientTest extends AbstractMongoClientTest { } @Override - int insert(String dbName, String collectionName) { + MongoCollection setupInsert(String dbName, String collectionName) { MongoCollection collection = runUnderTrace("setup") { MongoDatabase db = client.getDatabase(dbName) db.createCollection(collectionName) return db.getCollection(collectionName) } ignoreTracesAndClear(1) + return collection + } + + @Override + int insert(MongoCollection collection) { collection.insertOne(new Document("password", "SECRET")) return collection.estimatedDocumentCount() } @Override - int update(String dbName, String collectionName) { + MongoCollection setupUpdate(String dbName, String collectionName) { MongoCollection 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 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 setupDelete(String dbName, String collectionName) { MongoCollection 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 collection) { def result = collection.deleteOne(new BsonDocument("password", new BsonString("SECRET"))) collection.estimatedDocumentCount() return result.deletedCount } @Override - void getMore(String dbName, String collectionName) { + MongoCollection setupGetMore(String dbName, String collectionName) { MongoCollection 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 collection) { collection.find().filter(new Document("_id", new Document('$gte', 0))) .batchSize(2).into(new ArrayList()) } diff --git a/instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/MongoAsyncClientInstrumentationModule.java b/instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/MongoAsyncClientInstrumentationModule.java index 1e38cfea2b..61b9d2f0bf 100644 --- a/instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/MongoAsyncClientInstrumentationModule.java +++ b/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.Collections; +import java.util.HashMap; 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 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( - new TypeDescription.Latent( - "com.mongodb.event.CommandListener", - Modifier.PUBLIC, - null, - Collections.emptyList()))) - .and(isPublic()))) + takesArguments(1) + .and(takesArgument(0, named("com.mongodb.event.CommandListener")))))) .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 typeMatcher() { + return named("com.mongodb.connection.InternalStreamConnection"); + } + + @Override + public Map, String> transformers() { + Map, 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 typeMatcher() { + return named("com.mongodb.connection.BaseCluster"); + } + + @Override + public Map, 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 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 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 callback) { + callback = new SingleResultCallbackWrapper(Java8BytecodeBridge.currentContext(), callback); + } + } } diff --git a/instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/SingleResultCallbackWrapper.java b/instrumentation/mongo/mongo-async-3.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/mongoasync/v3_3/SingleResultCallbackWrapper.java new file mode 100644 index 0000000000..ebfbbe32aa --- /dev/null +++ b/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 { + private final Context context; + private final SingleResultCallback delegate; + + public SingleResultCallbackWrapper(Context context, SingleResultCallback delegate) { + this.context = context; + this.delegate = delegate; + } + + @Override + public void onResult(Object server, Throwable throwable) { + try (Scope ignored = context.makeCurrent()) { + delegate.onResult(server, throwable); + } + } +} diff --git a/instrumentation/mongo/mongo-async-3.3/javaagent/src/test/groovy/MongoAsyncClientTest.groovy b/instrumentation/mongo/mongo-async-3.3/javaagent/src/test/groovy/MongoAsyncClientTest.groovy index 1ea58f2994..b92bf786a9 100644 --- a/instrumentation/mongo/mongo-async-3.3/javaagent/src/test/groovy/MongoAsyncClientTest.groovy +++ b/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> { @Shared MongoClient client @@ -73,7 +73,7 @@ class MongoAsyncClientTest extends AbstractMongoClientTest { } @Override - int insert(String dbName, String collectionName) { + MongoCollection setupInsert(String dbName, String collectionName) { MongoCollection 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 collection) { def count = new CompletableFuture() 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 setupUpdate(String dbName, String collectionName) { MongoCollection 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 collection) { def result = new CompletableFuture() def count = new CompletableFuture() collection.updateOne( @@ -115,7 +125,7 @@ class MongoAsyncClientTest extends AbstractMongoClientTest { } @Override - int delete(String dbName, String collectionName) { + MongoCollection setupDelete(String dbName, String collectionName) { MongoCollection 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 collection) { def result = new CompletableFuture() 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 setupGetMore(String dbName, String collectionName) { + throw new AssumptionViolatedException("not tested on async") + } + + @Override + void getMore(MongoCollection collection) { throw new AssumptionViolatedException("not tested on async") } diff --git a/instrumentation/mongo/mongo-testing/src/main/groovy/AbstractMongoClientTest.groovy b/instrumentation/mongo/mongo-testing/src/main/groovy/AbstractMongoClientTest.groovy index bd3f0a6348..bfcf42522c 100644 --- a/instrumentation/mongo/mongo-testing/src/main/groovy/AbstractMongoClientTest.groovy +++ b/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 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) { - mongoSpan(it, 0, "create", collectionName, dbName) { + trace(0, 2) { + 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) { - mongoSpan(it, 0, "create", collectionName, dbName, { + trace(0, 2) { + 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) { - mongoSpan(it, 0, "count", collectionName, dbName) { + trace(0, 2) { + 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) { - trace(0, 1) { - mongoSpan(it, 0, "insert", collectionName, dbName) { + assertTraces(1) { + trace(0, 3) { + 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 } - } - trace(1, 1) { - mongoSpan(it, 0, "count", collectionName, dbName) { + mongoSpan(it, 2, "count", collectionName, dbName, span(0)) { 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) { - trace(0, 1) { - mongoSpan(it, 0, "update", collectionName, dbName) { + assertTraces(1) { + trace(0, 3) { + 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 } - } - trace(1, 1) { - mongoSpan(it, 0, "count", collectionName, dbName) { + mongoSpan(it, 2, "count", collectionName, dbName, span(0)) { 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) { - trace(0, 1) { - mongoSpan(it, 0, "delete", collectionName, dbName) { + assertTraces(1) { + trace(0, 3) { + 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 } - } - trace(1, 1) { - mongoSpan(it, 0, "count", collectionName, dbName) { + mongoSpan(it, 2, "count", collectionName, dbName, span(0)) { 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) { - trace(0, 1) { - mongoSpan(it, 0, "find", collectionName, dbName) { + assertTraces(1) { + trace(0, 3) { + basicSpan(it, 0, "parent") + mongoSpan(it, 1, "find", collectionName, dbName, span(0)) { assert it == '{"find":"' + collectionName + '","filter":{"_id":{"$gte":"?"}},"batchSize":"?"}' true } - } - trace(1, 1) { - mongoSpan(it, 0, "getMore", collectionName, dbName) { + mongoSpan(it, 2, "getMore", collectionName, dbName, span(0)) { 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) { - mongoSpan(it, 0, "create", collectionName, dbName) { + trace(0, 2) { + 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 statementEval, - Object parentSpan = null, Throwable exception = null) { + String dbName, Object parentSpan, + Closure statementEval, Throwable exception = null) { trace.span(index) { name { operation + " " + dbName + "." + collection } kind CLIENT diff --git a/settings.gradle b/settings.gradle index ed19bffe21..c2dc6f3d58 100644 --- a/settings.gradle +++ b/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'