Merge pull request #40191 from vkn/mongodb-client-opentelemetry

Add mongo commands to otel span attributes

extensions/mongodb-client/deployment/src/main/java/io/quarkus/mongodb/deployment/MongoClientBuildTimeConfig.java

@@ -24,6 +24,12 @@ public class MongoClientBuildTimeConfig {
     @ConfigItem(name = "force-default-clients")
     public boolean forceDefaultClients;
 
+    /**
+     * Whether or not tracing spans of driver commands are sent in case the quarkus-opentelemetry extension is present.
+     */
+    @ConfigItem(name = "tracing.enabled")
+    public boolean tracingEnabled;
+
     /**
      * Configuration for DevServices. DevServices allows Quarkus to automatically start MongoDB in dev and test mode.
      */

extensions/mongodb-client/deployment/src/main/java/io/quarkus/mongodb/deployment/MongoClientProcessor.java

@@ -73,6 +73,7 @@
 import io.quarkus.mongodb.runtime.MongodbConfig;
 import io.quarkus.mongodb.runtime.dns.MongoDnsClient;
 import io.quarkus.mongodb.runtime.dns.MongoDnsClientProvider;
+import io.quarkus.mongodb.tracing.MongoTracingCommandListener;
 import io.quarkus.runtime.metrics.MetricsFactory;
 import io.quarkus.smallrye.health.deployment.spi.HealthBuildItem;
 import io.quarkus.vertx.deployment.VertxBuildItem;
@@ -111,6 +112,14 @@ AdditionalIndexedClassesBuildItem includeDnsTypesToIndex() {
                 MongoDnsClient.class.getName());
     }
 
+    @BuildStep
+    AdditionalIndexedClassesBuildItem includeDnsTypesToIndex(MongoClientBuildTimeConfig buildTimeConfig) {
+        if (buildTimeConfig.tracingEnabled) {
+            return new AdditionalIndexedClassesBuildItem(MongoTracingCommandListener.class.getName());
+        }
+        return new AdditionalIndexedClassesBuildItem();
+    }
+
     @BuildStep
     public void registerDnsProvider(BuildProducer<NativeImageResourceBuildItem> nativeProducer) {
         nativeProducer.produce(new NativeImageResourceBuildItem("META-INF/services/" + DnsClientProvider.class.getName()));
@@ -143,8 +152,7 @@ BsonDiscriminatorBuildItem collectBsonDiscriminators(CombinedIndexBuildItem inde
     }
 
     @BuildStep
-    CommandListenerBuildItem collectCommandListeners(CombinedIndexBuildItem indexBuildItem,
-            MongoClientBuildTimeConfig buildTimeConfig, Capabilities capabilities) {
+    CommandListenerBuildItem collectCommandListeners(CombinedIndexBuildItem indexBuildItem) {
         Collection<ClassInfo> commandListenerClasses = indexBuildItem.getIndex()
                 .getAllKnownImplementors(DotName.createSimple(CommandListener.class.getName()));
         List<String> names = commandListenerClasses.stream()

extensions/mongodb-client/runtime/pom.xml

@@ -36,6 +36,11 @@
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-mutiny-reactive-streams-operators</artifactId>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-opentelemetry</artifactId>
+            <optional>true</optional>
+        </dependency>
         <dependency>
             <groupId>org.mongodb</groupId>
             <artifactId>mongodb-driver-sync</artifactId>

extensions/mongodb-client/runtime/src/main/java/io/quarkus/mongodb/tracing/MongoTracingCommandListener.java

@@ -0,0 +1,92 @@
+package io.quarkus.mongodb.tracing;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import javax.annotation.Nullable;
+
+import jakarta.inject.Inject;
+
+import org.jboss.logging.Logger;
+
+import com.mongodb.event.*;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+
+public class MongoTracingCommandListener implements CommandListener {
+    private static final org.jboss.logging.Logger LOGGER = Logger.getLogger(MongoTracingCommandListener.class);
+    private static final String KEY = "mongodb.command";
+    private final Map<Integer, ContextEvent> requestMap;
+    private final Instrumenter<CommandStartedEvent, Void> instrumenter;
+
+    private record ContextEvent(Context context, CommandStartedEvent commandEvent) {
+    }
+
+    @Inject
+    public MongoTracingCommandListener(OpenTelemetry openTelemetry) {
+        requestMap = new ConcurrentHashMap<>();
+        SpanNameExtractor<CommandStartedEvent> spanNameExtractor = CommandEvent::getCommandName;
+        instrumenter = Instrumenter.<CommandStartedEvent, Void> builder(
+                openTelemetry, "quarkus-mongodb-client", spanNameExtractor)
+                .addAttributesExtractor(new CommandEventAttrExtractor())
+                .buildInstrumenter(SpanKindExtractor.alwaysClient());
+        LOGGER.debugf("MongoTracingCommandListener created");
+    }
+
+    @Override
+    public void commandStarted(CommandStartedEvent event) {
+        LOGGER.tracef("commandStarted event %s", event.getCommandName());
+
+        Context parentContext = Context.current();
+        if (instrumenter.shouldStart(parentContext, event)) {
+            Context context = instrumenter.start(parentContext, event);
+            requestMap.put(event.getRequestId(), new ContextEvent(context, event));
+        }
+    }
+
+    @Override
+    public void commandSucceeded(CommandSucceededEvent event) {
+        LOGGER.tracef("commandSucceeded event %s", event.getCommandName());
+        ContextEvent contextEvent = requestMap.remove(event.getRequestId());
+        if (contextEvent != null) {
+            instrumenter.end(contextEvent.context(), contextEvent.commandEvent(), null, null);
+        }
+    }
+
+    @Override
+    public void commandFailed(CommandFailedEvent event) {
+        LOGGER.tracef("commandFailed event %s", event.getCommandName());
+        ContextEvent contextEvent = requestMap.remove(event.getRequestId());
+        if (contextEvent != null) {
+            instrumenter.end(
+                    contextEvent.context(),
+                    contextEvent.commandEvent(),
+                    null,
+                    event.getThrowable());
+        }
+    }
+
+    private static class CommandEventAttrExtractor implements AttributesExtractor<CommandStartedEvent, Void> {
+        @Override
+        public void onStart(AttributesBuilder attributesBuilder,
+                Context context,
+                CommandStartedEvent commandStartedEvent) {
+            attributesBuilder.put(KEY, commandStartedEvent.getCommand().toJson());
+        }
+
+        @Override
+        public void onEnd(AttributesBuilder attributesBuilder,
+                Context context,
+                CommandStartedEvent commandStartedEvent,
+                @Nullable Void unused,
+                @Nullable Throwable throwable) {
+
+        }
+    }
+}

extensions/mongodb-client/runtime/src/test/java/io/quarkus/mongodb/tracing/MongoTracingCommandListenerTest.java

@@ -0,0 +1,103 @@
+package io.quarkus.mongodb.tracing;
+
+import static org.assertj.core.api.Assertions.assertThatNoException;
+
+import org.bson.BsonDocument;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import com.mongodb.ServerAddress;
+import com.mongodb.connection.ClusterId;
+import com.mongodb.connection.ConnectionDescription;
+import com.mongodb.connection.ServerId;
+import com.mongodb.event.CommandFailedEvent;
+import com.mongodb.event.CommandStartedEvent;
+import com.mongodb.event.CommandSucceededEvent;
+
+import io.opentelemetry.api.OpenTelemetry;
+
+class MongoTracingCommandListenerTest {
+    private ConnectionDescription connDescr;
+    private MongoTracingCommandListener listener;
+    private BsonDocument command;
+
+    @BeforeEach
+    void setUp() {
+        connDescr = new ConnectionDescription(new ServerId(new ClusterId(), new ServerAddress()));
+        listener = new MongoTracingCommandListener(OpenTelemetry.noop());
+        command = new BsonDocument();
+    }
+
+    @Test
+    void commandStarted() {
+        var startEvent = new CommandStartedEvent(
+                null,
+                1L,
+                10,
+                connDescr,
+                "db",
+                "find",
+                command);
+        assertThatNoException().isThrownBy(() -> listener.commandStarted(startEvent));
+
+        CommandSucceededEvent successEvent = new CommandSucceededEvent(null,
+                startEvent.getOperationId(),
+                startEvent.getRequestId(),
+                connDescr,
+                startEvent.getDatabaseName(),
+                startEvent.getCommandName(),
+                startEvent.getCommand(),
+                10L);
+        assertThatNoException().isThrownBy(() -> listener.commandSucceeded(successEvent));
+    }
+
+    @Test
+    void commandSucceeded() {
+        CommandSucceededEvent cmd = new CommandSucceededEvent(null,
+                1L,
+                10,
+                connDescr,
+                "db",
+                "find",
+                command,
+                10L);
+        assertThatNoException().isThrownBy(() -> listener.commandSucceeded(cmd));
+    }
+
+    @Test
+    void commandFailed() {
+        var startedEvent = new CommandStartedEvent(
+                null,
+                1L,
+                10,
+                connDescr,
+                "db",
+                "find",
+                command);
+        assertThatNoException().isThrownBy(() -> listener.commandStarted(startedEvent));
+
+        CommandFailedEvent failedEvent = new CommandFailedEvent(null,
+                1L,
+                10,
+                connDescr,
+                "db",
+                "find",
+                10L,
+                new IllegalStateException("command failed"));
+        assertThatNoException().isThrownBy(() -> listener.commandFailed(failedEvent));
+    }
+
+    @Test
+    void commandFailedNoEvent() {
+        CommandFailedEvent cmd = new CommandFailedEvent(null,
+                1L,
+                10,
+                connDescr,
+                "db",
+                "find",
+                10L,
+                new IllegalStateException("command failed"));
+        assertThatNoException().isThrownBy(() -> listener.commandFailed(cmd));
+    }
+
+}