Merge pull request #750 from DataDog/tyler/mongo-decorator

Migrate Mongo to Decorator

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

@@ -27,9 +27,11 @@ dependencies {
   annotationProcessor deps.autoservice
   implementation deps.autoservice
 
-  testCompile project(':dd-trace-ot')
+  testCompile project(':dd-java-agent:testing')
+  
   testCompile group: 'org.assertj', name: 'assertj-core', version: '2.9.+'
   testCompile group: 'org.mockito', name: 'mockito-core', version: '2.19.0'
+
   testCompile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'
   latestDepTestCompile group: 'org.mongodb', name: 'mongo-java-driver', version: '+'
 }

dd-java-agent/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/DDTracingCommandListener.java

@@ -1,148 +1,53 @@
 package datadog.trace.instrumentation.mongo;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.mongo.MongoClientDecorator.DECORATE;
 
 import com.mongodb.event.CommandFailedEvent;
 import com.mongodb.event.CommandListener;
 import com.mongodb.event.CommandStartedEvent;
 import com.mongodb.event.CommandSucceededEvent;
-import datadog.trace.api.DDSpanTypes;
-import datadog.trace.api.DDTags;
 import io.opentracing.Span;
-import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
-import io.opentracing.tag.Tags;
-import java.net.Inet4Address;
-import java.net.InetAddress;
-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 final Map<Integer, Span> spanMap = new ConcurrentHashMap<>();
-
-  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(final Tracer tracer) {
-    this.tracer = tracer;
-  }
 
   @Override
   public void commandStarted(final CommandStartedEvent event) {
-    final Span span = buildSpan(event);
+    final Span span = GlobalTracer.get().buildSpan("mongo.query").start();
-    cache.put(event.getRequestId(), span);
+    DECORATE.afterStart(span);
+    DECORATE.onConnection(span, event);
+    if (event.getConnectionDescription() != null
+        && event.getConnectionDescription() != null
+        && event.getConnectionDescription().getServerAddress() != null) {
+      DECORATE.onPeerConnection(
+          span, event.getConnectionDescription().getServerAddress().getSocketAddress());
+    }
+    DECORATE.onStatement(span, event.getCommand());
+    spanMap.put(event.getRequestId(), span);
   }
 
   @Override
   public void commandSucceeded(final CommandSucceededEvent event) {
-    final Span span = cache.remove(event.getRequestId());
+    final Span span = spanMap.remove(event.getRequestId());
     if (span != null) {
+      DECORATE.beforeFinish(span);
       span.finish();
     }
   }
 
   @Override
   public void commandFailed(final CommandFailedEvent event) {
-    final Span span = cache.remove(event.getRequestId());
+    final Span span = spanMap.remove(event.getRequestId());
     if (span != null) {
-      Tags.ERROR.set(span, Boolean.TRUE);
+      DECORATE.onError(span, event.getThrowable());
-      span.log(Collections.singletonMap(ERROR_OBJECT, event.getThrowable()));
+      DECORATE.beforeFinish(span);
       span.finish();
     }
   }
-
-  private Span buildSpan(final CommandStartedEvent event) {
-    final Tracer.SpanBuilder spanBuilder =
-        tracer.buildSpan(MONGO_OPERATION).withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT);
-
-    final Span span = spanBuilder.start();
-    try {
-      decorate(span, event);
-    } catch (final Throwable e) {
-      log.debug("Couldn't decorate the mongo query: " + e.getMessage(), e);
-    }
-
-    return span;
-  }
-
-  public static void decorate(final Span span, final 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());
-
-    final InetAddress inetAddress =
-        event.getConnectionDescription().getServerAddress().getSocketAddress().getAddress();
-    if (inetAddress instanceof Inet4Address) {
-      Tags.PEER_HOST_IPV4.set(span, inetAddress.getHostAddress());
-    } 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, DDSpanTypes.MONGO);
-    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/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/MongoClientDecorator.java

@@ -0,0 +1,118 @@
+package datadog.trace.instrumentation.mongo;
+
+import com.mongodb.event.CommandStartedEvent;
+import datadog.trace.agent.decorator.DatabaseClientDecorator;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
+import io.opentracing.Span;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import org.bson.BsonArray;
+import org.bson.BsonDocument;
+import org.bson.BsonString;
+import org.bson.BsonValue;
+
+public class MongoClientDecorator extends DatabaseClientDecorator<CommandStartedEvent> {
+  public static final MongoClientDecorator DECORATE = new MongoClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"mongo"};
+  }
+
+  @Override
+  protected String service() {
+    return "mongo";
+  }
+
+  @Override
+  protected String component() {
+    return "java-mongo";
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.MONGO;
+  }
+
+  @Override
+  protected String dbType() {
+    return "mongo";
+  }
+
+  @Override
+  protected String dbUser(final CommandStartedEvent event) {
+    return null;
+  }
+
+  @Override
+  protected String dbInstance(final CommandStartedEvent event) {
+    return event.getDatabaseName();
+    // This would be the "proper" db.instance:
+    //    final ConnectionDescription connectionDescription = event.getConnectionDescription();
+    //    if (connectionDescription != null) {
+    //      final ConnectionId connectionId = connectionDescription.getConnectionId();
+    //      if (connectionId != null) {
+    //        final ServerId serverId = connectionId.getServerId();
+    //        if (serverId != null) {
+    //          return serverId.toString();
+    //        }
+    //      }
+    //    }
+    //    return null;
+  }
+
+  public Span onStatement(final Span span, final BsonDocument statement) {
+
+    // scrub the Mongo command so that parameters are removed from the string
+    final BsonDocument scrubbed = scrub(statement);
+    final String mongoCmd = scrubbed.toString();
+
+    span.setTag(DDTags.RESOURCE_NAME, mongoCmd);
+    return onStatement(span, mongoCmd);
+  }
+
+  /**
+   * 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 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/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/MongoClientInstrumentation.java

@@ -10,7 +10,6 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 import com.google.auto.service.AutoService;
 import com.mongodb.MongoClientOptions;
 import datadog.trace.agent.tooling.Instrumenter;
-import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Modifier;
 import java.util.Collections;
 import java.util.Map;
@@ -21,8 +20,6 @@ import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(Instrumenter.class)
 public final class MongoClientInstrumentation extends Instrumenter.Default {
-  public static final String[] HELPERS =
-      new String[] {"datadog.trace.instrumentation.mongo.DDTracingCommandListener"};
 
   public MongoClientInstrumentation() {
     super("mongo");
@@ -46,7 +43,13 @@ public final class MongoClientInstrumentation extends Instrumenter.Default {
 
   @Override
   public String[] helperClassNames() {
-    return HELPERS;
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.DatabaseClientDecorator",
+      packageName + ".MongoClientDecorator",
+      packageName + ".DDTracingCommandListener"
+    };
   }
 
   @Override
@@ -63,7 +66,7 @@ public final class MongoClientInstrumentation extends Instrumenter.Default {
       // referencing "this" in the method args causes the class to load under a transformer.
       // This bypasses the Builder instrumentation. Casting as a workaround.
       final MongoClientOptions.Builder builder = (MongoClientOptions.Builder) dis;
-      final DDTracingCommandListener listener = new DDTracingCommandListener(GlobalTracer.get());
+      final DDTracingCommandListener listener = new DDTracingCommandListener();
       builder.addCommandListener(listener);
     }
   }

dd-java-agent/instrumentation/mongo-3.1/src/test/groovy/MongoClientDecoratorTest.groovy

@@ -0,0 +1,45 @@
+import datadog.trace.api.DDTags
+import io.opentracing.Span
+import io.opentracing.tag.Tags
+import org.bson.BsonArray
+import org.bson.BsonDocument
+import org.bson.BsonString
+import spock.lang.Shared
+import spock.lang.Specification
+
+import static datadog.trace.instrumentation.mongo.MongoClientDecorator.DECORATE
+
+class MongoClientDecoratorTest extends Specification {
+
+  @Shared
+  def query1, query2
+
+  def setupSpec() {
+    query1 = new BsonDocument("find", new BsonString("show"))
+    query1.put("stuff", new BsonString("secret"))
+
+
+    query2 = new BsonDocument("insert", new BsonString("table"))
+    def nestedDoc = new BsonDocument("count", new BsonString("show"))
+    nestedDoc.put("id", new BsonString("secret"))
+    query2.put("docs", new BsonArray(Arrays.asList(new BsonString("secret"), nestedDoc)))
+  }
+
+  def "test query scrubbing"() {
+    setup:
+    def span = Mock(Span)
+    // all "secret" strings should be scrubbed out of these queries
+
+    when:
+    DECORATE.onStatement(span, query)
+
+    then:
+    1 * span.setTag(Tags.DB_STATEMENT.key, expected)
+    1 * span.setTag(DDTags.RESOURCE_NAME, expected)
+    0 * _
+
+    where:
+    query << [query1, query2]
+    expected = query.toString().replaceAll("secret", "?")
+  }
+}

dd-java-agent/instrumentation/mongo-3.1/src/test/java/datadog/trace/instrumentation/mongo/MongoClientInstrumentationTest.java

@@ -30,7 +30,8 @@ public class MongoClientInstrumentationTest {
         new CommandStartedEvent(1, makeConnection(), "databasename", "query", new BsonDocument());
 
     final DDSpan span = new DDTracer().buildSpan("foo").start();
-    DDTracingCommandListener.decorate(span, cmd);
+    MongoClientDecorator.DECORATE.afterStart(span);
+    MongoClientDecorator.DECORATE.onStatement(span, cmd.getCommand());
 
     assertThat(span.context().getSpanType()).isEqualTo("mongodb");
     assertThat(span.context().getResourceName())
@@ -53,7 +54,8 @@ public class MongoClientInstrumentationTest {
           new CommandStartedEvent(1, makeConnection(), "databasename", "query", query);
 
       final DDSpan span = new DDTracer().buildSpan("foo").start();
-      DDTracingCommandListener.decorate(span, cmd);
+      MongoClientDecorator.DECORATE.afterStart(span);
+      MongoClientDecorator.DECORATE.onStatement(span, cmd.getCommand());
 
       assertThat(span.getSpanType()).isEqualTo(DDSpanTypes.MONGO);
       assertThat(span.getTags().get(Tags.DB_STATEMENT.getKey()))

dd-java-agent/instrumentation/mongo-async-3.3/src/main/java/datadog/trace/instrumentation/mongo/MongoAsyncClientInstrumentation.java

@@ -10,7 +10,6 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 import com.google.auto.service.AutoService;
 import com.mongodb.async.client.MongoClientSettings;
 import datadog.trace.agent.tooling.Instrumenter;
-import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Modifier;
 import java.util.Collections;
 import java.util.Map;
@@ -44,7 +43,13 @@ public final class MongoAsyncClientInstrumentation extends Instrumenter.Default
 
   @Override
   public String[] helperClassNames() {
-    return MongoClientInstrumentation.HELPERS;
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.DatabaseClientDecorator",
+      packageName + ".MongoClientDecorator",
+      packageName + ".DDTracingCommandListener"
+    };
   }
 
   @Override
@@ -61,7 +66,7 @@ public final class MongoAsyncClientInstrumentation extends Instrumenter.Default
       // 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());
+      final DDTracingCommandListener listener = new DDTracingCommandListener();
       builder.addCommandListener(listener);
     }
   }

dd-java-agent/instrumentation/ratpack-1.4/src/test/groovy/RatpackTest.groovy

@@ -17,7 +17,9 @@ import ratpack.http.HttpUrlBuilder
 import ratpack.http.client.HttpClient
 import ratpack.path.PathBinding
 import ratpack.test.exec.ExecHarness
+import spock.lang.Retry
 
+@Retry
 class RatpackTest extends AgentTestRunner {
   static {
     System.setProperty("dd.integration.ratpack.enabled", "true")

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy

@@ -18,11 +18,17 @@ class SpanAssert {
                          @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.SpanAssert'])
                          @DelegatesTo(value = SpanAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     def asserter = new SpanAssert(span)
+    asserter.assertSpan spec
+  }
+
+  void assertSpan(
+    @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.SpanAssert'])
+    @DelegatesTo(value = SpanAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     def clone = (Closure) spec.clone()
-    clone.delegate = asserter
+    clone.delegate = this
     clone.resolveStrategy = Closure.DELEGATE_FIRST
-    clone(asserter)
+    clone(this)
-    asserter.assertDefaults()
+    assertDefaults()
   }
 
   def assertSpanNameContains(String spanName, String... shouldContainArr) {

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy

@@ -5,6 +5,8 @@ import datadog.trace.api.Config
 import groovy.transform.stc.ClosureParams
 import groovy.transform.stc.SimpleType
 
+import java.util.regex.Pattern
+
 class TagsAssert {
   private final String spanParentId
   private final Map<String, Object> tags
@@ -65,7 +67,9 @@ class TagsAssert {
       return
     }
     assertedTags.add(name)
-    if (value instanceof Class) {
+    if (value instanceof Pattern) {
+      assert tags[name] =~ value
+    } else if (value instanceof Class) {
       assert ((Class) value).isInstance(tags[name])
     } else if (value instanceof Closure) {
       assert ((Closure) value).call(tags[name])

dd-trace-ot/dd-trace-ot.gradle

@@ -41,7 +41,7 @@ dependencies {
   testCompile project(':utils:gc-utils')
   testCompile group: 'org.assertj', name: 'assertj-core', version: '2.9.+'
   testCompile group: 'org.mockito', name: 'mockito-core', version: '2.19.0'
-  testCompile group: 'org.objenesis', name: 'objenesis', version: '2.6'
+  testCompile group: 'org.objenesis', name: 'objenesis', version: '2.6' // Last version to support Java7
   testCompile group: 'cglib', name: 'cglib-nodep', version: '3.2.5'
   testCompile 'com.github.stefanbirkner:system-rules:1.17.1'
 }