Add current pool size metric

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraClientPoolImpl.java

@@ -359,6 +359,7 @@ void setServersInPoolTo(ImmutableMap<CassandraServer, CassandraServerOrigin> des
 
         if (!(validatedServersToAdd.isEmpty() && absentServers.isEmpty())) { // if we made any changes
             cassandra.refreshTokenRangesAndGetServers();
+            metrics.recordPoolSize(getCurrentPools().size());
         }
 
         Preconditions.checkState(

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/pool/CassandraClientPoolMetrics.java

@@ -18,6 +18,7 @@
 import com.codahale.metrics.Counter;
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Meter;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.palantir.atlasdb.keyvalue.cassandra.CassandraClientPool;
 import com.palantir.atlasdb.keyvalue.cassandra.CassandraClientPoolingContainer;
@@ -26,10 +27,13 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 public class CassandraClientPoolMetrics {
+
+    public static final String POOL_SIZE_METRIC_NAME = "pool-size";
     private final MetricsManager metricsManager;
     private final RequestMetrics aggregateRequestMetrics;
     private final Map<CassandraServer, RequestMetrics> metricsByHost = new HashMap<>();
@@ -39,12 +43,15 @@ public class CassandraClientPoolMetrics {
     // Not bundled in with request metrics, as we seek to not produce host-level metrics for economic reasons.
     private final Counter poolExhaustionCounter;
 
+    private final AtomicLong poolSize = new AtomicLong(0L);
+
     public CassandraClientPoolMetrics(MetricsManager metricsManager) {
         this.metricsManager = metricsManager;
         this.aggregateRequestMetrics = new RequestMetrics(metricsManager);
         this.poolExhaustionCounter =
                 metricsManager.registerOrGetCounter(CassandraClientPoolMetrics.class, "pool-exhaustion");
         this.outlierControllers = createOutlierControllers(metricsManager);
+        metricsManager.registerMetric(CassandraClientPoolMetrics.class, POOL_SIZE_METRIC_NAME, poolSize::get);
     }
 
     private static Map<CassandraClientPoolHostLevelMetric, DistributionOutlierController> createOutlierControllers(
@@ -93,6 +100,15 @@ public void recordPoolExhaustion() {
         poolExhaustionCounter.inc();
     }
 
+    public void recordPoolSize(long desiredPoolSize) {
+        poolSize.set(desiredPoolSize);
+    }
+
+    @VisibleForTesting
+    Long getPoolSize() {
+        return poolSize.get();
+    }
+
     @SuppressWarnings("unchecked") // Guaranteed to have the correct type
     public void registerPoolMetric(CassandraClientPoolHostLevelMetric metric, Gauge<Long> gauge, int poolNumber) {
         MetricPublicationFilter filter = outlierControllers.get(metric).registerAndCreateFilter(gauge);

atlasdb-cassandra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraClientPoolTest.java

@@ -35,6 +35,7 @@
 import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceRuntimeConfig;
 import com.palantir.atlasdb.cassandra.CassandraServersConfigs;
 import com.palantir.atlasdb.cassandra.ImmutableDefaultConfig;
+import com.palantir.atlasdb.keyvalue.cassandra.pool.CassandraClientPoolMetrics;
 import com.palantir.atlasdb.keyvalue.cassandra.pool.CassandraServer;
 import com.palantir.atlasdb.keyvalue.cassandra.pool.CassandraService;
 import com.palantir.atlasdb.util.MetricsManagers;
@@ -224,7 +225,7 @@ private void runTwoNoopsOnTwoHostsAndThrowFromSecondRunOnFirstHost(Exception exc
     }
 
     private CassandraServer getServerForIndex(int index) {
-        return CassandraServer.of(InetSocketAddress.createUnresolved(Integer.toString(index), index));
+        return CassandraServer.of(InetSocketAddress.createUnresolved(Integer.toString(index), 8000));
     }
 
     @Test
@@ -310,6 +311,7 @@ public void hostIsAutomaticallyRemovedOnStartup() {
 
         createClientPool();
         assertThat(poolServers).containsOnlyKeys(CASS_SERVER_1);
+        assertPoolSizeMetricIsEqualTo(1);
     }
 
     @Test
@@ -321,10 +323,12 @@ public void hostIsAutomaticallyRemovedOnRefresh() {
 
         createClientPool();
         assertThat(poolServers).containsOnlyKeys(CASS_SERVER_1, CASS_SERVER_2, CASS_SERVER_3);
+        assertPoolSizeMetricIsEqualTo(3);
 
         setCassandraServersTo(CASS_SERVER_1, CASS_SERVER_2);
         refreshPool();
         assertThat(poolServers).containsOnlyKeys(CASS_SERVER_1, CASS_SERVER_2);
+        assertPoolSizeMetricIsEqualTo(2);
     }
 
     @Test
@@ -495,7 +499,7 @@ private void setCassandraServersTo(CassandraServer... servers) {
 
     private CassandraClientPoolImpl createClientPool() {
         return CassandraClientPoolImpl.createImplForTest(
-                MetricsManagers.createForTests(),
+                MetricsManagers.of(metricRegistry, taggedMetricRegistry),
                 config,
                 refreshableRuntimeConfig,
                 CassandraClientPoolImpl.StartupChecks.DO_NOT_RUN,
@@ -687,4 +691,16 @@ private void setupHostsWithInconsistentTopology(CassandraServer... cassandraServ
     private void refreshPool() {
         deterministicExecutor.tick(POOL_REFRESH_INTERVAL_SECONDS, TimeUnit.SECONDS);
     }
+
+    private void assertPoolSizeMetricIsEqualTo(long expected) {
+        assertThat(taggedMetricRegistry
+                        .<Long>gauge(MetricName.builder()
+                                .safeName(MetricRegistry.name(
+                                        CassandraClientPoolMetrics.class,
+                                        CassandraClientPoolMetrics.POOL_SIZE_METRIC_NAME))
+                                .build())
+                        .get()
+                        .getValue())
+                .isEqualTo(expected);
+    }
 }

atlasdb-cassandra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/pool/CassandraClientPoolMetricsTest.java

@@ -33,6 +33,8 @@ public class CassandraClientPoolMetricsTest {
     private final MetricsManager metricsManager =
             MetricsManagers.of(new MetricRegistry(), new DefaultTaggedMetricRegistry(), Refreshable.only(true));
 
+    private final CassandraClientPoolMetrics metrics = new CassandraClientPoolMetrics(metricsManager);
+
     @Test
     public void metricsAreProducedAndFiltered() {
         CassandraClientPoolMetrics metrics = new CassandraClientPoolMetrics(metricsManager);
@@ -57,6 +59,18 @@ public void metricsAreProducedAndFiltered() {
                 .containsKey(createMeanActiveTimeMillisMetric("mean"));
     }
 
+    @Test
+    public void recordPoolSizeSetsSizeToLastReportedValue() {
+        metrics.recordPoolSize(100);
+        assertThat(metrics.getPoolSize()).isEqualTo(100);
+        metrics.recordPoolSize(-13);
+        assertThat(metrics.getPoolSize()).isEqualTo(-13);
+        metrics.recordPoolSize(25);
+        assertThat(metrics.getPoolSize()).isEqualTo(25);
+        metrics.recordPoolSize(0);
+        assertThat(metrics.getPoolSize()).isEqualTo(0);
+    }
+
     private static MetricName createMeanActiveTimeMillisMetric(String pool) {
         return MetricName.builder()
                 .safeName(MetricRegistry.name(

changelog/@unreleased/pr-6554.v2.yml

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: Adds metric for tracking the size of the cassandra client pool
+  links:
+  - https://github.com/palantir/atlasdb/pull/6554