Merge pull request #1095 from dropwizard/shared_health_checks

Add support for the default shared health check registry name

metrics-healthchecks/src/main/java/com/codahale/metrics/health/SharedHealthCheckRegistries.java

@@ -1,10 +1,9 @@
 package com.codahale.metrics.health;
 
-import com.codahale.metrics.MetricRegistry;
-
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * A map of shared, named health registries.
@@ -13,6 +12,13 @@ public class SharedHealthCheckRegistries {
     private static final ConcurrentMap<String, HealthCheckRegistry> REGISTRIES =
             new ConcurrentHashMap<String, HealthCheckRegistry>();
 
+    private static AtomicReference<String> defaultRegistryName = new AtomicReference<String>();
+
+    /* Visible for testing */
+    static void setDefaultRegistryName(AtomicReference<String> defaultRegistryName) {
+        SharedHealthCheckRegistries.defaultRegistryName = defaultRegistryName;
+    }
+
     private SharedHealthCheckRegistries() { /* singleton */ }
 
     public static void clear() {
@@ -43,4 +49,45 @@ public static HealthCheckRegistry getOrCreate(String name) {
         }
         return existing;
     }
+
+    /**
+     * Creates a new registry and sets it as the default one under the provided name.
+     *
+     * @param name the registry name
+     * @return the default registry
+     * @throws IllegalStateException if the name has already been set
+     */
+    public synchronized static HealthCheckRegistry setDefault(String name) {
+        final HealthCheckRegistry registry = getOrCreate(name);
+        return setDefault(name, registry);
+    }
+
+    /**
+     * Sets the provided registry as the default one under the provided name
+     *
+     * @param name                the default registry name
+     * @param healthCheckRegistry the default registry
+     * @throws IllegalStateException if the default registry has already been set
+     */
+    public static HealthCheckRegistry setDefault(String name, HealthCheckRegistry healthCheckRegistry) {
+        if (defaultRegistryName.compareAndSet(null, name)) {
+            add(name, healthCheckRegistry);
+            return healthCheckRegistry;
+        }
+        throw new IllegalStateException("Default health check registry is already set.");
+    }
+
+    /**
+     * Gets the name of the default registry, if it has been set
+     *
+     * @return the default registry
+     * @throws IllegalStateException if the default has not been set
+     */
+    public static HealthCheckRegistry getDefault() {
+        final String name = defaultRegistryName.get();
+        if (name != null) {
+            return getOrCreate(name);
+        }
+        throw new IllegalStateException("Default registry name has not been set.");
+    }
 }

metrics-healthchecks/src/test/java/com/codahale/metrics/health/SharedHealthCheckRegistriesTest.java

@@ -0,0 +1,97 @@
+package com.codahale.metrics.health;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class SharedHealthCheckRegistriesTest {
+
+    @Rule
+    public ExpectedException expectedException = ExpectedException.none();
+
+    @Before
+    public void setUp() {
+        SharedHealthCheckRegistries.setDefaultRegistryName(new AtomicReference<String>());
+        SharedHealthCheckRegistries.clear();
+    }
+
+    @Test
+    public void savesCreatedRegistry() {
+        final HealthCheckRegistry one = SharedHealthCheckRegistries.getOrCreate("db");
+        final HealthCheckRegistry two = SharedHealthCheckRegistries.getOrCreate("db");
+
+        assertThat(one).isSameAs(two);
+    }
+
+    @Test
+    public void returnsSetOfCreatedRegistries() {
+        SharedHealthCheckRegistries.getOrCreate("db");
+
+        assertThat(SharedHealthCheckRegistries.names()).containsOnly("db");
+    }
+
+    @Test
+    public void registryCanBeRemoved() {
+        final HealthCheckRegistry first = SharedHealthCheckRegistries.getOrCreate("db");
+        SharedHealthCheckRegistries.remove("db");
+
+        assertThat(SharedHealthCheckRegistries.names()).isEmpty();
+        assertThat(SharedHealthCheckRegistries.getOrCreate("db")).isNotEqualTo(first);
+    }
+
+    @Test
+    public void registryCanBeCleared() {
+        SharedHealthCheckRegistries.getOrCreate("db");
+        SharedHealthCheckRegistries.getOrCreate("web");
+
+        SharedHealthCheckRegistries.clear();
+
+        assertThat(SharedHealthCheckRegistries.names()).isEmpty();
+    }
+
+    @Test
+    public void defaultRegistryIsNotSetByDefault() {
+        expectedException.expect(IllegalStateException.class);
+        expectedException.expectMessage("Default registry name has not been set.");
+
+        SharedHealthCheckRegistries.getDefault();
+    }
+
+    @Test
+    public void defaultRegistryCanBeSet() {
+        HealthCheckRegistry registry = SharedHealthCheckRegistries.setDefault("default");
+
+        assertThat(SharedHealthCheckRegistries.getDefault()).isEqualTo(registry);
+    }
+
+    @Test
+    public void specificRegistryCanBeSetAsDefault() {
+        HealthCheckRegistry registry = new HealthCheckRegistry();
+        SharedHealthCheckRegistries.setDefault("default", registry);
+
+        assertThat(SharedHealthCheckRegistries.getDefault()).isEqualTo(registry);
+    }
+
+    @Test
+    public void unableToSetDefaultRegistryTwice() {
+        expectedException.expect(IllegalStateException.class);
+        expectedException.expectMessage("Default health check registry is already set.");
+
+        SharedHealthCheckRegistries.setDefault("default");
+        SharedHealthCheckRegistries.setDefault("default");
+    }
+
+    @Test
+    public void unableToSetCustomDefaultRegistryTwice() {
+        expectedException.expect(IllegalStateException.class);
+        expectedException.expectMessage("Default health check registry is already set.");
+
+        SharedHealthCheckRegistries.setDefault("default", new HealthCheckRegistry());
+        SharedHealthCheckRegistries.setDefault("default", new HealthCheckRegistry());
+    }
+}