From 5c1bbe6bc9ff10f8b2059e28ce2fd0282fd0434d Mon Sep 17 00:00:00 2001
From: Brett Wooldridge <brett.wooldridge@gmail.com>
Date: Mon, 19 Dec 2016 00:32:41 +0900
Subject: [PATCH] Cleanup registration of Prometheus metrics (by unregistering
 on close()).

---
 .../metrics/prometheus/HikariCPCollector.java | 10 ++---
 .../prometheus/PrometheusMetricsTracker.java  | 39 ++++++++++++++-----
 .../PrometheusMetricsTrackerFactory.java      |  6 ++-
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/src/main/java/com/zaxxer/hikari/metrics/prometheus/HikariCPCollector.java b/src/main/java/com/zaxxer/hikari/metrics/prometheus/HikariCPCollector.java
index 899bf7dff..3bff9748d 100644
--- a/src/main/java/com/zaxxer/hikari/metrics/prometheus/HikariCPCollector.java
+++ b/src/main/java/com/zaxxer/hikari/metrics/prometheus/HikariCPCollector.java
@@ -44,11 +44,11 @@ public List<MetricFamilySamples> collect() {
       );
    }
 
-   private MetricFamilySamples createSample(String name, String helpMessage, double value) {
-      List<MetricFamilySamples.Sample> samples = Collections.singletonList(new MetricFamilySamples.Sample(name,
-         labelNames,
-         labelValues,
-         value));
+   private MetricFamilySamples createSample(String name, String helpMessage, double value)
+   {
+      List<MetricFamilySamples.Sample> samples = Collections.singletonList(
+         new MetricFamilySamples.Sample(name, labelNames, labelValues, value)
+      );
 
       return new MetricFamilySamples(name, Type.GAUGE, helpMessage, samples);
    }
diff --git a/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTracker.java b/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTracker.java
index c9446f66b..dfcc9945e 100644
--- a/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTracker.java
+++ b/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTracker.java
@@ -17,6 +17,9 @@
 package com.zaxxer.hikari.metrics.prometheus;
 
 import com.zaxxer.hikari.metrics.MetricsTracker;
+
+import io.prometheus.client.Collector;
+import io.prometheus.client.CollectorRegistry;
 import io.prometheus.client.Counter;
 import io.prometheus.client.Summary;
 
@@ -27,38 +30,54 @@ class PrometheusMetricsTracker extends MetricsTracker
    private final Summary.Child elapsedBorrowedSummary;
    private final Summary.Child elapsedCreationSummary;
 
-   PrometheusMetricsTracker(String poolName)
+   private final Counter ctCounter;
+   private final Summary eaSummary;
+   private final Summary ebSummary;
+   private final Summary ecSummary;
+   private final Collector collector;
+
+   PrometheusMetricsTracker(String poolName, Collector collector)
    {
-      super();
+      this.collector = collector;
 
-      Counter counter = Counter.build()
+      ctCounter = Counter.build()
          .name("hikaricp_connection_timeout_count")
          .labelNames("pool")
          .help("Connection timeout count")
          .register();
 
-      this.connectionTimeoutCounter = counter.labels(poolName);
+      this.connectionTimeoutCounter = ctCounter.labels(poolName);
 
-      Summary elapsedAcquiredSummary = Summary.build()
+      eaSummary = Summary.build()
          .name("hikaricp_connection_acquired_nanos")
          .labelNames("pool")
          .help("Connection acquired time (ns)")
          .register();
-      this.elapsedAcquiredSummary = elapsedAcquiredSummary.labels(poolName);
+      this.elapsedAcquiredSummary = eaSummary.labels(poolName);
 
-      Summary elapsedBorrowedSummary = Summary.build()
+      ebSummary = Summary.build()
          .name("hikaricp_connection_usage_millis")
          .labelNames("pool")
          .help("Connection usage (ms)")
          .register();
-      this.elapsedBorrowedSummary = elapsedBorrowedSummary.labels(poolName);
+      this.elapsedBorrowedSummary = ebSummary.labels(poolName);
 
-      Summary elapsedCreationSummary = Summary.build()
+      ecSummary = Summary.build()
             .name("hikaricp_connection_creation_millis")
             .labelNames("pool")
             .help("Connection creation (ms)")
             .register();
-      this.elapsedCreationSummary = elapsedCreationSummary.labels(poolName);
+      this.elapsedCreationSummary = ecSummary.labels(poolName);
+   }
+
+   @Override
+   public void close()
+   {
+      CollectorRegistry.defaultRegistry.unregister(ctCounter);
+      CollectorRegistry.defaultRegistry.unregister(eaSummary);
+      CollectorRegistry.defaultRegistry.unregister(ebSummary);
+      CollectorRegistry.defaultRegistry.unregister(ecSummary);
+      CollectorRegistry.defaultRegistry.unregister(collector);
    }
 
    @Override
diff --git a/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTrackerFactory.java b/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTrackerFactory.java
index b38c3e4ce..77c52d6c5 100644
--- a/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTrackerFactory.java
+++ b/src/main/java/com/zaxxer/hikari/metrics/prometheus/PrometheusMetricsTrackerFactory.java
@@ -20,6 +20,8 @@
 import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
 import com.zaxxer.hikari.metrics.PoolStats;
 
+import io.prometheus.client.Collector;
+
 /**
  * <pre>{@code
  * HikariConfig config = new HikariConfig();
@@ -29,7 +31,7 @@
 public class PrometheusMetricsTrackerFactory implements MetricsTrackerFactory {
    @Override
    public MetricsTracker create(String poolName, PoolStats poolStats) {
-      new HikariCPCollector(poolName, poolStats).register();
-      return new PrometheusMetricsTracker(poolName);
+      Collector collector = new HikariCPCollector(poolName, poolStats).register();
+      return new PrometheusMetricsTracker(poolName, collector);
    }
 }