Merge pull request #7099 from Postremus/mongo-connection-pool-metrics

Add Mongo Connection Pool Metrics

docs/src/main/asciidoc/mongodb.adoc

@@ -542,6 +542,14 @@ So when you access the `/health/ready` endpoint of your application you will hav
 
 This behavior can be disabled by setting the `quarkus.mongodb.health.enabled` property to `false` in your `application.properties`.
 
+== Metrics
+
+If you are using the `quarkus-smallrye-metrics` extension, `quarkus-mongodb` can provide metrics about the connection pools.
+This behavior must first be enabled by setting the `quarkus.mongodb.metrics.enabled` property to `true` in your `application.properties`.
+
+So when you access the `/metrics` endpoint of your application you will have information about the connection pool status in the `vendor` scope.
+
+
 == The legacy client
 
 We don't include the legacy MongoDB client by default. It contains the now retired MongoDB Java API (DB, DBCollection,... )

extensions/mongodb-client/deployment/pom.xml

@@ -38,6 +38,11 @@
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-mongodb-client</artifactId>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-smallrye-metrics</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>io.quarkus</groupId>
             <artifactId>quarkus-junit5-internal</artifactId>

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

@@ -11,4 +11,10 @@ public class MongoClientBuildTimeConfig {
      */
     @ConfigItem(name = "health.enabled", defaultValue = "true")
     public boolean healthEnabled;
+
+    /**
+     * Whether or not metrics are published in case the smallrye-metrics extension is present.
+     */
+    @ConfigItem(name = "metrics.enabled")
+    public boolean metricsEnabled;
 }

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

@@ -24,6 +24,7 @@
 import org.jboss.jandex.IndexView;
 
 import com.mongodb.client.MongoClient;
+import com.mongodb.event.ConnectionPoolListener;
 
 import io.quarkus.arc.Unremovable;
 import io.quarkus.arc.deployment.BeanContainerListenerBuildItem;
@@ -32,6 +33,7 @@
 import io.quarkus.arc.deployment.GeneratedBeanGizmoAdaptor;
 import io.quarkus.arc.deployment.UnremovableBeanBuildItem;
 import io.quarkus.arc.processor.DotNames;
+import io.quarkus.deployment.Capabilities;
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
 import io.quarkus.deployment.annotations.Record;
@@ -49,6 +51,7 @@
 import io.quarkus.gizmo.MethodDescriptor;
 import io.quarkus.gizmo.ResultHandle;
 import io.quarkus.mongodb.ReactiveMongoClient;
+import io.quarkus.mongodb.metrics.MongoMetricsConnectionPoolListener;
 import io.quarkus.mongodb.runtime.AbstractMongoClientProducer;
 import io.quarkus.mongodb.runtime.MongoClientConfig;
 import io.quarkus.mongodb.runtime.MongoClientName;
@@ -75,9 +78,16 @@ UnremovableBeanBuildItem markBeansAsUnremovable() {
     void configureRuntimeProperties(MongoClientRecorder recorder,
             CodecProviderBuildItem codecProvider,
             BsonDiscriminatorBuildItem bsonDiscriminator,
-            MongodbConfig config) {
+            MongodbConfig config,
+            List<MongoConnectionPoolListenerBuildItem> connectionPoolListenerProvider) {
+        List<ConnectionPoolListener> poolListenerList = connectionPoolListenerProvider.stream()
+                .map(MongoConnectionPoolListenerBuildItem::getConnectionPoolListener)
+                .collect(Collectors.toList());
+
         recorder.configureRuntimeProperties(codecProvider.getCodecProviderClassNames(),
-                bsonDiscriminator.getBsonDiscriminatorClassNames(), config);
+                bsonDiscriminator.getBsonDiscriminatorClassNames(),
+                config,
+                poolListenerList);
     }
 
     @BuildStep
@@ -392,4 +402,14 @@ HealthBuildItem addHealthCheck(MongoClientBuildTimeConfig buildTimeConfig) {
         return new HealthBuildItem("io.quarkus.mongodb.health.MongoHealthCheck",
                 buildTimeConfig.healthEnabled, "mongodb");
     }
+
+    @BuildStep
+    void setupMetrics(
+            MongoClientBuildTimeConfig buildTimeConfig, Capabilities capabilities,
+            BuildProducer<MongoConnectionPoolListenerBuildItem> producer) {
+
+        if (buildTimeConfig.metricsEnabled && capabilities.isCapabilityPresent(Capabilities.METRICS)) {
+            producer.produce(new MongoConnectionPoolListenerBuildItem(new MongoMetricsConnectionPoolListener()));
+        }
+    }
 }

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

@@ -0,0 +1,17 @@
+package io.quarkus.mongodb.deployment;
+
+import com.mongodb.event.ConnectionPoolListener;
+
+import io.quarkus.builder.item.MultiBuildItem;
+
+public final class MongoConnectionPoolListenerBuildItem extends MultiBuildItem {
+    private ConnectionPoolListener connectionPoolListener;
+
+    public MongoConnectionPoolListenerBuildItem(ConnectionPoolListener connectionPoolListener) {
+        this.connectionPoolListener = connectionPoolListener;
+    }
+
+    public ConnectionPoolListener getConnectionPoolListener() {
+        return connectionPoolListener;
+    }
+}

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

@@ -0,0 +1,77 @@
+package io.quarkus.mongodb;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import javax.inject.Inject;
+
+import org.eclipse.microprofile.metrics.Metric;
+import org.eclipse.microprofile.metrics.MetricID;
+import org.eclipse.microprofile.metrics.MetricRegistry;
+import org.eclipse.microprofile.metrics.Tag;
+import org.eclipse.microprofile.metrics.annotation.RegistryType;
+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.mongodb.metrics.ConnectionPoolGauge;
+import io.quarkus.test.QuarkusUnitTest;
+
+public class MongoMetricsTest extends MongoTestBase {
+
+    @Inject
+    MongoClient client;
+
+    @Inject
+    @RegistryType(type = MetricRegistry.Type.VENDOR)
+    MetricRegistry registry;
+
+    @RegisterExtension
+    static final QuarkusUnitTest config = new QuarkusUnitTest()
+            .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class).addClasses(MongoTestBase.class))
+            .withConfigurationResource("application-metrics-mongo.properties");
+
+    @AfterEach
+    void cleanup() {
+        if (client != null) {
+            client.close();
+        }
+    }
+
+    @Test
+    void testMetricsInitialization() {
+        assertNull(getGaugeValueOrNull("mongodb.connection-pool.size", getTags()));
+        assertNull(getGaugeValueOrNull("mongodb.connection-pool.checked-out-count", getTags()));
+
+        // Just need to execute something so that an connection is opened
+        String name = client.listDatabaseNames().first();
+
+        assertEquals(1L, getGaugeValueOrNull("mongodb.connection-pool.size", getTags()));
+        assertEquals(0L, getGaugeValueOrNull("mongodb.connection-pool.checked-out-count", getTags()));
+
+        client.close();
+        assertEquals(0L, getGaugeValueOrNull("mongodb.connection-pool.size", getTags()));
+        assertEquals(0L, getGaugeValueOrNull("mongodb.connection-pool.checked-out-count", getTags()));
+    }
+
+    private Long getGaugeValueOrNull(String metricName, Tag[] tags) {
+        MetricID metricID = new MetricID(metricName, tags);
+        Metric metric = registry.getMetrics().get(metricID);
+
+        if (metric == null) {
+            return null;
+        }
+        return ((ConnectionPoolGauge) metric).getValue();
+    }
+
+    private Tag[] getTags() {
+        return new Tag[] {
+                new Tag("host", "localhost"),
+                new Tag("port", "27018"),
+        };
+    }
+}

extensions/mongodb-client/deployment/src/test/resources/application-metrics-mongo.properties

@@ -0,0 +1,2 @@
+quarkus.mongodb.connection-string=mongodb://localhost:27018
+quarkus.mongodb.metrics.enabled=true

extensions/mongodb-client/runtime/pom.xml

@@ -48,6 +48,12 @@
             <optional>true</optional>
         </dependency>
 
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-smallrye-metrics</artifactId>
+            <optional>true</optional>
+        </dependency>
+
         <dependency>
             <groupId>org.graalvm.nativeimage</groupId>
             <artifactId>svm</artifactId>

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

@@ -0,0 +1,22 @@
+package io.quarkus.mongodb.metrics;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.eclipse.microprofile.metrics.Gauge;
+
+public class ConnectionPoolGauge implements Gauge<Long> {
+    private AtomicLong value = new AtomicLong();
+
+    public void decrement() {
+        value.decrementAndGet();
+    }
+
+    public void increment() {
+        value.incrementAndGet();
+    }
+
+    @Override
+    public Long getValue() {
+        return value.longValue();
+    }
+}

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

@@ -0,0 +1,117 @@
+package io.quarkus.mongodb.metrics;
+
+import org.eclipse.microprofile.metrics.Metadata;
+import org.eclipse.microprofile.metrics.Metric;
+import org.eclipse.microprofile.metrics.MetricID;
+import org.eclipse.microprofile.metrics.MetricRegistry;
+import org.eclipse.microprofile.metrics.MetricType;
+import org.eclipse.microprofile.metrics.Tag;
+
+import com.mongodb.connection.ServerId;
+import com.mongodb.event.ConnectionAddedEvent;
+import com.mongodb.event.ConnectionCheckedInEvent;
+import com.mongodb.event.ConnectionCheckedOutEvent;
+import com.mongodb.event.ConnectionPoolClosedEvent;
+import com.mongodb.event.ConnectionPoolListener;
+import com.mongodb.event.ConnectionPoolOpenedEvent;
+import com.mongodb.event.ConnectionPoolWaitQueueEnteredEvent;
+import com.mongodb.event.ConnectionPoolWaitQueueExitedEvent;
+import com.mongodb.event.ConnectionRemovedEvent;
+
+import io.smallrye.metrics.MetricRegistries;
+
+public class MongoMetricsConnectionPoolListener implements ConnectionPoolListener {
+    private final static String SIZE_NAME = "mongodb.connection-pool.size";
+    private final static String CHECKED_OUT_COUNT_NAME = "mongodb.connection-pool.checked-out-count";
+
+    @Override
+    public void connectionPoolOpened(ConnectionPoolOpenedEvent event) {
+        Tag[] tags = createTags(event.getServerId());
+
+        registerGauge(SIZE_NAME, "the current size of the pool, including idle and and in-use members", tags);
+        registerGauge(CHECKED_OUT_COUNT_NAME, "the current count of connections that are currently in use", tags);
+    }
+
+    @Override
+    public void connectionPoolClosed(ConnectionPoolClosedEvent event) {
+    }
+
+    @Override
+    public void connectionCheckedOut(ConnectionCheckedOutEvent event) {
+        MetricID metricID = createMetricID(CHECKED_OUT_COUNT_NAME, event.getConnectionId().getServerId());
+
+        Metric metric = getMetricRegistry().getMetrics().get(metricID);
+
+        if (metric != null) {
+            ((ConnectionPoolGauge) metric).increment();
+        }
+    }
+
+    @Override
+    public void connectionCheckedIn(ConnectionCheckedInEvent event) {
+        MetricID metricID = createMetricID(CHECKED_OUT_COUNT_NAME, event.getConnectionId().getServerId());
+
+        Metric metric = getMetricRegistry().getMetrics().get(metricID);
+
+        if (metric != null) {
+            ((ConnectionPoolGauge) metric).decrement();
+        }
+    }
+
+    @Override
+    public void waitQueueEntered(ConnectionPoolWaitQueueEnteredEvent connectionPoolWaitQueueEnteredEvent) {
+        // Not supported, since the event is deprecated and will be removed anyway
+    }
+
+    @Override
+    public void waitQueueExited(ConnectionPoolWaitQueueExitedEvent connectionPoolWaitQueueExitedEvent) {
+        // Not supported, since the event is deprecated and will be removed anyway
+    }
+
+    @Override
+    public void connectionAdded(ConnectionAddedEvent event) {
+
+        MetricID metricID = createMetricID(SIZE_NAME, event.getConnectionId().getServerId());
+
+        Metric metric = getMetricRegistry().getMetrics().get(metricID);
+
+        if (metric != null) {
+            ((ConnectionPoolGauge) metric).increment();
+        }
+    }
+
+    @Override
+    public void connectionRemoved(ConnectionRemovedEvent event) {
+
+        MetricID metricID = createMetricID(SIZE_NAME, event.getConnectionId().getServerId());
+
+        Metric metric = getMetricRegistry().getMetrics().get(metricID);
+
+        if (metric != null) {
+            ((ConnectionPoolGauge) metric).decrement();
+        }
+    }
+
+    private void registerGauge(String metricName, String description, Tag[] tags) {
+        getMetricRegistry().remove(new MetricID(metricName, tags));
+
+        Metadata metaData = Metadata.builder().withName(metricName).withType(MetricType.GAUGE)
+                .withDescription(description).build();
+        getMetricRegistry().register(metaData, new ConnectionPoolGauge(), tags);
+    }
+
+    private MetricRegistry getMetricRegistry() {
+        return MetricRegistries.get(MetricRegistry.Type.VENDOR);
+    }
+
+    private Tag[] createTags(ServerId server) {
+        return new Tag[] {
+                new Tag("host", server.getAddress().getHost()),
+                new Tag("port", String.valueOf(server.getAddress().getPort())),
+        };
+    }
+
+    private MetricID createMetricID(String metricName, ServerId server) {
+        return new MetricID(metricName, createTags(server));
+    }
+}