Merge pull request #1585 from gjesse/pluggable-healthcheck-registry

allow setting a custom health check registry per #1584
dropwizard · Jun 3, 2016 · 9a6e7d6 · 9a6e7d6
2 parents 8ba1842 + eedca3a
commit 9a6e7d6
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 3 deletions.
diff --git a/dropwizard-core/src/main/java/io/dropwizard/cli/EnvironmentCommand.java b/dropwizard-core/src/main/java/io/dropwizard/cli/EnvironmentCommand.java
@@ -33,7 +33,8 @@ protected void run(Bootstrap<T> bootstrap, Namespace namespace, T configuration)
                                                         bootstrap.getObjectMapper(),
                                                         bootstrap.getValidatorFactory().getValidator(),
                                                         bootstrap.getMetricRegistry(),
-                                                        bootstrap.getClassLoader());
+                                                        bootstrap.getClassLoader(),
+                                                        bootstrap.getHealthCheckRegistry());
         configuration.getMetricsFactory().configure(environment.lifecycle(),
                                                     bootstrap.getMetricRegistry());
         configuration.getServerFactory().configure(environment);

diff --git a/dropwizard-core/src/main/java/io/dropwizard/setup/Bootstrap.java b/dropwizard-core/src/main/java/io/dropwizard/setup/Bootstrap.java
@@ -3,6 +3,7 @@
 import com.codahale.metrics.JmxReporter;
 import com.codahale.metrics.JvmAttributeGaugeSet;
 import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.health.HealthCheckRegistry;
 import com.codahale.metrics.jvm.BufferPoolMetricSet;
 import com.codahale.metrics.jvm.ClassLoadingGaugeSet;
 import com.codahale.metrics.jvm.FileDescriptorRatioGauge;
@@ -51,6 +52,7 @@ public class Bootstrap<T extends Configuration> {
     private ValidatorFactory validatorFactory;
 
     private boolean metricsAreRegistered;
+    private HealthCheckRegistry healthCheckRegistry;
 
     /**
      * Creates a new {@link Bootstrap} for the given application.
@@ -68,6 +70,7 @@ public Bootstrap(Application<T> application) {
         this.configurationSourceProvider = new FileConfigurationSourceProvider();
         this.classLoader = Thread.currentThread().getContextClassLoader();
         this.configurationFactoryFactory = new DefaultConfigurationFactoryFactory<>();
+        this.healthCheckRegistry = new HealthCheckRegistry();
     }
 
     /**
@@ -239,4 +242,15 @@ public ConfigurationFactoryFactory<T> getConfigurationFactoryFactory() {
     public void setConfigurationFactoryFactory(ConfigurationFactoryFactory<T> configurationFactoryFactory) {
         this.configurationFactoryFactory = configurationFactoryFactory;
     }
+
+    /**
+     * returns the health check registry
+     */
+    public HealthCheckRegistry getHealthCheckRegistry() {
+        return healthCheckRegistry;
+    }
+
+    public void setHealthCheckRegistry(HealthCheckRegistry healthCheckRegistry) {
+        this.healthCheckRegistry = healthCheckRegistry;
+    }
 }
diff --git a/dropwizard-core/src/main/java/io/dropwizard/setup/Environment.java b/dropwizard-core/src/main/java/io/dropwizard/setup/Environment.java
@@ -55,11 +55,12 @@ public Environment(String name,
                        ObjectMapper objectMapper,
                        Validator validator,
                        MetricRegistry metricRegistry,
-                       ClassLoader classLoader) {
+                       ClassLoader classLoader,
+                       HealthCheckRegistry healthCheckRegistry) {
         this.name = name;
         this.objectMapper = objectMapper;
         this.metricRegistry = metricRegistry;
-        this.healthCheckRegistry = new HealthCheckRegistry();
+        this.healthCheckRegistry = healthCheckRegistry;
         this.validator = validator;
 
         this.servletContext = new MutableServletContextHandler();
@@ -89,6 +90,17 @@ public Environment(String name,
         SharedMetricRegistries.add("default", metricRegistry);
     }
 
+    /**
+     * Creates an environment with default health check registry
+     */
+    public Environment(String name,
+                       ObjectMapper objectMapper,
+                       Validator validator,
+                       MetricRegistry metricRegistry,
+                       ClassLoader classLoader) {
+        this(name, objectMapper, validator, metricRegistry, classLoader, new HealthCheckRegistry());
+    }
+
     /**
      * Returns the application's {@link JerseyEnvironment}.
      */

diff --git a/dropwizard-core/src/test/java/io/dropwizard/setup/BootstrapTest.java b/dropwizard-core/src/test/java/io/dropwizard/setup/BootstrapTest.java
@@ -3,6 +3,7 @@
 import com.codahale.metrics.Histogram;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.UniformReservoir;
+import com.codahale.metrics.health.HealthCheckRegistry;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.dropwizard.Application;
 import io.dropwizard.Configuration;
@@ -51,6 +52,12 @@ public void hasAnObjectMapper() throws Exception {
                 .isNotNull();
     }
 
+    @Test
+    public void hasHealthCheckRegistry() {
+        assertThat(bootstrap.getHealthCheckRegistry())
+            .isNotNull();
+    }
+
     @Test
     public void defaultsToUsingFilesForConfiguration() throws Exception {
         assertThat(bootstrap.getConfigurationSourceProvider())
@@ -132,4 +139,12 @@ public void canUseCustomObjectMapper() {
         bootstrap.setObjectMapper(minimalObjectMapper);
         assertThat(bootstrap.getObjectMapper()).isSameAs(minimalObjectMapper);
     }
+
+    @Test
+    public void canUseCustomHealthCheckRegistry() {
+        final HealthCheckRegistry healthCheckRegistry = new HealthCheckRegistry();
+        bootstrap.setHealthCheckRegistry(healthCheckRegistry);
+        assertThat(bootstrap.getHealthCheckRegistry()).isSameAs(healthCheckRegistry);
+    }
+
 }