Refactor/fix Prometheus metrics for multiple connection pools, add un…

…it tests (brettwooldridge#1331)

* Refactor/fix Prometheus metrics for multiple connection pools

Changes:
* Fix "Collector already registered that provides name:
hikaricp_connection_timeout_total" error
* Register only one HikariCPCollector instance in one CollectorRegistry
instance
* Add ability to remove metrics when connection pool is shutting down
* Add/update unit tests

* Refactor/add unit tests - metrics package

* Re-format curly braces to be inline with the whole code base

src/main/java/com/zaxxer/hikari/metrics/micrometer/MicrometerMetricsTrackerFactory.java

@@ -5,7 +5,8 @@
 import com.zaxxer.hikari.metrics.PoolStats;
 import io.micrometer.core.instrument.MeterRegistry;
 
-public class MicrometerMetricsTrackerFactory implements MetricsTrackerFactory {
+public class MicrometerMetricsTrackerFactory implements MetricsTrackerFactory
+{
 
    private final MeterRegistry registry;
 

src/main/java/com/zaxxer/hikari/metrics/prometheus/HikariCPCollector.java

@@ -12,28 +12,31 @@
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
-*/
+ */
 
 package com.zaxxer.hikari.metrics.prometheus;
 
 import com.zaxxer.hikari.metrics.PoolStats;
 import io.prometheus.client.Collector;
 import io.prometheus.client.GaugeMetricFamily;
+
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
-class HikariCPCollector extends Collector {
+class HikariCPCollector extends Collector
+{
 
    private static final List<String> LABEL_NAMES = Collections.singletonList("pool");
 
    private final Map<String, PoolStats> poolStatsMap = new ConcurrentHashMap<>();
 
    @Override
-   public List<MetricFamilySamples> collect() {
+   public List<MetricFamilySamples> collect()
+   {
       return Arrays.asList(
          createGauge("hikaricp_active_connections", "Active connections",
             PoolStats::getActiveConnections),
@@ -50,13 +53,19 @@ public List<MetricFamilySamples> collect() {
       );
    }
 
-   protected HikariCPCollector add(String name, PoolStats poolStats) {
+   void add(String name, PoolStats poolStats)
+   {
       poolStatsMap.put(name, poolStats);
-      return this;
+   }
+
+   void remove(String name)
+   {
+      poolStatsMap.remove(name);
    }
 
    private GaugeMetricFamily createGauge(String metric, String help,
-      Function<PoolStats, Integer> metricValueFunction) {
+                                         Function<PoolStats, Integer> metricValueFunction)
+   {
       GaugeMetricFamily metricFamily = new GaugeMetricFamily(metric, help, LABEL_NAMES);
       poolStatsMap.forEach((k, v) -> metricFamily.addMetric(
          Collections.singletonList(k),

src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTracker.java

@@ -12,66 +12,69 @@
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
-*/
+ */
 
 package com.zaxxer.hikari.metrics.prometheus;
 
 import com.zaxxer.hikari.metrics.IMetricsTracker;
+import com.zaxxer.hikari.metrics.prometheus.PrometheusMetricsTrackerFactory.RegistrationStatus;
 import io.prometheus.client.CollectorRegistry;
 import io.prometheus.client.Counter;
 import io.prometheus.client.Summary;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 
+import static com.zaxxer.hikari.metrics.prometheus.PrometheusMetricsTrackerFactory.RegistrationStatus.REGISTERED;
+
 class PrometheusMetricsTracker implements IMetricsTracker
 {
-   private final Counter CONNECTION_TIMEOUT_COUNTER = Counter.build()
+   private final static Counter CONNECTION_TIMEOUT_COUNTER = Counter.build()
       .name("hikaricp_connection_timeout_total")
       .labelNames("pool")
       .help("Connection timeout total count")
       .create();
 
-   private final Summary ELAPSED_ACQUIRED_SUMMARY =
-      registerSummary("hikaricp_connection_acquired_nanos", "Connection acquired time (ns)");
+   private final static Summary ELAPSED_ACQUIRED_SUMMARY =
+      createSummary("hikaricp_connection_acquired_nanos", "Connection acquired time (ns)");
 
-   private final Summary ELAPSED_BORROWED_SUMMARY =
-      registerSummary("hikaricp_connection_usage_millis", "Connection usage (ms)");
+   private final static Summary ELAPSED_USAGE_SUMMARY =
+      createSummary("hikaricp_connection_usage_millis", "Connection usage (ms)");
 
-   private final Summary ELAPSED_CREATION_SUMMARY =
-      registerSummary("hikaricp_connection_creation_millis", "Connection creation (ms)");
+   private final static Summary ELAPSED_CREATION_SUMMARY =
+      createSummary("hikaricp_connection_creation_millis", "Connection creation (ms)");
 
-   private final Counter.Child connectionTimeoutCounterChild;
+   private final static Map<CollectorRegistry, RegistrationStatus> registrationStatuses = new ConcurrentHashMap<>();
 
-   private Summary registerSummary(String name, String help) {
-      return Summary.build()
-         .name(name)
-         .labelNames("pool")
-         .help(help)
-         .quantile(0.5, 0.05)
-         .quantile(0.95, 0.01)
-         .quantile(0.99, 0.001)
-         .maxAgeSeconds(TimeUnit.MINUTES.toSeconds(5))
-         .ageBuckets(5)
-         .create();
-   }
+   private final String poolName;
+   private final HikariCPCollector hikariCPCollector;
+
+   private final Counter.Child connectionTimeoutCounterChild;
 
    private final Summary.Child elapsedAcquiredSummaryChild;
-   private final Summary.Child elapsedBorrowedSummaryChild;
+   private final Summary.Child elapsedUsageSummaryChild;
    private final Summary.Child elapsedCreationSummaryChild;
 
-   PrometheusMetricsTracker(String poolName, CollectorRegistry collectorRegistry) {
+   PrometheusMetricsTracker(String poolName, CollectorRegistry collectorRegistry, HikariCPCollector hikariCPCollector)
+   {
       registerMetrics(collectorRegistry);
+      this.poolName = poolName;
+      this.hikariCPCollector = hikariCPCollector;
       this.connectionTimeoutCounterChild = CONNECTION_TIMEOUT_COUNTER.labels(poolName);
       this.elapsedAcquiredSummaryChild = ELAPSED_ACQUIRED_SUMMARY.labels(poolName);
-      this.elapsedBorrowedSummaryChild = ELAPSED_BORROWED_SUMMARY.labels(poolName);
+      this.elapsedUsageSummaryChild = ELAPSED_USAGE_SUMMARY.labels(poolName);
       this.elapsedCreationSummaryChild = ELAPSED_CREATION_SUMMARY.labels(poolName);
    }
 
-   private void registerMetrics(CollectorRegistry collectorRegistry){
-      CONNECTION_TIMEOUT_COUNTER.register(collectorRegistry);
-      ELAPSED_ACQUIRED_SUMMARY.register(collectorRegistry);
-      ELAPSED_BORROWED_SUMMARY.register(collectorRegistry);
-      ELAPSED_CREATION_SUMMARY.register(collectorRegistry);
+   private void registerMetrics(CollectorRegistry collectorRegistry)
+   {
+      if (registrationStatuses.putIfAbsent(collectorRegistry, REGISTERED) == null) {
+         CONNECTION_TIMEOUT_COUNTER.register(collectorRegistry);
+         ELAPSED_ACQUIRED_SUMMARY.register(collectorRegistry);
+         ELAPSED_USAGE_SUMMARY.register(collectorRegistry);
+         ELAPSED_CREATION_SUMMARY.register(collectorRegistry);
+      }
    }
 
    @Override
@@ -83,7 +86,7 @@ public void recordConnectionAcquiredNanos(long elapsedAcquiredNanos)
    @Override
    public void recordConnectionUsageMillis(long elapsedBorrowedMillis)
    {
-      elapsedBorrowedSummaryChild.observe(elapsedBorrowedMillis);
+      elapsedUsageSummaryChild.observe(elapsedBorrowedMillis);
    }
 
    @Override
@@ -97,4 +100,28 @@ public void recordConnectionTimeout()
    {
       connectionTimeoutCounterChild.inc();
    }
+
+   private static Summary createSummary(String name, String help)
+   {
+      return Summary.build()
+         .name(name)
+         .labelNames("pool")
+         .help(help)
+         .quantile(0.5, 0.05)
+         .quantile(0.95, 0.01)
+         .quantile(0.99, 0.001)
+         .maxAgeSeconds(TimeUnit.MINUTES.toSeconds(5))
+         .ageBuckets(5)
+         .create();
+   }
+
+   @Override
+   public void close()
+   {
+      hikariCPCollector.remove(poolName);
+      CONNECTION_TIMEOUT_COUNTER.remove(poolName);
+      ELAPSED_ACQUIRED_SUMMARY.remove(poolName);
+      ELAPSED_USAGE_SUMMARY.remove(poolName);
+      ELAPSED_CREATION_SUMMARY.remove(poolName);
+   }
 }

src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTrackerFactory.java

@@ -12,59 +12,78 @@
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
-*/
+ */
 
 package com.zaxxer.hikari.metrics.prometheus;
 
 import com.zaxxer.hikari.metrics.IMetricsTracker;
 import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
 import com.zaxxer.hikari.metrics.PoolStats;
+import io.prometheus.client.Collector;
 import io.prometheus.client.CollectorRegistry;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static com.zaxxer.hikari.metrics.prometheus.PrometheusMetricsTrackerFactory.RegistrationStatus.REGISTERED;
+
 /**
  * <pre>{@code
  * HikariConfig config = new HikariConfig();
  * config.setMetricsTrackerFactory(new PrometheusMetricsTrackerFactory());
  * }</pre>
+ * or
+ * <pre>{@code
+ * config.setMetricsTrackerFactory(new PrometheusMetricsTrackerFactory(new CollectorRegistry()));
+ * }</pre>
  *
  * Note: the internal {@see io.prometheus.client.Summary} requires heavy locks. Consider using
  * {@see PrometheusHistogramMetricsTrackerFactory} if performance plays a role and you don't need the summary per se.
  */
-public class PrometheusMetricsTrackerFactory implements MetricsTrackerFactory {
+public class PrometheusMetricsTrackerFactory implements MetricsTrackerFactory
+{
 
-   private HikariCPCollector collector;
+   private final static Map<CollectorRegistry, RegistrationStatus> registrationStatuses = new ConcurrentHashMap<>();
 
-   private CollectorRegistry collectorRegistry;
+   private final HikariCPCollector collector = new HikariCPCollector();
+
+   private final CollectorRegistry collectorRegistry;
+
+   public enum RegistrationStatus
+   {
+      REGISTERED;
+   }
 
    /**
     * Default Constructor. The Hikari metrics are registered to the default
     * collector registry ({@code CollectorRegistry.defaultRegistry}).
     */
-   public PrometheusMetricsTrackerFactory() {
-      this.collectorRegistry = CollectorRegistry.defaultRegistry;
+   public PrometheusMetricsTrackerFactory()
+   {
+      this(CollectorRegistry.defaultRegistry);
    }
 
    /**
     * Constructor that allows to pass in a {@link CollectorRegistry} to which the
     * Hikari metrics are registered.
     */
-   public PrometheusMetricsTrackerFactory(CollectorRegistry collectorRegistry) {
+   public PrometheusMetricsTrackerFactory(CollectorRegistry collectorRegistry)
+   {
       this.collectorRegistry = collectorRegistry;
    }
 
    @Override
-   public IMetricsTracker create(String poolName, PoolStats poolStats) {
-      getCollector().add(poolName, poolStats);
-      return new PrometheusMetricsTracker(poolName, this.collectorRegistry);
+   public IMetricsTracker create(String poolName, PoolStats poolStats)
+   {
+      registerCollector(this.collector, this.collectorRegistry);
+      this.collector.add(poolName, poolStats);
+      return new PrometheusMetricsTracker(poolName, this.collectorRegistry, this.collector);
    }
 
-   /**
-    * initialize and register collector if it isn't initialized yet
-    */
-   private HikariCPCollector getCollector() {
-      if (collector == null) {
-         collector = new HikariCPCollector().register(this.collectorRegistry);
+   private void registerCollector(Collector collector, CollectorRegistry collectorRegistry)
+   {
+      if (registrationStatuses.putIfAbsent(collectorRegistry, REGISTERED) == null) {
+         collector.register(collectorRegistry);
       }
-      return collector;
    }
 }

src/test/java/com/zaxxer/hikari/metrics/dropwizard/CodaHaleMetricsTrackerTest.java

@@ -1,31 +1,33 @@
 package com.zaxxer.hikari.metrics.dropwizard;
 
-import static org.mockito.Mockito.verify;
-
-import com.zaxxer.hikari.metrics.PoolStats;
+import com.codahale.metrics.MetricRegistry;
+import com.zaxxer.hikari.mocks.StubPoolStats;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
-import com.codahale.metrics.MetricRegistry;
+import static org.mockito.Mockito.verify;
 
 @RunWith(MockitoJUnitRunner.class)
-public class CodaHaleMetricsTrackerTest {
+public class CodaHaleMetricsTrackerTest
+{
 
    @Mock
    public MetricRegistry mockMetricRegistry;
 
    private CodaHaleMetricsTracker testee;
 
    @Before
-   public void setup() {
-      testee = new CodaHaleMetricsTracker("mypool", poolStats(), mockMetricRegistry);
+   public void setup()
+   {
+      testee = new CodaHaleMetricsTracker("mypool", new StubPoolStats(0), mockMetricRegistry);
    }
 
    @Test
-   public void close() throws Exception {
+   public void close()
+   {
       testee.close();
 
       verify(mockMetricRegistry).remove("mypool.pool.Wait");
@@ -39,13 +41,4 @@ public void close() throws Exception {
       verify(mockMetricRegistry).remove("mypool.pool.MaxConnections");
       verify(mockMetricRegistry).remove("mypool.pool.MinConnections");
    }
-
-   private PoolStats poolStats() {
-      return new PoolStats(0) {
-         @Override
-         protected void update() {
-            // do nothing
-         }
-      };
-   }
 }

src/test/java/com/zaxxer/hikari/metrics/micrometer/MicrometerMetricsTrackerTest.java

@@ -1,30 +1,28 @@
 package com.zaxxer.hikari.metrics.micrometer;
 
-import com.zaxxer.hikari.metrics.PoolStats;
+import com.zaxxer.hikari.mocks.StubPoolStats;
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
-public class MicrometerMetricsTrackerTest {
+public class MicrometerMetricsTrackerTest
+{
 
    private MeterRegistry mockMeterRegistry = new SimpleMeterRegistry();
 
    private MicrometerMetricsTracker testee;
 
    @Before
-   public void setup(){
-      testee = new MicrometerMetricsTracker("mypool", new PoolStats(1000L) {
-         @Override
-         protected void update() {
-            // nothing
-         }
-      }, mockMeterRegistry);
+   public void setup()
+   {
+      testee = new MicrometerMetricsTracker("mypool", new StubPoolStats(1000L), mockMeterRegistry);
    }
 
    @Test
-   public void close() throws Exception {
+   public void close()
+   {
       Assert.assertNotNull(mockMeterRegistry.find("hikaricp.connections.acquire").tag("pool", "mypool").timer());
       Assert.assertNotNull(mockMeterRegistry.find("hikaricp.connections.usage").tag("pool", "mypool").timer());
       Assert.assertNotNull(mockMeterRegistry.find("hikaricp.connections.creation").tag("pool", "mypool").timer());