Merge pull request #40714 from vkn/mongodb-otel-enhancements

Enhance MongoDb otel integration

.github/native-tests.json

@@ -117,7 +117,7 @@
         {
             "category": "Misc4",
             "timeout": 130,
-            "test-modules": "picocli-native, gradle, micrometer-mp-metrics, micrometer-prometheus, logging-json, jaxp, jaxb, opentelemetry, opentelemetry-jdbc-instrumentation, opentelemetry-redis-instrumentation, web-dependency-locator",
+            "test-modules": "picocli-native, gradle, micrometer-mp-metrics, micrometer-prometheus, logging-json, jaxp, jaxb, opentelemetry, opentelemetry-jdbc-instrumentation, opentelemetry-mongodb-client-instrumentation, opentelemetry-redis-instrumentation, web-dependency-locator",
             "os-name": "ubuntu-latest"
         },
         {

docs/src/main/asciidoc/mongodb.adoc

@@ -616,6 +616,18 @@ This behavior must first be enabled by setting the `quarkus.mongodb.metrics.enab
 So when you access the `/q/metrics` endpoint of your application you will have information about the connection pool status.
 When using xref:smallrye-metrics.adoc[SmallRye Metrics], connection pool metrics will be available under the `vendor` scope.
 
+
+== Tracing
+
+To use tracing with MongoDB, you need to add the xref:opentelemetry.adoc[`quarkus-opentelemetry`] extension to your project.
+
+Even with all the tracing infrastructure in place the mongodb tracing is not enabled by default, and you need to enable it by setting this property:
+[source, properties]
+----
+# enable tracing
+quarkus.mongodb.tracing.enabled=true
+----
+
 == Testing helpers
 
 xref:#dev-services[Dev Services for MongoDB] is your best option to start a MongoDB database for your unit tests.

docs/src/main/asciidoc/opentelemetry.adoc

@@ -684,6 +684,7 @@ Additional exporters will be available in the Quarkiverse https://docs.quarkiver
 * https://quarkus.io/guides/resteasy-client[`quarkus-resteasy-client`]
 * https://quarkus.io/guides/scheduler[`quarkus-scheduler`]
 * https://quarkus.io/guides/smallrye-graphql[`quarkus-smallrye-graphql`]
+* https://quarkus.io/extensions/io.quarkus/quarkus-mongodb-client[`quarkus-mongodb-client`]
 * https://quarkus.io/extensions/io.quarkus/quarkus-messaging[`quarkus-messaging`]
 ** AMQP 1.0
 ** RabbitMQ

extensions/mongodb-client/deployment/pom.xml

@@ -34,6 +34,11 @@
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-mongodb-client</artifactId>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-opentelemetry-deployment</artifactId>
+            <optional>true</optional>
+        </dependency>
         <dependency>
             <groupId>org.testcontainers</groupId>
             <artifactId>mongodb</artifactId>

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

@@ -0,0 +1,23 @@
+package io.quarkus.mongodb.deployment;
+
+import java.util.List;
+
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
+
+import io.quarkus.builder.item.SimpleBuildItem;
+
+/**
+ * Register additional {@link ReactiveContextProvider}s for the MongoDB clients.
+ */
+public final class ContextProviderBuildItem extends SimpleBuildItem {
+
+    private final List<String> classNames;
+
+    public ContextProviderBuildItem(List<String> classNames) {
+        this.classNames = classNames == null ? List.of() : classNames;
+    }
+
+    public List<String> getContextProviderClassNames() {
+        return classNames;
+    }
+}

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

@@ -32,6 +32,7 @@
 import com.mongodb.client.model.changestream.UpdateDescription;
 import com.mongodb.event.CommandListener;
 import com.mongodb.event.ConnectionPoolListener;
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
 import com.mongodb.spi.dns.DnsClientProvider;
 
 import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
@@ -69,6 +70,7 @@
 import io.quarkus.mongodb.runtime.MongoClientRecorder;
 import io.quarkus.mongodb.runtime.MongoClientSupport;
 import io.quarkus.mongodb.runtime.MongoClients;
+import io.quarkus.mongodb.runtime.MongoReactiveContextProvider;
 import io.quarkus.mongodb.runtime.MongoServiceBindingConverter;
 import io.quarkus.mongodb.runtime.MongodbConfig;
 import io.quarkus.mongodb.runtime.dns.MongoDnsClient;
@@ -113,9 +115,11 @@ AdditionalIndexedClassesBuildItem includeDnsTypesToIndex() {
     }
 
     @BuildStep
-    AdditionalIndexedClassesBuildItem includeDnsTypesToIndex(MongoClientBuildTimeConfig buildTimeConfig) {
+    AdditionalIndexedClassesBuildItem includeMongoCommandListener(MongoClientBuildTimeConfig buildTimeConfig) {
         if (buildTimeConfig.tracingEnabled) {
-            return new AdditionalIndexedClassesBuildItem(MongoTracingCommandListener.class.getName());
+            return new AdditionalIndexedClassesBuildItem(
+                    MongoTracingCommandListener.class.getName(),
+                    MongoReactiveContextProvider.class.getName());
         }
         return new AdditionalIndexedClassesBuildItem();
     }
@@ -161,15 +165,27 @@ CommandListenerBuildItem collectCommandListeners(CombinedIndexBuildItem indexBui
         return new CommandListenerBuildItem(names);
     }
 
+    @BuildStep
+    ContextProviderBuildItem collectContextProviders(CombinedIndexBuildItem indexBuildItem) {
+        Collection<ClassInfo> contextProviders = indexBuildItem.getIndex()
+                .getAllKnownImplementors(DotName.createSimple(ReactiveContextProvider.class.getName()));
+        List<String> names = contextProviders.stream()
+                .map(ci -> ci.name().toString())
+                .collect(Collectors.toList());
+        return new ContextProviderBuildItem(names);
+    }
+
     @BuildStep
     List<ReflectiveClassBuildItem> addExtensionPointsToNative(CodecProviderBuildItem codecProviders,
             PropertyCodecProviderBuildItem propertyCodecProviders, BsonDiscriminatorBuildItem bsonDiscriminators,
-            CommandListenerBuildItem commandListeners) {
+            CommandListenerBuildItem commandListeners,
+            ContextProviderBuildItem contextProviders) {
         List<String> reflectiveClassNames = new ArrayList<>();
         reflectiveClassNames.addAll(codecProviders.getCodecProviderClassNames());
         reflectiveClassNames.addAll(propertyCodecProviders.getPropertyCodecProviderClassNames());
         reflectiveClassNames.addAll(bsonDiscriminators.getBsonDiscriminatorClassNames());
         reflectiveClassNames.addAll(commandListeners.getCommandListenerClassNames());
+        reflectiveClassNames.addAll(contextProviders.getContextProviderClassNames());
 
         List<ReflectiveClassBuildItem> reflectiveClass = reflectiveClassNames.stream()
                 .map(s -> ReflectiveClassBuildItem.builder(s).methods().build())
@@ -256,6 +272,7 @@ void build(
             PropertyCodecProviderBuildItem propertyCodecProvider,
             BsonDiscriminatorBuildItem bsonDiscriminator,
             CommandListenerBuildItem commandListener,
+            ContextProviderBuildItem contextProvider,
             List<MongoConnectionPoolListenerBuildItem> connectionPoolListenerProvider,
             BuildProducer<AdditionalBeanBuildItem> additionalBeanBuildItemProducer,
             BuildProducer<SyntheticBeanBuildItem> syntheticBeanBuildItemBuildProducer) {
@@ -277,6 +294,9 @@ void build(
         for (String name : commandListener.getCommandListenerClassNames()) {
             additionalBeansBuilder.addBeanClass(name);
         }
+        for (String name : contextProvider.getContextProviderClassNames()) {
+            additionalBeansBuilder.addBeanClass(name);
+        }
         additionalBeanBuildItemProducer.produce(additionalBeansBuilder.build());
 
         // create MongoClientSupport as a synthetic bean as it's used in AbstractMongoClientProducer

extensions/mongodb-client/deployment/src/test/java/io/quarkus/mongodb/MockReactiveContextProvider.java

@@ -0,0 +1,22 @@
+package io.quarkus.mongodb;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.reactivestreams.Subscriber;
+
+import com.mongodb.RequestContext;
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
+
+import io.quarkus.mongodb.runtime.MongoRequestContext;
+
+public class MockReactiveContextProvider implements ReactiveContextProvider {
+
+    public static final List<String> EVENTS = new ArrayList<>();
+
+    @Override
+    public RequestContext getContext(Subscriber<?> subscriber) {
+        EVENTS.add(MongoRequestContext.class.getName());
+        return new MongoRequestContext(null);
+    }
+}

extensions/mongodb-client/deployment/src/test/java/io/quarkus/mongodb/MongoTracingEnabled.java

@@ -0,0 +1,49 @@
+package io.quarkus.mongodb;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+import java.time.Duration;
+
+import jakarta.inject.Inject;
+
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.mongodb.reactive.ReactiveMongoClient;
+import io.quarkus.test.QuarkusUnitTest;
+
+public class MongoTracingEnabled extends MongoTestBase {
+
+    @Inject
+    ReactiveMongoClient reactiveClient;
+
+    @RegisterExtension
+    static final QuarkusUnitTest config = new QuarkusUnitTest().setArchiveProducer(
+            () -> ShrinkWrap.create(JavaArchive.class)
+                    .addClasses(MongoTestBase.class, MockReactiveContextProvider.class, MockCommandListener.class))
+            .withConfigurationResource("application-tracing-mongoclient.properties");
+
+    @AfterEach
+    void cleanup() {
+        if (reactiveClient != null) {
+            reactiveClient.close();
+        }
+    }
+
+    @Test
+    void invokeReactiveContextProvider() {
+        String dbNames = reactiveClient.listDatabaseNames().toUni().await().atMost(Duration.ofSeconds(30L));
+        assertThat(dbNames).as("expect db names available").isNotBlank();
+        await().atMost(Duration.ofSeconds(30L)).untilAsserted(
+                () -> assertThat(MockReactiveContextProvider.EVENTS)
+                        .as("reactive context provider must be called")
+                        .isNotEmpty());
+        assertThat(MockCommandListener.EVENTS).isNotEmpty();
+
+    }
+
+}

extensions/mongodb-client/deployment/src/test/java/io/quarkus/mongodb/MongoTracingNotEnabledTest.java

@@ -0,0 +1,42 @@
+package io.quarkus.mongodb;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import jakarta.inject.Inject;
+
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import com.mongodb.client.MongoClient;
+
+import io.quarkus.test.QuarkusUnitTest;
+
+public class MongoTracingNotEnabledTest extends MongoTestBase {
+
+    @Inject
+    MongoClient client;
+
+    @RegisterExtension
+    static final QuarkusUnitTest config = new QuarkusUnitTest()
+            .setArchiveProducer(
+                    () -> ShrinkWrap.create(JavaArchive.class).addClasses(MongoTestBase.class,
+                            MockReactiveContextProvider.class))
+            .withConfigurationResource("default-mongoclient.properties");
+
+    @AfterEach
+    void cleanup() {
+        if (client != null) {
+            client.close();
+        }
+    }
+
+    @Test
+    void contextProviderMustNotBeCalledIfNoOpenTelemetryIsAvailable() {
+        assertThat(client.listDatabaseNames().first()).isNotEmpty();
+        assertThat(MockReactiveContextProvider.EVENTS).isEmpty();
+    }
+
+}

extensions/mongodb-client/deployment/src/test/java/io/quarkus/mongodb/deployment/ContextProviderBuildItemTest.java

@@ -0,0 +1,28 @@
+package io.quarkus.mongodb.deployment;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+class ContextProviderBuildItemTest {
+
+    @Test
+    void getContextProviderClassNames() {
+        ContextProviderBuildItem item = new ContextProviderBuildItem(List.of("foo.bar"));
+        assertThat(item.getContextProviderClassNames())
+                .hasSize(1)
+                .first()
+                .isEqualTo("foo.bar");
+    }
+
+    @Test
+    void emptyOrNull() {
+        ContextProviderBuildItem withNull = new ContextProviderBuildItem(null);
+        assertThat(withNull.getContextProviderClassNames()).isEmpty();
+
+        ContextProviderBuildItem empty = new ContextProviderBuildItem(List.of());
+        assertThat(empty.getContextProviderClassNames()).isEmpty();
+    }
+}

extensions/mongodb-client/deployment/src/test/resources/application-tracing-mongoclient.properties

@@ -0,0 +1,3 @@
+quarkus.mongodb.connection-string=mongodb://127.0.0.1:27018
+quarkus.mongodb.tracing.enabled=true
+

extensions/mongodb-client/runtime/pom.xml

@@ -58,6 +58,11 @@
             <artifactId>quarkus-kubernetes-service-binding</artifactId>
             <optional>true</optional>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-opentelemetry</artifactId>
+            <optional>true</optional>
+        </dependency>
 
         <!-- Add the health extension as optional as we will produce the health check only if it's included -->
         <dependency>

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

@@ -56,6 +56,7 @@
 import com.mongodb.connection.SslSettings;
 import com.mongodb.event.CommandListener;
 import com.mongodb.event.ConnectionPoolListener;
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
 
 import io.quarkus.arc.Arc;
 import io.quarkus.arc.InstanceHandle;
@@ -86,18 +87,21 @@ public class MongoClients {
 
     private final Map<String, MongoClient> mongoclients = new HashMap<>();
     private final Map<String, ReactiveMongoClient> reactiveMongoClients = new HashMap<>();
+    private final Instance<ReactiveContextProvider> reactiveContextProviders;
     private final Instance<MongoClientCustomizer> customizers;
 
     public MongoClients(MongodbConfig mongodbConfig, MongoClientSupport mongoClientSupport,
             Instance<CodecProvider> codecProviders,
             Instance<PropertyCodecProvider> propertyCodecProviders,
             Instance<CommandListener> commandListeners,
+            Instance<ReactiveContextProvider> reactiveContextProviders,
             @Any Instance<MongoClientCustomizer> customizers) {
         this.mongodbConfig = mongodbConfig;
         this.mongoClientSupport = mongoClientSupport;
         this.codecProviders = codecProviders;
         this.propertyCodecProviders = propertyCodecProviders;
         this.commandListeners = commandListeners;
+        this.reactiveContextProviders = reactiveContextProviders;
         this.customizers = customizers;
 
         try {
@@ -121,15 +125,17 @@ public MongoClients(MongodbConfig mongodbConfig, MongoClientSupport mongoClientS
     }
 
     public MongoClient createMongoClient(String clientName) throws MongoException {
-        MongoClientSettings mongoConfiguration = createMongoConfiguration(clientName, getMatchingMongoClientConfig(clientName));
+        MongoClientSettings mongoConfiguration = createMongoConfiguration(clientName, getMatchingMongoClientConfig(clientName),
+                false);
         MongoClient client = com.mongodb.client.MongoClients.create(mongoConfiguration);
         mongoclients.put(clientName, client);
         return client;
     }
 
     public ReactiveMongoClient createReactiveMongoClient(String clientName)
             throws MongoException {
-        MongoClientSettings mongoConfiguration = createMongoConfiguration(clientName, getMatchingMongoClientConfig(clientName));
+        MongoClientSettings mongoConfiguration = createMongoConfiguration(clientName, getMatchingMongoClientConfig(clientName),
+                true);
         com.mongodb.reactivestreams.client.MongoClient client = com.mongodb.reactivestreams.client.MongoClients
                 .create(mongoConfiguration);
         ReactiveMongoClientImpl reactive = new ReactiveMongoClientImpl(client);
@@ -254,14 +260,18 @@ public void apply(ServerSettings.Builder builder) {
         }
     }
 
-    private MongoClientSettings createMongoConfiguration(String name, MongoClientConfig config) {
+    private MongoClientSettings createMongoConfiguration(String name, MongoClientConfig config, boolean isReactive) {
         if (config == null) {
             throw new RuntimeException("mongo config is missing for creating mongo client.");
         }
         CodecRegistry defaultCodecRegistry = MongoClientSettings.getDefaultCodecRegistry();
 
         MongoClientSettings.Builder settings = MongoClientSettings.builder();
 
+        if (isReactive) {
+            reactiveContextProviders.stream().findAny().ifPresent(settings::contextProvider);
+        }
+
         ConnectionString connectionString;
         Optional<String> maybeConnectionString = config.connectionString;
         if (maybeConnectionString.isPresent()) {

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

@@ -0,0 +1,17 @@
+package io.quarkus.mongodb.runtime;
+
+import org.reactivestreams.Subscriber;
+
+import com.mongodb.RequestContext;
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
+
+import io.opentelemetry.context.Context;
+
+public class MongoReactiveContextProvider implements ReactiveContextProvider {
+
+    @Override
+    public RequestContext getContext(Subscriber<?> subscriber) {
+        return new MongoRequestContext(Context.current());
+    }
+
+}