Merge pull request #331 from arthurdandrea/dev

Allow the use of a user provided MetricsTracker

src/main/java/com/zaxxer/hikari/HikariConfig.java

@@ -42,6 +42,8 @@ import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.health.HealthCheckRegistry;
 import com.zaxxer.hikari.pool.PoolElf;
 import com.zaxxer.hikari.util.PropertyElf;
+import com.zaxxer.hikari.metrics.CodahaleMetricsTrackerFactory;
+import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
 
 public class HikariConfig implements HikariConfigMXBean
 {
@@ -88,7 +90,7 @@ public class HikariConfig implements HikariConfigMXBean
    private Properties dataSourceProperties;
    private ThreadFactory threadFactory;
    private ScheduledThreadPoolExecutor scheduledExecutor;
-   private Object metricRegistry;
+   private MetricsTrackerFactory metricsTrackerFactory;
    private Object healthCheckRegistry;
    private Properties healthCheckProperties;
 
@@ -442,6 +444,16 @@ public class HikariConfig implements HikariConfigMXBean
       LOGGER.warn("The jdbcConnectionTest property is now deprecated, see the documentation for connectionTestQuery");
    }
 
+   public MetricsTrackerFactory getMetricsTrackerFactory()
+   {
+      return metricsTrackerFactory;
+   }
+
+   public void setMetricsTrackerFactory(MetricsTrackerFactory metricsTrackerFactory)
+   {
+      this.metricsTrackerFactory = metricsTrackerFactory;
+   }
+
    /**
     * Get the Codahale MetricRegistry, could be null.
     *
@@ -449,7 +461,12 @@ public class HikariConfig implements HikariConfigMXBean
     */
    public Object getMetricRegistry()
    {
-      return metricRegistry;
+      if (metricsTrackerFactory instanceof CodahaleMetricsTrackerFactory) {
+         return ((CodahaleMetricsTrackerFactory) metricsTrackerFactory).getRegistry();
+      }
+      else {
+         return null;
+      }
    }
 
    /**
@@ -459,6 +476,7 @@ public class HikariConfig implements HikariConfigMXBean
     */
    public void setMetricRegistry(Object metricRegistry)
    {
+      MetricsTrackerFactory metricsTrackerFactory;
       if (metricRegistry != null) {
          if (metricRegistry instanceof String) {
             try {
@@ -471,11 +489,15 @@ public class HikariConfig implements HikariConfigMXBean
          }
 
          if (!(metricRegistry instanceof MetricRegistry)) {
-            throw new IllegalArgumentException("Class must be an instance of com.codahale.metrics.MetricRegistry");            
+            throw new IllegalArgumentException("Class must be an instance of com.codahale.metrics.MetricRegistry");
          }
+         metricsTrackerFactory = new CodahaleMetricsTrackerFactory((MetricRegistry) metricRegistry);
       }
-
-      this.metricRegistry = metricRegistry;
+      else {
+         metricsTrackerFactory = null;
+      }
+      // use setter to allow HikariDataSource to alter the behavior of both setters in a single override
+      setMetricsTrackerFactory(metricsTrackerFactory);
    }
 
    /**

src/main/java/com/zaxxer/hikari/HikariDataSource.java

@@ -29,6 +29,7 @@ import javax.sql.DataSource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
 import com.zaxxer.hikari.pool.HikariPool;
 import com.zaxxer.hikari.proxy.IHikariConnectionProxy;
 import com.zaxxer.hikari.util.DriverDataSource;
@@ -192,17 +193,17 @@ public class HikariDataSource extends HikariConfig implements DataSource, Closea
 
    /** {@inheritDoc} */
    @Override
-   public void setMetricRegistry(Object metricRegistry)
+   public void setMetricsTrackerFactory(MetricsTrackerFactory metricsTrackerFactory)
    {
-      boolean isAlreadySet = getMetricRegistry() != null;
-      super.setMetricRegistry(metricRegistry);
+      boolean isAlreadySet = getMetricsTrackerFactory() != null;
+      super.setMetricsTrackerFactory(metricsTrackerFactory);
 
       if (pool != null) {
          if (isAlreadySet) {
-            throw new IllegalStateException("MetricRegistry can only be set one time");
+            throw new IllegalStateException("MetricsTrackerFactory can only be set one time");
          }
          else {
-            pool.setMetricRegistry(super.getMetricRegistry());
+            pool.setMetricsTrackerFactory(super.getMetricsTrackerFactory());
          }
       }
    }

src/main/java/com/zaxxer/hikari/metrics/CodahaleMetricsTrackerFactory.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2013,2014 Brett Wooldridge
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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;
+
+import com.codahale.metrics.MetricRegistry;
+
+public final class CodahaleMetricsTrackerFactory implements MetricsTrackerFactory
+{
+   private final MetricRegistry registry;
+
+   public CodahaleMetricsTrackerFactory(MetricRegistry registry) {
+      this.registry = registry;
+   }
+
+   public MetricRegistry getRegistry()
+   {
+      return registry;
+   }
+
+   @Override
+   public MetricsTracker create(String poolName, PoolStats poolStats)
+   {
+      return new CodaHaleMetricsTracker(poolName, poolStats, registry);
+   }
+}

src/main/java/com/zaxxer/hikari/metrics/MetricsTrackerFactory.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2013,2014 Brett Wooldridge
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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;
+
+public interface MetricsTrackerFactory {
+   MetricsTracker create(String poolName, PoolStats poolStats);
+}

src/main/java/com/zaxxer/hikari/pool/HikariPool.java

@@ -46,7 +46,9 @@ import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariPoolMXBean;
 import com.zaxxer.hikari.metrics.CodaHaleMetricsTracker;
 import com.zaxxer.hikari.metrics.CodahaleHealthChecker;
+import com.zaxxer.hikari.metrics.CodahaleMetricsTrackerFactory;
 import com.zaxxer.hikari.metrics.MetricsTracker;
+import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
 import com.zaxxer.hikari.metrics.MetricsTracker.MetricsContext;
 import com.zaxxer.hikari.metrics.PoolStats;
 import com.zaxxer.hikari.proxy.ConnectionProxy;
@@ -285,9 +287,19 @@ public class HikariPool implements HikariPoolMXBean, IBagStateListener
 
    public void setMetricRegistry(Object metricRegistry)
    {
-      this.isRecordMetrics = metricRegistry != null;
+      if (metricRegistry != null) {
+         setMetricsTrackerFactory(new CodahaleMetricsTrackerFactory((MetricRegistry) metricRegistry));
+      }
+      else {
+         setMetricsTrackerFactory(null);
+      }
+   }
+
+   public void setMetricsTrackerFactory(MetricsTrackerFactory metricsTrackerFactory)
+   {
+      this.isRecordMetrics = metricsTrackerFactory != null;
       if (isRecordMetrics) {
-         this.metricsTracker = new CodaHaleMetricsTracker(poolName, getPoolStats(), (MetricRegistry) metricRegistry);
+         this.metricsTracker = metricsTrackerFactory.create(config.getPoolName(), getPoolStats());
       }
       else {
          this.metricsTracker = new MetricsTracker();

src/test/java/com/zaxxer/hikari/TestMetrics.java

@@ -31,6 +31,8 @@ import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.Timer;
 import com.codahale.metrics.health.HealthCheck.Result;
 import com.codahale.metrics.health.HealthCheckRegistry;
+import com.zaxxer.hikari.metrics.CodahaleMetricsTrackerFactory;
+import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
 import com.zaxxer.hikari.util.UtilityElf;
 
 /**
@@ -218,4 +220,79 @@ public class TestMetrics
          ds.close();
       }
    }
+
+   @Test
+   public void testSetters3() throws Exception
+   {
+      HikariDataSource ds = new HikariDataSource();
+      ds.setMaximumPoolSize(1);
+      ds.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
+
+      MetricRegistry metricRegistry = new MetricRegistry();
+      MetricsTrackerFactory metricsTrackerFactory = new CodahaleMetricsTrackerFactory(metricRegistry);
+
+      try {
+         Connection connection = ds.getConnection();
+         connection.close();
+
+         // After the pool as started, we can only set them once...
+         ds.setMetricsTrackerFactory(metricsTrackerFactory);
+
+         // and never again...
+         ds.setMetricsTrackerFactory(metricsTrackerFactory);
+         Assert.fail("Should not have been allowed to set metricsTrackerFactory after pool started");
+      }
+      catch (IllegalStateException ise) {
+         // pass
+         try {
+            // and never again... (even when calling another method)
+            ds.setMetricRegistry(metricRegistry);
+            Assert.fail("Should not have been allowed to set registry after pool started");
+         }
+         catch (IllegalStateException ise2) {
+            // pass
+         }
+      }
+      finally {
+         ds.close();
+      }
+   }
+
+   @Test
+   public void testSetters4() throws Exception
+   {
+      HikariDataSource ds = new HikariDataSource();
+      ds.setMaximumPoolSize(1);
+      ds.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
+
+      MetricRegistry metricRegistry = new MetricRegistry();
+      MetricsTrackerFactory metricsTrackerFactory = new CodahaleMetricsTrackerFactory(metricRegistry);
+
+      // before the pool is started, we can set it any number of times using either setter
+      ds.setMetricRegistry(metricRegistry);
+      ds.setMetricsTrackerFactory(metricsTrackerFactory);
+
+      try {
+         Connection connection = ds.getConnection();
+         connection.close();
+
+         // after the pool is started, we cannot set it any more
+         ds.setMetricRegistry(metricRegistry);
+         Assert.fail("Should not have been allowed to set registry after pool started");
+      }
+      catch (IllegalStateException ise) {
+         // pass
+         try {
+            // after the pool is started, we cannot set it any more
+            ds.setMetricsTrackerFactory(metricsTrackerFactory);
+            Assert.fail("Should not have been allowed to set metricsTrackerFactory after pool started");
+         }
+         catch (IllegalStateException ise2) {
+            // pass
+         }
+      }
+      finally {
+         ds.close();
+      }
+   }
 }