Mongo Async Instrumentation and embedded mongo tests.
MODIFY
dd-java-agent-ittests/dd-java-agent-ittests.gradle
dd-java-agent-ittests/dd-java-agent-ittests.gradle
@@ -10,6 +10,10 @@ dependencies {
testCompile deps.testLogging
testCompile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.4.2'
+ testCompile group: 'org.mongodb', name: 'mongodb-driver-async', version: '3.4.2'
+ // run embeded mongodb for integration testing
+ testCompile group: 'de.flapdoodle.embed', name: 'de.flapdoodle.embed.mongo', version: '2.0.0'
+
testCompile group: 'org.eclipse.jetty', name: 'jetty-server', version: '9.4.1.v20170120'
testCompile group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '9.4.1.v20170120'
testCompile group: 'org.apache.tomcat.embed', name: 'tomcat-embed-core', version: '8.0.41'

MODIFY
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationsTest.java
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationsTest.java
@@ -2,6 +2,7 @@ package com.datadoghq.agent.instrumentation.annotation;
import static org.assertj.core.api.Assertions.assertThat;
+import com.datadoghq.agent.integration.TestUtils;
import com.datadoghq.agent.test.SayTracedHello;
import com.datadoghq.trace.DDBaseSpan;
import com.datadoghq.trace.DDTracer;
@@ -10,12 +11,10 @@ import com.datadoghq.trace.writer.ListWriter;
import io.opentracing.util.GlobalTracer;
import java.io.PrintWriter;
import java.io.StringWriter;
-import java.lang.reflect.Field;
import org.junit.Before;
import org.junit.Test;
public class TraceAnnotationsTest {
-
private final ListWriter writer = new ListWriter();
private final DDTracer tracer = new DDTracer(writer);
@@ -24,14 +23,8 @@ public class TraceAnnotationsTest {
Class.forName("com.datadoghq.agent.InstrumentationRulesManager")
.getMethod("registerClassLoad")
.invoke(null);
- try {
- GlobalTracer.register(tracer);
- } catch (final Exception e) {
- // Force it anyway using reflection
- final Field field = GlobalTracer.class.getDeclaredField("tracer");
- field.setAccessible(true);
- field.set(null, tracer);
- }
+ TestUtils.registerOrReplaceGlobalTracer(tracer);
+
writer.start();
assertThat(GlobalTracer.isRegistered()).isTrue();
}

ADD
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/MongoAsyncClientInstrumentationTest.java
+package com.datadoghq.agent.integration;
+
+import static com.datadoghq.agent.integration.MongoClientInstrumentationTest.MONGO_DB_NAME;
+import static com.datadoghq.agent.integration.MongoClientInstrumentationTest.MONGO_HOST;
+import static com.datadoghq.agent.integration.MongoClientInstrumentationTest.MONGO_PORT;
+
+import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.DDTracer;
+import com.datadoghq.trace.writer.ListWriter;
+import com.mongodb.async.SingleResultCallback;
+import com.mongodb.async.client.MongoClient;
+import com.mongodb.async.client.MongoClients;
+import com.mongodb.async.client.MongoDatabase;
+import io.opentracing.tag.Tags;
+import java.net.InetAddress;
+import java.nio.ByteBuffer;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.bson.Document;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public class MongoAsyncClientInstrumentationTest {
+ private static MongoClient client;
+ private static final ListWriter writer = new ListWriter();
+ private static final DDTracer tracer = new DDTracer(writer);
+
+ @BeforeClass
+ public static void setup() throws Exception {
+ TestUtils.registerOrReplaceGlobalTracer(tracer);
+ MongoClientInstrumentationTest.startLocalMongo();
+ client = MongoClients.create("mongodb://" + MONGO_HOST + ":" + MONGO_PORT);
+ }
+
+ @AfterClass
+ public static void destroy() throws Exception {
+ if (null != client) {
+ client.close();
+ client = null;
+ }
+ MongoClientInstrumentationTest.stopLocalMongo();
+ }
+
+ @Test
+ public void asyncClientHasListener() {
+ Assert.assertEquals(1, client.getSettings().getCommandListeners().size());
+ Assert.assertEquals(
+ "DDTracingCommandListener",
+ client.getSettings().getCommandListeners().get(0).getClass().getSimpleName());
+ }
+
+ @Test
+ public void insertOperation() throws InterruptedException, Exception {
+ MongoDatabase db = client.getDatabase(MONGO_DB_NAME);
+ final String collectionName = "asyncCollection";
+ final AtomicBoolean done = new AtomicBoolean(false);
+
+ db.createCollection(
+ collectionName,
+ new SingleResultCallback<Void>() {
+ @Override
+ public void onResult(Void result, Throwable t) {
+ done.set(true);
+ }
+ });
+ while (!done.get()) {
+ Thread.sleep(1);
+ }
+
+ db.getCollection(collectionName)
+ .insertOne(
+ new Document("foo", "bar"),
+ new SingleResultCallback<Void>() {
+ @Override
+ public void onResult(final Void result, final Throwable t) {
+ done.set(true);
+ }
+ });
+ while (!done.get()) {
+ Thread.sleep(1);
+ }
+
+ done.set(false);
+ db.getCollection(collectionName)
+ .count(
+ new SingleResultCallback<Long>() {
+ @Override
+ public void onResult(Long result, Throwable t) {
+ Assert.assertEquals(1, result.longValue());
+ done.set(true);
+ }
+ });
+
+ while (!done.get()) {
+ Thread.sleep(1);
+ }
+
+ // the final trace may still be reporting to the ListWriter,
+ // but we're only testing the first trace.
+ Assert.assertTrue(writer.getList().size() >= 1);
+
+ final String createCollectionQuery =
+ "{ \"create\" : \"asyncCollection\", \"autoIndexId\" : \"?\", \"capped\" : \"?\" }";
+ final DDBaseSpan<?> trace0 = writer.get(0).get(0);
+ Assert.assertEquals("mongo.query", trace0.getOperationName());
+ Assert.assertEquals(createCollectionQuery, trace0.getResourceName());
+ Assert.assertEquals("mongodb", trace0.getType());
+ Assert.assertEquals("mongo", trace0.getServiceName());
+
+ Assert.assertEquals("java-mongo", trace0.getTags().get(Tags.COMPONENT.getKey()));
+ Assert.assertEquals(createCollectionQuery, trace0.getTags().get(Tags.DB_STATEMENT.getKey()));
+ Assert.assertEquals(MONGO_DB_NAME, trace0.getTags().get(Tags.DB_INSTANCE.getKey()));
+ Assert.assertEquals(MONGO_HOST, trace0.getTags().get(Tags.PEER_HOSTNAME.getKey()));
+ Assert.assertEquals(
+ ByteBuffer.wrap(InetAddress.getByName("127.0.0.1").getAddress()).getInt(),
+ trace0.getTags().get(Tags.PEER_HOST_IPV4.getKey()));
+ Assert.assertEquals(MONGO_PORT, trace0.getTags().get(Tags.PEER_PORT.getKey()));
+ Assert.assertEquals("mongo", trace0.getTags().get(Tags.DB_TYPE.getKey()));
+ }
+}

MODIFY
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/MongoClientInstrumentationTest.java
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/MongoClientInstrumentationTest.java
package com.datadoghq.agent.integration;
-import static org.assertj.core.api.Assertions.assertThat;
-
+import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.DDTracer;
+import com.datadoghq.trace.writer.ListWriter;
import com.mongodb.MongoClient;
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.MongoDatabase;
+import de.flapdoodle.embed.mongo.*;
+import de.flapdoodle.embed.mongo.config.IMongodConfig;
+import de.flapdoodle.embed.mongo.config.MongodConfigBuilder;
+import de.flapdoodle.embed.mongo.config.Net;
+import de.flapdoodle.embed.mongo.distribution.Version;
+import de.flapdoodle.embed.process.runtime.Network;
+import io.opentracing.tag.Tags;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.nio.ByteBuffer;
+import org.bson.Document;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
import org.junit.Test;
public class MongoClientInstrumentationTest {
+ public static final String MONGO_DB_NAME = "embedded";
+ public static final String MONGO_HOST = "localhost";
+ public static final int MONGO_PORT = 12345;
+ private static MongodExecutable mongodExe;
+ private static MongodProcess mongod;
+
+ private static MongoClient client;
+ private static final ListWriter writer = new ListWriter();
+ private static final DDTracer tracer = new DDTracer(writer);
+
+ public static void startLocalMongo() throws Exception {
+ MongodStarter starter = MongodStarter.getDefaultInstance();
+
+ IMongodConfig mongodConfig =
+ new MongodConfigBuilder()
+ .version(Version.Main.PRODUCTION)
+ .net(new Net(MONGO_HOST, MONGO_PORT, Network.localhostIsIPv6()))
+ .build();
+
+ mongodExe = starter.prepare(mongodConfig);
+ mongod = mongodExe.start();
+ }
+
+ public static void stopLocalMongo() throws Exception {
+ if (null != mongod) {
+ mongod.stop();
+ mongod = null;
+ }
+ if (null != mongodExe) {
+ mongodExe.stop();
+ mongodExe = null;
+ }
+ }
+
+ @BeforeClass
+ public static void setup() throws Exception {
+ TestUtils.registerOrReplaceGlobalTracer(tracer);
+ startLocalMongo();
+
+ client = new MongoClient(MONGO_HOST, MONGO_PORT);
+ }
+
+ @AfterClass
+ public static void destroy() throws Exception {
+ if (null != client) {
+ client.close();
+ client = null;
+ }
+ stopLocalMongo();
+ }
+
+ @Test
+ public void syncClientHasListener() {
+ Assert.assertEquals(1, client.getMongoClientOptions().getCommandListeners().size());
+ Assert.assertEquals(
+ "DDTracingCommandListener",
+ client.getMongoClientOptions().getCommandListeners().get(0).getClass().getSimpleName());
+ }
@Test
- public void test() {
- final MongoClient mongoClient = new MongoClient();
-
- assertThat(mongoClient.getMongoClientOptions().getCommandListeners().size()).isEqualTo(1);
- assertThat(
- mongoClient
- .getMongoClientOptions()
- .getCommandListeners()
- .get(0)
- .getClass()
- .getSimpleName())
- .isEqualTo("DDTracingCommandListener");
-
- mongoClient.close();
+ public void insertOperation() throws UnknownHostException {
+ MongoDatabase db = client.getDatabase(MONGO_DB_NAME);
+ final String collectionName = "testCollection";
+ db.createCollection(collectionName);
+ MongoCollection<Document> collection = db.getCollection(collectionName);
+
+ collection.insertOne(new Document("foo", "bar"));
+
+ Assert.assertEquals(1, collection.count());
+
+ Assert.assertEquals(3, writer.getList().size());
+
+ final String createCollectionQuery =
+ "{ \"create\" : \"testCollection\", \"autoIndexId\" : \"?\", \"capped\" : \"?\" }";
+ final DDBaseSpan<?> trace0 = writer.get(0).get(0);
+ Assert.assertEquals("mongo.query", trace0.getOperationName());
+ Assert.assertEquals(createCollectionQuery, trace0.getResourceName());
+ Assert.assertEquals("mongodb", trace0.getType());
+ Assert.assertEquals("mongo", trace0.getServiceName());
+
+ Assert.assertEquals("java-mongo", trace0.getTags().get(Tags.COMPONENT.getKey()));
+ Assert.assertEquals(createCollectionQuery, trace0.getTags().get(Tags.DB_STATEMENT.getKey()));
+ Assert.assertEquals(MONGO_DB_NAME, trace0.getTags().get(Tags.DB_INSTANCE.getKey()));
+ Assert.assertEquals(MONGO_HOST, trace0.getTags().get(Tags.PEER_HOSTNAME.getKey()));
+ Assert.assertEquals(
+ ByteBuffer.wrap(InetAddress.getByName("127.0.0.1").getAddress()).getInt(),
+ trace0.getTags().get(Tags.PEER_HOST_IPV4.getKey()));
+ Assert.assertEquals(MONGO_PORT, trace0.getTags().get(Tags.PEER_PORT.getKey()));
+ Assert.assertEquals("mongo", trace0.getTags().get(Tags.DB_TYPE.getKey()));
}
}

ADD
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/TestUtils.java
+package com.datadoghq.agent.integration;
+
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Field;
+
+public class TestUtils {
+ public static void registerOrReplaceGlobalTracer(Tracer tracer) throws Exception {
+ try {
+ GlobalTracer.register(tracer);
+ } catch (final Exception e) {
+ // Force it anyway using reflection
+ final Field field = GlobalTracer.class.getDeclaredField("tracer");
+ field.setAccessible(true);
+ field.set(null, tracer);
+ }
+ }
+}

MODIFY
dd-java-agent/dd-java-agent.gradle
dd-java-agent/dd-java-agent.gradle
@@ -30,7 +30,10 @@ dependencies {
compile(project(':dd-java-agent:integrations:jms-2')) {
transitive = false
}
- compile(project(':dd-java-agent:integrations:mongo')) {
+ compile(project(':dd-java-agent:integrations:mongo-3.1')) {
+ transitive = false
+ }
+ compile(project(':dd-java-agent:integrations:mongo-async-3.3')) {
transitive = false
}
compile(project(':dd-java-agent:integrations:servlet-2')) {

ADD
dd-java-agent/integrations/helpers/src/main/java/com/datadoghq/agent/integration/DDTracingCommandListener.java
+package com.datadoghq.agent.integration;
+
+import com.datadoghq.trace.DDTags;
+import com.mongodb.event.CommandFailedEvent;
+import com.mongodb.event.CommandListener;
+import com.mongodb.event.CommandStartedEvent;
+import com.mongodb.event.CommandSucceededEvent;
+import io.opentracing.Span;
+import io.opentracing.Tracer;
+import io.opentracing.tag.Tags;
+import java.net.Inet4Address;
+import java.net.InetAddress;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import lombok.extern.slf4j.Slf4j;
+import org.bson.BsonArray;
+import org.bson.BsonDocument;
+import org.bson.BsonString;
+import org.bson.BsonValue;
+
+@Slf4j
+public class DDTracingCommandListener implements CommandListener {
+ /**
+ * The values of these mongo fields will not be scrubbed out. This allows the non-sensitive
+ * collection names to be captured.
+ */
+ private static final List<String> UNSCRUBBED_FIELDS =
+ Arrays.asList("ordered", "insert", "count", "find", "create");
+
+ private static final BsonValue HIDDEN_CHAR = new BsonString("?");
+
+ private static final String MONGO_OPERATION = "mongo.query";
+ private static final String COMPONENT_NAME = "java-mongo";
+
+ private final Tracer tracer;
+ /** requestID -> span */
+ private final Map<Integer, Span> cache = new ConcurrentHashMap<>();
+
+ public DDTracingCommandListener(Tracer tracer) {
+ this.tracer = tracer;
+ }
+
+ @Override
+ public void commandStarted(CommandStartedEvent event) {
+ Span span = buildSpan(event);
+ cache.put(event.getRequestId(), span);
+ }
+
+ @Override
+ public void commandSucceeded(CommandSucceededEvent event) {
+ Span span = cache.remove(event.getRequestId());
+ if (span != null) {
+ span.finish();
+ }
+ }
+
+ @Override
+ public void commandFailed(CommandFailedEvent event) {
+ Span span = cache.remove(event.getRequestId());
+ if (span != null) {
+ Tags.ERROR.set(span, Boolean.TRUE);
+ span.log(Collections.singletonMap("error.object", event.getThrowable()));
+ span.finish();
+ }
+ }
+
+ private Span buildSpan(CommandStartedEvent event) {
+ Tracer.SpanBuilder spanBuilder =
+ tracer.buildSpan(MONGO_OPERATION).withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT);
+
+ final Span span = spanBuilder.startManual();
+ try {
+ decorate(span, event);
+ } catch (final Throwable e) {
+ log.warn("Couldn't decorate the mongo query: " + e.getMessage(), e);
+ }
+
+ return span;
+ }
+
+ public static void decorate(Span span, CommandStartedEvent event) {
+ // scrub the Mongo command so that parameters are removed from the string
+ final BsonDocument scrubbed = scrub(event.getCommand());
+ final String mongoCmd = scrubbed.toString();
+
+ Tags.COMPONENT.set(span, COMPONENT_NAME);
+ Tags.DB_STATEMENT.set(span, mongoCmd);
+ Tags.DB_INSTANCE.set(span, event.getDatabaseName());
+
+ Tags.PEER_HOSTNAME.set(span, event.getConnectionDescription().getServerAddress().getHost());
+
+ InetAddress inetAddress =
+ event.getConnectionDescription().getServerAddress().getSocketAddress().getAddress();
+ if (inetAddress instanceof Inet4Address) {
+ byte[] address = inetAddress.getAddress();
+ Tags.PEER_HOST_IPV4.set(span, ByteBuffer.wrap(address).getInt());
+ } else {
+ Tags.PEER_HOST_IPV6.set(span, inetAddress.getHostAddress());
+ }
+
+ Tags.PEER_PORT.set(span, event.getConnectionDescription().getServerAddress().getPort());
+ Tags.DB_TYPE.set(span, "mongo");
+
+ // dd-specific tags
+ span.setTag(DDTags.RESOURCE_NAME, mongoCmd);
+ span.setTag(DDTags.SPAN_TYPE, "mongodb");
+ span.setTag(DDTags.SERVICE_NAME, "mongo");
+ }
+
+ private static BsonDocument scrub(final BsonDocument origin) {
+ final BsonDocument scrub = new BsonDocument();
+ for (final Map.Entry<String, BsonValue> entry : origin.entrySet()) {
+ if (UNSCRUBBED_FIELDS.contains(entry.getKey()) && entry.getValue().isString()) {
+ scrub.put(entry.getKey(), entry.getValue());
+ } else {
+ final BsonValue child = scrub(entry.getValue());
+ scrub.put(entry.getKey(), child);
+ }
+ }
+ return scrub;
+ }
+
+ private static BsonValue scrub(final BsonArray origin) {
+ final BsonArray scrub = new BsonArray();
+ for (final BsonValue value : origin) {
+ final BsonValue child = scrub(value);
+ scrub.add(child);
+ }
+ return scrub;
+ }
+
+ private static BsonValue scrub(final BsonValue origin) {
+ final BsonValue scrubbed;
+ if (origin.isDocument()) {
+ scrubbed = scrub(origin.asDocument());
+ } else if (origin.isArray()) {
+ scrubbed = scrub(origin.asArray());
+ } else {
+ scrubbed = HIDDEN_CHAR;
+ }
+ return scrubbed;
+ }
+}

RENAME
dd-java-agent/integrations/mongo/mongo.gradle
dd-java-agent/integrations/mongo-3.1/mongo-3.1.gradle
@@ -3,10 +3,10 @@ apply plugin: 'version-scan'
versionScan {
group = "org.mongodb"
module = "mongo-java-driver"
- versions = "[3.0,)"
+ scanMethods = true
+ versions = "[3.1,)"
verifyPresent = [
- "com.mongodb.operation.AsyncReadOperation": null,
- "com.mongodb.client.model.MapReduceAction": null,
+ 'com.mongodb.MongoClientOptions$Builder': 'addCommandListener'
]
}
@@ -16,6 +16,7 @@ dependencies {
compile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.4.2'
compile project(':dd-trace')
+ compile project(':dd-java-agent:integrations:helpers')
compile project(':dd-java-agent:tooling')
compile deps.bytebuddy

ADD
dd-java-agent/integrations/mongo-3.1/src/main/java/dd/inst/mongo/MongoClientInstrumentation.java
+package dd.inst.mongo;
+
+import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.datadoghq.agent.integration.DDTracingCommandListener;
+import com.google.auto.service.AutoService;
+import com.mongodb.MongoClientOptions;
+import dd.trace.Instrumenter;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Modifier;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+
+@AutoService(Instrumenter.class)
+public final class MongoClientInstrumentation implements Instrumenter {
+
+ @Override
+ public AgentBuilder instrument(AgentBuilder agentBuilder) {
+ return agentBuilder
+ .type(
+ named("com.mongodb.MongoClientOptions$Builder")
+ .and(
+ declaresMethod(
+ named("addCommandListener")
+ .and(isPublic())
+ .and(
+ takesArguments(
+ new TypeDescription.Latent(
+ "com.mongodb.event.CommandListener",
+ Modifier.PUBLIC,
+ null,
+ new TypeDescription.Generic[] {}))))))
+ .transform(
+ new AgentBuilder.Transformer.ForAdvice()
+ .advice(
+ isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)),
+ MongoClientAdvice.class.getName())
+ .withExceptionHandler(defaultExceptionHandler()))
+ .asDecorator();
+ }
+
+ public static class MongoClientAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static void injectTraceListener(@Advice.This final Object dis) {
+ // referencing "this" in the method args causes the class to load under a transformer.
+ // This bypasses the Builder instrumentation. Casting as a workaround.
+ MongoClientOptions.Builder builder = (MongoClientOptions.Builder) dis;
+ final DDTracingCommandListener listener = new DDTracingCommandListener(GlobalTracer.get());
+ builder.addCommandListener(listener);
+ }
+ }
+}

RENAME
dd-java-agent/integrations/mongo/src/test/java/dd/inst/mongo/MongoClientInstrumentationTest.java
dd-java-agent/integrations/mongo-3.1/src/test/java/dd/inst/mongo/MongoClientInstrumentationTest.java
@@ -2,6 +2,7 @@ package dd.inst.mongo;
import static org.assertj.core.api.Java6Assertions.assertThat;
+import com.datadoghq.agent.integration.DDTracingCommandListener;
import com.datadoghq.trace.DDSpan;
import com.datadoghq.trace.DDTracer;
import com.mongodb.ServerAddress;
@@ -9,7 +10,6 @@ import com.mongodb.connection.ClusterId;
import com.mongodb.connection.ConnectionDescription;
import com.mongodb.connection.ServerId;
import com.mongodb.event.CommandStartedEvent;
-import dd.inst.mongo.MongoClientInstrumentation.DDTracingCommandListener;
import io.opentracing.tag.Tags;
import java.util.Arrays;
import java.util.List;

ADD
dd-java-agent/integrations/mongo-async-3.3/mongo-async-3.3.gradle
+apply plugin: 'version-scan'
+
+versionScan {
+ group = "org.mongodb"
+ module = "mongodb-driver-async"
+ scanMethods = true
+ versions = "[3.3,)"
+ verifyPresent = [
+ 'com.mongodb.async.client.MongoClientSettings$Builder': 'addCommandListener'
+ ]
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+dependencies {
+ compile group: 'org.mongodb', name: 'mongodb-driver-async', version: '3.4.2'
+
+ compile project(':dd-trace')
+ compile project(':dd-java-agent:integrations:helpers')
+ compile project(':dd-java-agent:tooling')
+
+ compile deps.bytebuddy
+ compile deps.opentracing
+}
\ No newline at end of file

ADD
dd-java-agent/integrations/mongo-async-3.3/src/main/java/dd/inst/mongo/MongoAsyncClientInstrumentation.java
+package dd.inst.mongo;
+
+import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.datadoghq.agent.integration.DDTracingCommandListener;
+import com.google.auto.service.AutoService;
+import com.mongodb.async.client.MongoClientSettings;
+import dd.trace.Instrumenter;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Modifier;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+
+@AutoService(Instrumenter.class)
+public final class MongoAsyncClientInstrumentation implements Instrumenter {
+
+ @Override
+ public AgentBuilder instrument(AgentBuilder agentBuilder) {
+ return agentBuilder
+ .type(
+ named("com.mongodb.async.client.MongoClientSettings$Builder")
+ .and(
+ declaresMethod(
+ named("addCommandListener")
+ .and(isPublic())
+ .and(
+ takesArguments(
+ new TypeDescription.Latent(
+ "com.mongodb.event.CommandListener",
+ Modifier.PUBLIC,
+ null,
+ new TypeDescription.Generic[] {}))))))
+ .transform(
+ new AgentBuilder.Transformer.ForAdvice()
+ .advice(
+ isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)),
+ MongoAsyncClientAdvice.class.getName())
+ .withExceptionHandler(defaultExceptionHandler()))
+ .asDecorator();
+ }
+
+ public static class MongoAsyncClientAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static void injectTraceListener(@Advice.This final Object dis) {
+ // referencing "this" in the method args causes the class to load under a transformer.
+ // This bypasses the Builder instrumentation. Casting as a workaround.
+ final MongoClientSettings.Builder builder = (MongoClientSettings.Builder) dis;
+ final DDTracingCommandListener listener = new DDTracingCommandListener(GlobalTracer.get());
+ builder.addCommandListener(listener);
+ }
+ }
+}

DELETE
dd-java-agent/integrations/mongo-async/mongo-async.gradle
-apply plugin: 'version-scan'
-
-versionScan {
- group = "org.mongodb"
- module = "mongodb-driver-async"
- versions = "[3.0,)"
- verifyPresent = [
- "com.mongodb.operation.AsyncReadOperation": null,
- "com.mongodb.client.model.MapReduceAction": null,
- ]
-}

DELETE
dd-java-agent/integrations/mongo/src/main/java/dd/inst/mongo/MongoClientInstrumentation.java
-package dd.inst.mongo;
-
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
-import static net.bytebuddy.matcher.ElementMatchers.*;
-
-import com.datadoghq.trace.DDTags;
-import com.google.auto.service.AutoService;
-import com.mongodb.MongoClientOptions;
-import com.mongodb.event.CommandFailedEvent;
-import com.mongodb.event.CommandListener;
-import com.mongodb.event.CommandStartedEvent;
-import com.mongodb.event.CommandSucceededEvent;
-import dd.trace.Instrumenter;
-import io.opentracing.Span;
-import io.opentracing.Tracer;
-import io.opentracing.tag.Tags;
-import io.opentracing.util.GlobalTracer;
-import java.net.Inet4Address;
-import java.net.InetAddress;
-import java.nio.ByteBuffer;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import lombok.extern.slf4j.Slf4j;
-import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.asm.Advice;
-import org.bson.BsonArray;
-import org.bson.BsonDocument;
-import org.bson.BsonString;
-import org.bson.BsonValue;
-
-@Slf4j
-@AutoService(Instrumenter.class)
-public final class MongoClientInstrumentation implements Instrumenter {
-
- @Override
- public AgentBuilder instrument(AgentBuilder agentBuilder) {
- return agentBuilder
- .type(named("com.mongodb.MongoClientOptions$Builder"))
- .transform(
- new AgentBuilder.Transformer.ForAdvice()
- .advice(
- isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)),
- MongoClientAdvice.class.getName())
- .withExceptionHandler(defaultExceptionHandler()))
- .asDecorator();
- }
-
- public static class MongoClientAdvice {
-
- @Advice.OnMethodEnter(suppress = Throwable.class)
- public static void injectTraceListener(@Advice.This final Object dis) {
- // referencing "this" in the method args causes the class to load under a transformer.
- // This bypasses the Builder instrumentation. Casting as a workaround.
- MongoClientOptions.Builder builder = (MongoClientOptions.Builder) dis;
- final DDTracingCommandListener listener = new DDTracingCommandListener(GlobalTracer.get());
- builder.addCommandListener(listener);
- }
- }
-
- public static class DDTracingCommandListener implements CommandListener {
- /**
- * The values of these mongo fields will not be scrubbed out. This allows the non-sensitive
- * collection names to be captured.
- */
- private static final List<String> UNSCRUBBED_FIELDS =
- Arrays.asList("ordered", "insert", "count", "find", "create");
-
- private static final BsonValue HIDDEN_CHAR = new BsonString("?");
- private static final String MONGO_OPERATION = "mongo.query";
-
- static final String COMPONENT_NAME = "java-mongo";
- private final Tracer tracer;
- /** Cache for (request id, span) pairs */
- private final Map<Integer, Span> cache = new ConcurrentHashMap<>();
-
- public DDTracingCommandListener(Tracer tracer) {
- this.tracer = tracer;
- }
-
- @Override
- public void commandStarted(CommandStartedEvent event) {
- Span span = buildSpan(event);
- cache.put(event.getRequestId(), span);
- }
-
- @Override
- public void commandSucceeded(CommandSucceededEvent event) {
- Span span = cache.remove(event.getRequestId());
- if (span != null) {
- span.finish();
- }
- }
-
- @Override
- public void commandFailed(CommandFailedEvent event) {
- Span span = cache.remove(event.getRequestId());
- if (span != null) {
- onError(span, event.getThrowable());
- span.finish();
- }
- }
-
- private Span buildSpan(CommandStartedEvent event) {
- Tracer.SpanBuilder spanBuilder =
- tracer.buildSpan(MONGO_OPERATION).withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT);
-
- Span span = spanBuilder.startManual();
- try {
- decorate(span, event);
- } catch (final Throwable e) {
- log.warn("Couldn't decorate the mongo query: " + e.getMessage(), e);
- }
-
- return span;
- }
-
- private static void onError(Span span, Throwable throwable) {
- Tags.ERROR.set(span, Boolean.TRUE);
- span.log(Collections.singletonMap("error.object", throwable));
- }
-
- public static void decorate(Span span, CommandStartedEvent event) {
- // scrub the Mongo command so that parameters are removed from the string
- final BsonDocument scrubbed = scrub(event.getCommand());
- final String mongoCmd = scrubbed.toString();
-
- Tags.COMPONENT.set(span, COMPONENT_NAME);
- Tags.DB_STATEMENT.set(span, mongoCmd);
- Tags.DB_INSTANCE.set(span, event.getDatabaseName());
- // add specific resource name
- span.setTag(DDTags.RESOURCE_NAME, mongoCmd);
- span.setTag(DDTags.SPAN_TYPE, "mongodb");
- span.setTag(DDTags.SERVICE_NAME, "mongo");
-
- Tags.PEER_HOSTNAME.set(span, event.getConnectionDescription().getServerAddress().getHost());
-
- InetAddress inetAddress =
- event.getConnectionDescription().getServerAddress().getSocketAddress().getAddress();
-
- if (inetAddress instanceof Inet4Address) {
- byte[] address = inetAddress.getAddress();
- Tags.PEER_HOST_IPV4.set(span, ByteBuffer.wrap(address).getInt());
- } else {
- Tags.PEER_HOST_IPV6.set(span, inetAddress.getHostAddress());
- }
-
- Tags.PEER_PORT.set(span, event.getConnectionDescription().getServerAddress().getPort());
- Tags.DB_TYPE.set(span, "mongo");
- }
-
- private static BsonDocument scrub(final BsonDocument origin) {
- final BsonDocument scrub = new BsonDocument();
- for (final Map.Entry<String, BsonValue> entry : origin.entrySet()) {
- if (UNSCRUBBED_FIELDS.contains(entry.getKey()) && entry.getValue().isString()) {
- scrub.put(entry.getKey(), entry.getValue());
- } else {
- final BsonValue child = scrub(entry.getValue());
- scrub.put(entry.getKey(), child);
- }
- }
- return scrub;
- }
-
- private static BsonValue scrub(final BsonArray origin) {
- final BsonArray scrub = new BsonArray();
- for (final BsonValue value : origin) {
- final BsonValue child = scrub(value);
- scrub.add(child);
- }
- return scrub;
- }
-
- private static BsonValue scrub(final BsonValue origin) {
- final BsonValue scrubbed;
- if (origin.isDocument()) {
- scrubbed = scrub(origin.asDocument());
- } else if (origin.isArray()) {
- scrubbed = scrub(origin.asArray());
- } else {
- scrubbed = HIDDEN_CHAR;
- }
- return scrubbed;
- }
- }
-}

MODIFY
dd-trace/dd-trace.gradle
dd-trace/dd-trace.gradle
@@ -31,6 +31,9 @@ dependencies {
testCompile group: 'org.objenesis', name: 'objenesis', version: '2.6'
testCompile group: 'cglib', name: 'cglib-nodep', version: '3.2.5'
+
+ testCompile 'org.openjdk.jmh:jmh-core:1.19'
+ testCompile 'org.openjdk.jmh:jmh-generator-annprocess:1.19'
}
jmh {

MODIFY
settings.gradle
settings.gradle
@@ -18,8 +18,8 @@ include ':dd-java-agent:integrations:aws-sdk'
include ':dd-java-agent:integrations:cassandra'
include ':dd-java-agent:integrations:jms-1'
include ':dd-java-agent:integrations:jms-2'
-include ':dd-java-agent:integrations:mongo'
-include ':dd-java-agent:integrations:mongo-async'
+include ':dd-java-agent:integrations:mongo-3.1'
+include ':dd-java-agent:integrations:mongo-async-3.3'
include ':dd-java-agent:integrations:okhttp'
include ':dd-java-agent:integrations:servlet-2'
include ':dd-java-agent:integrations:servlet-3'

Mongo Async Instrumentation and embedded mongo tests.