PR comments: build traceparent incl. sampling flag, exception logging

Author: Yoone

dd-java-agent/instrumentation/mongo/common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommentInjector.java

@@ -46,8 +46,7 @@ public static BsonDocument injectComment(String dbmComment, CommandStartedEvent
       log.warn(
           "Linking Database Monitoring profiles to spans is not supported for the following query type: {}. "
               + "To disable this feature please set the following environment variable: DD_DBM_PROPAGATION_MODE=disabled",
-          event.getCommand().getClass().getSimpleName(),
-          e);
+          event.getCommand().getClass().getSimpleName());
       return event.getCommand();
     }
   }
@@ -115,9 +114,12 @@ private static String getHostnameFromEvent(CommandStartedEvent event) {
 
   private static String buildTraceParent(AgentSpan span) {
     // W3C traceparent format: version-traceId-spanId-flags
-    String traceIdHex = span.getTraceId().toHexStringPadded(32);
-    String spanIdHex = String.format("%016x", span.getSpanId());
-    String flags = "00"; // '01' if sampled, '00' if not
-    return String.format("00-%s-%s-%s", traceIdHex, spanIdHex, flags);
+    StringBuilder sb = new StringBuilder(2 + 1 + 32 + 1 + 16 + 1 + 2);
+    sb.append("00-"); // version
+    sb.append(span.getTraceId().toHexStringPadded(32)); // traceId
+    sb.append("-");
+    sb.append(String.format("%016x", span.getSpanId())); // spanId
+    sb.append(span.context().getSamplingPriority() > 0 ? "-01" : "-00");
+    return sb.toString();
   }
 }

dd-java-agent/instrumentation/mongo/common/src/test/groovy/MongoCommentInjectorTest.groovy

@@ -2,6 +2,7 @@ package datadog.trace.instrumentation.mongo
 
 import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.api.config.TraceInstrumentationConfig
+import datadog.trace.api.sampling.PrioritySampling
 import org.bson.BsonDocument
 import org.bson.BsonString
 
@@ -25,6 +26,66 @@ class MongoCommentInjectorTest extends InstrumentationSpecification {
     comment == null
   }
 
+  def "getComment with full mode builds traceParent in W3C format with sampled flag"() {
+    setup:
+    // Set the propagation mode to full to trigger buildTraceParent call
+    injectSysConfig("dd." + TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "full")
+    def span = TEST_TRACER.buildSpan("test-op").start()
+    span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP, 0)
+
+    // Create a mock event
+    def event = Mock(com.mongodb.event.CommandStartedEvent) {
+      getDatabaseName() >> "testdb"
+      getConnectionDescription() >> Mock(com.mongodb.connection.ConnectionDescription) {
+        getServerAddress() >> Mock(com.mongodb.ServerAddress) {
+          getHost() >> "localhost"
+        }
+      }
+    }
+
+    when:
+    String comment = MongoCommentInjector.getComment(span, event)
+
+    then:
+    comment != null
+    comment.contains("dddbs='test-mongo-service'")
+    comment.contains("dde='test'")
+    comment.contains("ddh='localhost'")
+    comment.contains("dddb='testdb'")
+    // W3C traceparent format: 00-{32 hex chars}-{16 hex chars}-01 (sampled)
+    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-01'.*/
+  }
+
+  def "getComment with full mode builds traceParent in W3C format with not sampled flag"() {
+    setup:
+    // Set the propagation mode to full to trigger buildTraceParent call
+    injectSysConfig("dd." + TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "full")
+    def span = TEST_TRACER.buildSpan("test-op").start()
+    span.setSamplingPriority(PrioritySampling.SAMPLER_DROP, 0)
+
+    // Create a mock event
+    def event = Mock(com.mongodb.event.CommandStartedEvent) {
+      getDatabaseName() >> "testdb"
+      getConnectionDescription() >> Mock(com.mongodb.connection.ConnectionDescription) {
+        getServerAddress() >> Mock(com.mongodb.ServerAddress) {
+          getHost() >> "localhost"
+        }
+      }
+    }
+
+    when:
+    String comment = MongoCommentInjector.getComment(span, event)
+
+    then:
+    comment != null
+    comment.contains("dddbs='test-mongo-service'")
+    comment.contains("dde='test'")
+    comment.contains("ddh='localhost'")
+    comment.contains("dddb='testdb'")
+    // W3C traceparent format: 00-{32 hex chars}-{16 hex chars}-00 (not sampled)
+    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-00'.*/
+  }
+
   def "injectComment returns null when event is null"() {
     when:
     BsonDocument result = MongoCommentInjector.injectComment("test-comment", null)

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

@@ -8,15 +8,20 @@ import com.mongodb.event.CommandFailedEvent
 import com.mongodb.event.CommandListener
 import com.mongodb.event.CommandStartedEvent
 import com.mongodb.event.CommandSucceededEvent
+import datadog.trace.api.sampling.PrioritySampling
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.core.DDSpan
 import org.bson.BsonDocument
 import org.bson.BsonString
 import org.bson.Document
 import spock.lang.Shared
 
+import java.util.concurrent.ConcurrentLinkedQueue
+
 import static datadog.trace.agent.test.utils.PortUtils.UNUSABLE_PORT
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
+import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 
 abstract class MongoJava31ClientTest extends MongoBaseTest {
@@ -231,6 +236,132 @@ abstract class MongoJava31ClientTest extends MongoBaseTest {
     // Unfortunately not caught by our instrumentation.
     assertTraces(0) {}
   }
+
+  def "test DBM propagation with sampled flag"() {
+    setup:
+    injectSysConfig("dd." + DB_DBM_PROPAGATION_MODE_MODE, "full")
+    injectSysConfig("service.name", "test-mongo-service")
+    injectSysConfig("dd.env", "test")
+
+    String collectionName = randomCollectionName()
+    def capturedEvents = new ConcurrentLinkedQueue<CommandStartedEvent>()
+
+    def dbmClient = new MongoClient(new ServerAddress(mongoDbContainer.getHost(), port),
+      MongoClientOptions.builder()
+      .description("dbm-test-client")
+      .addCommandListener(new CommandListener() {
+        @Override
+        void commandStarted(CommandStartedEvent event) {
+          capturedEvents.add(event)
+        }
+
+        @Override
+        void commandSucceeded(CommandSucceededEvent event) {}
+
+        @Override
+        void commandFailed(CommandFailedEvent event) {}
+      })
+      .build())
+
+    DDSpan setupSpan = null
+    MongoCollection<Document> collection = runUnderTrace("setup") {
+      setupSpan = activeSpan() as DDSpan
+      MongoDatabase db = dbmClient.getDatabase(databaseName)
+      db.createCollection(collectionName)
+      return db.getCollection(collectionName)
+    }
+    TEST_WRITER.waitUntilReported(setupSpan)
+    TEST_WRITER.clear()
+    capturedEvents.clear()
+
+    when:
+    runUnderTrace("dbm-test") {
+      AgentSpan span = activeSpan() as AgentSpan
+      span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP, 0)
+      collection.insertOne(new Document("test", "value"))
+    }
+
+    then:
+    // Find the insert command event
+    def insertEvent = capturedEvents.find { it.commandName == "insert" }
+    insertEvent != null
+    def command = insertEvent.command
+    command.containsKey('$comment')
+    def comment = command.getString('$comment').value
+
+    // Verify comment contains DBM tags
+    comment.contains("dddbs='test-mongo-service'")
+    comment.contains("dde='test'")
+
+    // Verify traceparent has sampled flag (01)
+    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-01'.*/
+
+    cleanup:
+    dbmClient?.close()
+  }
+
+  def "test DBM propagation with not sampled flag"() {
+    setup:
+    injectSysConfig("dd." + DB_DBM_PROPAGATION_MODE_MODE, "full")
+    injectSysConfig("service.name", "test-mongo-service")
+    injectSysConfig("dd.env", "test")
+
+    String collectionName = randomCollectionName()
+    def capturedEvents = new ConcurrentLinkedQueue<CommandStartedEvent>()
+
+    def dbmClient = new MongoClient(new ServerAddress(mongoDbContainer.getHost(), port),
+      MongoClientOptions.builder()
+      .description("dbm-test-client")
+      .addCommandListener(new CommandListener() {
+        @Override
+        void commandStarted(CommandStartedEvent event) {
+          capturedEvents.add(event)
+        }
+
+        @Override
+        void commandSucceeded(CommandSucceededEvent event) {}
+
+        @Override
+        void commandFailed(CommandFailedEvent event) {}
+      })
+      .build())
+
+    DDSpan setupSpan = null
+    MongoCollection<Document> collection = runUnderTrace("setup") {
+      setupSpan = activeSpan() as DDSpan
+      MongoDatabase db = dbmClient.getDatabase(databaseName)
+      db.createCollection(collectionName)
+      return db.getCollection(collectionName)
+    }
+    TEST_WRITER.waitUntilReported(setupSpan)
+    TEST_WRITER.clear()
+    capturedEvents.clear()
+
+    when:
+    runUnderTrace("dbm-test") {
+      AgentSpan span = activeSpan() as AgentSpan
+      span.setSamplingPriority(PrioritySampling.SAMPLER_DROP, 0)
+      collection.insertOne(new Document("test", "value"))
+    }
+
+    then:
+    // Find the insert command event
+    def insertEvent = capturedEvents.find { it.commandName == "insert" }
+    insertEvent != null
+    def command = insertEvent.command
+    command.containsKey('$comment')
+    def comment = command.getString('$comment').value
+
+    // Verify comment contains DBM tags
+    comment.contains("dddbs='test-mongo-service'")
+    comment.contains("dde='test'")
+
+    // Verify traceparent has not sampled flag (00)
+    comment ==~ /.*ddtp='00-[0-9a-f]{32}-[0-9a-f]{16}-00'.*/
+
+    cleanup:
+    dbmClient?.close()
+  }
 }
 
 class MongoJava31ClientV0Test extends MongoJava31ClientTest {