Migrate Mongo to Decorator
MODIFY
dd-java-agent/instrumentation/mongo-3.1/mongo-3.1.gradle
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: '+'
}

MODIFY
dd-java-agent/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/DDTracingCommandListener.java
dd-java-agent/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/DDTracingCommandListener.java
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.tag.Tags;
-import java.net.Inet4Address;
-import java.net.InetAddress;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import io.opentracing.util.GlobalTracer;
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(final Tracer tracer) {
- this.tracer = tracer;
- }
+ private final Map<Integer, Span> spanMap = new ConcurrentHashMap<>();
@Override
public void commandStarted(final CommandStartedEvent event) {
- final Span span = buildSpan(event);
- cache.put(event.getRequestId(), span);
+ final Span span = GlobalTracer.get().buildSpan("mongo.query").start();
+ 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);
- span.log(Collections.singletonMap(ERROR_OBJECT, event.getThrowable()));
+ DECORATE.onError(span, 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;
- }
}

ADD
dd-java-agent/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/MongoClientDecorator.java
+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;
+ }
+}

MODIFY
dd-java-agent/instrumentation/mongo-3.1/src/main/java/datadog/trace/instrumentation/mongo/MongoClientInstrumentation.java
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);
}
}

ADD
dd-java-agent/instrumentation/mongo-3.1/src/test/groovy/MongoClientDecoratorTest.groovy
+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", "?")
+ }
+}

MODIFY
dd-java-agent/instrumentation/mongo-3.1/src/test/java/datadog/trace/instrumentation/mongo/MongoClientInstrumentationTest.java
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()))

MODIFY
dd-java-agent/instrumentation/mongo-async-3.3/src/main/java/datadog/trace/instrumentation/mongo/MongoAsyncClientInstrumentation.java
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);
}
}

MODIFY
dd-java-agent/instrumentation/ratpack-1.4/src/test/groovy/RatpackTest.groovy
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")

MODIFY
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy
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)
- asserter.assertDefaults()
+ clone(this)
+ assertDefaults()
}
def assertSpanNameContains(String spanName, String... shouldContainArr) {

MODIFY
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy
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])

MODIFY
dd-trace-ot/dd-trace-ot.gradle
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'
}

Migrate Mongo to Decorator