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 @@
 
 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 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 void beforeTest() throws Exception {
     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();
   }

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 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()));
   }
 }

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