Replace TracingCommandListener with DDTracingCommandListener

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

@@ -19,7 +19,7 @@ public class MongoClientInstrumentationTest {
                 .get(0)
                 .getClass()
                 .getSimpleName())
-        .isEqualTo("TracingCommandListener");
+        .isEqualTo("DDTracingCommandListener");
 
     mongoClient.close();
   }

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

@@ -2,14 +2,21 @@ 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.contrib.mongo.TracingCommandListener;
-import io.opentracing.contrib.mongo.TracingCommandListenerFactory;
+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;
@@ -21,10 +28,6 @@ import org.jboss.byteman.rule.Rule;
 @Slf4j
 public class MongoHelper extends DDAgentTracingHelper<MongoClientOptions.Builder> {
 
-  private static final List<String> WHILDCARD_FIELDS =
-      Arrays.asList("ordered", "insert", "count", "find");
-  private static final BsonValue HIDDEN_CAR = new BsonString("?");
-
   public MongoHelper(final Rule rule) {
     super(rule);
   }
@@ -42,7 +45,7 @@ public class MongoHelper extends DDAgentTracingHelper<MongoClientOptions.Builder
   protected MongoClientOptions.Builder doPatch(final MongoClientOptions.Builder builder)
       throws Exception {
 
-    final TracingCommandListener listener = TracingCommandListenerFactory.create(tracer);
+    final DDTracingCommandListener listener = new DDTracingCommandListener(tracer);
     builder.addCommandListener(listener);
 
     setState(builder, 1);
@@ -50,54 +53,129 @@ public class MongoHelper extends DDAgentTracingHelper<MongoClientOptions.Builder
     return builder;
   }
 
-  public void decorate(final Span span, final CommandStartedEvent event) {
-    try {
-      // normalize the Mongo command so that parameters are removed from the string
-      final BsonDocument normalized = norm(event.getCommand());
-      final String mongoCmd = normalized.toString();
+  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");
 
-      // add specific resource name and replace the `db.statement` OpenTracing
-      // tag with the quantized version of the Mongo command
-      span.setTag(DDTags.RESOURCE_NAME, mongoCmd);
-      span.setTag(Tags.DB_STATEMENT.getKey(), mongoCmd);
-      span.setTag(DDTags.SPAN_TYPE, "mongodb");
-    } catch (final Throwable e) {
-      log.warn("Couldn't decorate the mongo query: " + e.getMessage(), e);
+    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;
     }
-  }
 
-  private BsonDocument norm(final BsonDocument origin) {
-    final BsonDocument normalized = new BsonDocument();
-    for (final Map.Entry<String, BsonValue> entry : origin.entrySet()) {
-      if (WHILDCARD_FIELDS.contains(entry.getKey())) {
-        normalized.put(entry.getKey(), entry.getValue());
-      } else {
-        final BsonValue child = norm(entry.getValue());
-        normalized.put(entry.getKey(), child);
+    @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();
       }
     }
-    return normalized;
-  }
 
-  private BsonValue norm(final BsonArray origin) {
-    final BsonArray normalized = new BsonArray();
-    for (final BsonValue value : origin) {
-      final BsonValue child = norm(value);
-      normalized.add(child);
+    @Override
+    public void commandFailed(CommandFailedEvent event) {
+      Span span = cache.remove(event.getRequestId());
+      if (span != null) {
+        onError(span, event.getThrowable());
+        span.finish();
+      }
     }
-    return normalized;
-  }
 
-  private BsonValue norm(final BsonValue origin) {
+    private Span buildSpan(CommandStartedEvent event) {
+      Tracer.SpanBuilder spanBuilder =
+          tracer.buildSpan(MONGO_OPERATION).withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT);
 
-    final BsonValue normalized;
-    if (origin.isDocument()) {
-      normalized = norm(origin.asDocument());
-    } else if (origin.isArray()) {
-      normalized = norm(origin.asArray());
-    } else {
-      normalized = HIDDEN_CAR;
+      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;
     }
-    return normalized;
   }
 }

dd-java-agent/integrations/helpers/src/main/java/io/opentracing/contrib/mongo/TracingCommandListenerFactory.java

@@ -1,14 +0,0 @@
-package io.opentracing.contrib.mongo;
-
-import io.opentracing.Tracer;
-
-/**
- * This class exists purely to bypass the reduction in constructor visibility of
- * TracingCommandListener.
- */
-public class TracingCommandListenerFactory {
-
-  public static TracingCommandListener create(final Tracer tracer) {
-    return new TracingCommandListener(tracer);
-  }
-}

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

@@ -48,15 +48,6 @@ DO
   com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
 ENDRULE
 
-RULE TracingCommandListener-init
-CLASS io.opentracing.contrib.mongo.TracingCommandListener
-METHOD <init>
-AT EXIT
-IF TRUE
-DO
-  com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
-ENDRULE
-
 
 # Instrument OkHttp
 # ===========================

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

@@ -70,16 +70,6 @@ DO
   patch($0);
 ENDRULE
 
-RULE opentracing-mongo-driver-helper
-CLASS io.opentracing.contrib.mongo.TracingCommandListener
-METHOD decorate
-HELPER com.datadoghq.agent.integration.MongoHelper
-AT EXIT
-IF TRUE
-DO
-  decorate($1, $2);
-ENDRULE
-
 
 # Instrument OkHttp
 # ===========================