Merge pull request #154 from DataDog/ark/mongo_bytebuddy

Mongo Bytebuddy and Mongo Async

.gitignore

@@ -51,3 +51,4 @@ Thumbs.db
 */out
 dd-java-agent/integrations/*/out
 dd-trace-examples/*/out
+derby.log

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'

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 {
+    TestUtils.registerOrReplaceGlobalTracer(tracer);
-      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);
-    }
     writer.start();
     assertThat(GlobalTracer.isRegistered()).isTrue();
   }

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/MongoAsyncClientInstrumentationTest.java

@@ -0,0 +1,121 @@
+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()));
+  }
+}

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/MongoClientInstrumentationTest.java

@@ -1,26 +1,116 @@
 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 test() {
+  public void syncClientHasListener() {
-    final MongoClient mongoClient = new MongoClient();
+    Assert.assertEquals(1, client.getMongoClientOptions().getCommandListeners().size());
+    Assert.assertEquals(
+        "DDTracingCommandListener",
+        client.getMongoClientOptions().getCommandListeners().get(0).getClass().getSimpleName());
+  }
 
-    assertThat(mongoClient.getMongoClientOptions().getCommandListeners().size()).isEqualTo(1);
+  @Test
-    assertThat(
+  public void insertOperation() throws UnknownHostException {
-            mongoClient
+    MongoDatabase db = client.getDatabase(MONGO_DB_NAME);
-                .getMongoClientOptions()
+    final String collectionName = "testCollection";
-                .getCommandListeners()
+    db.createCollection(collectionName);
-                .get(0)
+    MongoCollection<Document> collection = db.getCollection(collectionName);
-                .getClass()
-                .getSimpleName())
-        .isEqualTo("DDTracingCommandListener");
 
-    mongoClient.close();
+    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()));
   }
 }

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/TestUtils.java

@@ -0,0 +1,18 @@
+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);
+    }
+  }
+}

dd-java-agent/dd-java-agent.gradle

@@ -30,6 +30,12 @@ dependencies {
   compile(project(':dd-java-agent:integrations:jms-2')) {
     transitive = false
   }
+  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')) {
     transitive = false
   }
@@ -51,7 +57,6 @@ dependencies {
 
   testCompile deps.testLogging
   testCompile deps.opentracingMock
-  testCompile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.4.2'
 
   testCompile(project(path: ':dd-java-agent:integrations:helpers')) {
     transitive = false

dd-java-agent/integrations/helpers/src/main/java/com/datadoghq/agent/integration/DDTracingCommandListener.java

@@ -0,0 +1,147 @@
+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;
+  }
+}

dd-java-agent/integrations/helpers/src/main/java/com/datadoghq/agent/integration/MongoHelper.java

@@ -1,181 +0,0 @@
-package com.datadoghq.agent.integration;
-
-import com.datadoghq.trace.DDTags;
-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 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;
-import org.jboss.byteman.rule.Rule;
-
-/** Patch the Mongo builder before constructing the final client */
-@Slf4j
-public class MongoHelper extends DDAgentTracingHelper<MongoClientOptions.Builder> {
-
-  public MongoHelper(final Rule rule) {
-    super(rule);
-  }
-
-  /**
-   * Strategy: Just before com.mongodb.MongoClientOptions$Builder.build() method is called, we add a
-   * new command listener in charge of the tracing.
-   *
-   * @param builder The builder instance
-   * @return The same builder instance with a new tracing command listener that will be use for the
-   *     client construction
-   * @throws Exception
-   */
-  @Override
-  protected MongoClientOptions.Builder doPatch(final MongoClientOptions.Builder builder)
-      throws Exception {
-
-    final DDTracingCommandListener listener = new DDTracingCommandListener(tracer);
-    builder.addCommandListener(listener);
-
-    setState(builder, 1);
-
-    return builder;
-  }
-
-  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;
-    }
-  }
-}

dd-java-agent/integrations/mongo-3.1/mongo-3.1.gradle

@@ -0,0 +1,24 @@
+apply plugin: 'version-scan'
+
+versionScan {
+  group = "org.mongodb"
+  module = "mongo-java-driver"
+  scanMethods = true
+  versions = "[3.1,)"
+  verifyPresent = [
+    'com.mongodb.MongoClientOptions$Builder': 'addCommandListener'
+  ]
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+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
+  compile deps.opentracing
+}

dd-java-agent/integrations/mongo-3.1/src/main/java/dd/inst/mongo/MongoClientInstrumentation.java

@@ -0,0 +1,55 @@
+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);
+    }
+  }
+}

dd-java-agent/integrations/mongo-3.1/src/test/java/dd/inst/mongo/MongoClientInstrumentationTest.java

@@ -1,7 +1,8 @@
-package com.datadoghq.agent.integration;
+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;
@@ -17,7 +18,7 @@ import org.bson.BsonDocument;
 import org.bson.BsonString;
 import org.junit.Test;
 
-public class MongoHelperTest {
+public class MongoClientInstrumentationTest {
 
   private static ConnectionDescription makeConnection() {
     return new ConnectionDescription(new ServerId(new ClusterId(), new ServerAddress()));
@@ -29,7 +30,7 @@ public class MongoHelperTest {
         new CommandStartedEvent(1, makeConnection(), "databasename", "query", new BsonDocument());
 
     final DDSpan span = new DDTracer().buildSpan("foo").startManual();
-    MongoHelper.DDTracingCommandListener.decorate(span, cmd);
+    DDTracingCommandListener.decorate(span, cmd);
 
     assertThat(span.context().getSpanType()).isEqualTo("mongodb");
     assertThat(span.context().getResourceName())
@@ -51,7 +52,7 @@ public class MongoHelperTest {
           new CommandStartedEvent(1, makeConnection(), "databasename", "query", query);
 
       final DDSpan span = new DDTracer().buildSpan("foo").startManual();
-      MongoHelper.DDTracingCommandListener.decorate(span, cmd);
+      DDTracingCommandListener.decorate(span, cmd);
 
       assertThat(span.getTags().get(Tags.DB_STATEMENT.getKey()))
           .isEqualTo(query.toString().replaceAll("secret", "?"));

dd-java-agent/integrations/mongo-async-3.3/mongo-async-3.3.gradle

@@ -0,0 +1,24 @@
+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
+}

dd-java-agent/integrations/mongo-async-3.3/src/main/java/dd/inst/mongo/MongoAsyncClientInstrumentation.java

@@ -0,0 +1,55 @@
+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);
+    }
+  }
+}

dd-java-agent/integrations/mongo-async/mongo-async.gradle

@@ -1,11 +0,0 @@
-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,
-  ]
-}

dd-java-agent/integrations/mongo/mongo.gradle

@@ -1,11 +0,0 @@
-apply plugin: 'version-scan'
-
-versionScan {
-  group = "org.mongodb"
-  module = "mongo-java-driver"
-  versions = "[3.0,)"
-  verifyPresent = [
-    "com.mongodb.operation.AsyncReadOperation": null,
-    "com.mongodb.client.model.MapReduceAction": null,
-  ]
-}

dd-java-agent/src/main/resources/dd-trace-supported-framework.yaml

@@ -47,32 +47,6 @@ opentracing-jms-2_consumer:
       javax.jms.JMSContext:
       javax.jms.CompletionListener:
 
-opentracing-mongo-driver:
-  - artifact: mongo-java-driver
-    supported_version: 3\..*
-    identifying_present_classes:
-      com.mongodb.operation.AsyncReadOperation:
-      com.mongodb.client.model.MapReduceAction:
-
-  - artifact: mongodb-driver-async
-    supported_version: 3\..*
-    identifying_present_classes:
-      com.mongodb.operation.AsyncReadOperation:
-      com.mongodb.client.model.MapReduceAction:
-
-opentracing-mongo-driver-helper:
-  - artifact: mongo-java-driver
-    supported_version: 3\..*
-    identifying_present_classes:
-      com.mongodb.operation.AsyncReadOperation:
-      com.mongodb.client.model.MapReduceAction:
-
-  - artifact: mongodb-driver-async
-    supported_version: 3\..*
-    identifying_present_classes:
-      com.mongodb.operation.AsyncReadOperation:
-      com.mongodb.client.model.MapReduceAction:
-
 opentracing-okhttp3:
   - artifact: okhttp
     supported_version: 3\..*

dd-java-agent/src/main/resources/initializer-rules.btm

@@ -25,18 +25,6 @@ DO
 ENDRULE
 
 
-# Instrument Mongo client
-# ========================
-RULE MongoClientOptions$Builder-init
-CLASS com.mongodb.MongoClientOptions$Builder
-METHOD <init>
-AT EXIT
-IF TRUE
-DO
-  com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
-ENDRULE
-
-
 # Instrument OkHttp
 # ===========================
 RULE OkHttpClient$Builder-init

dd-java-agent/src/main/resources/integration-rules.btm

@@ -34,19 +34,6 @@ DO
 ENDRULE
 
 
-# Instrument Mongo client
-# ========================
-RULE opentracing-mongo-driver
-CLASS com.mongodb.MongoClientOptions$Builder
-METHOD build
-HELPER com.datadoghq.agent.integration.MongoHelper
-AT ENTRY
-IF getState($0) == 0
-DO
-  patch($0);
-ENDRULE
-
-
 # Instrument OkHttp
 # ===========================
 RULE opentracing-okhttp3

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 {

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-3.1'
-include ':dd-java-agent:integrations:mongo-async'
+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'