Merge pull request #7105 from lucassaldanha/gh-5199

* pr/7105:
  Polish ReservoirFactory support
  Enable custom Reservoir with Dropwizard metrics

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/MetricsDropwizardAutoConfiguration.java

@@ -18,10 +18,12 @@ package org.springframework.boot.actuate.autoconfigure;
 
 import com.codahale.metrics.MetricRegistry;
 
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.actuate.endpoint.MetricReaderPublicMetrics;
 import org.springframework.boot.actuate.metrics.CounterService;
 import org.springframework.boot.actuate.metrics.GaugeService;
 import org.springframework.boot.actuate.metrics.dropwizard.DropwizardMetricServices;
+import org.springframework.boot.actuate.metrics.dropwizard.ReservoirFactory;
 import org.springframework.boot.actuate.metrics.reader.MetricRegistryMetricReader;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -41,6 +43,13 @@ import org.springframework.context.annotation.Configuration;
 @AutoConfigureBefore(MetricRepositoryAutoConfiguration.class)
 public class MetricsDropwizardAutoConfiguration {
 
+	private final ReservoirFactory reservoirFactory;
+
+	public MetricsDropwizardAutoConfiguration(
+			ObjectProvider<ReservoirFactory> reservoirFactory) {
+		this.reservoirFactory = reservoirFactory.getIfAvailable();
+	}
+
 	@Bean
 	@ConditionalOnMissingBean
 	public MetricRegistry metricRegistry() {
@@ -52,7 +61,12 @@ public class MetricsDropwizardAutoConfiguration {
 			GaugeService.class })
 	public DropwizardMetricServices dropwizardMetricServices(
 			MetricRegistry metricRegistry) {
-		return new DropwizardMetricServices(metricRegistry);
+		if (this.reservoirFactory == null) {
+			return new DropwizardMetricServices(metricRegistry);
+		}
+		else {
+			return new DropwizardMetricServices(metricRegistry, this.reservoirFactory);
+		}
 	}
 
 	@Bean

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/dropwizard/DropwizardMetricServices.java

@@ -24,11 +24,15 @@ import com.codahale.metrics.Counter;
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Histogram;
 import com.codahale.metrics.Meter;
+import com.codahale.metrics.Metric;
 import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.Reservoir;
 import com.codahale.metrics.Timer;
 
 import org.springframework.boot.actuate.metrics.CounterService;
 import org.springframework.boot.actuate.metrics.GaugeService;
+import org.springframework.core.ResolvableType;
+import org.springframework.util.Assert;
 
 /**
  * A {@link GaugeService} and {@link CounterService} that sends data to a Dropwizard
@@ -53,6 +57,8 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 
 	private final MetricRegistry registry;
 
+	private final ReservoirFactory reservoirFactory;
+
 	private final ConcurrentMap<String, SimpleGauge> gauges = new ConcurrentHashMap<String, SimpleGauge>();
 
 	private final ConcurrentHashMap<String, String> names = new ConcurrentHashMap<String, String>();
@@ -62,7 +68,20 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 	 * @param registry the underlying metric registry
 	 */
 	public DropwizardMetricServices(MetricRegistry registry) {
+		this(registry, null);
+	}
+
+	/**
+	 * Create a new {@link DropwizardMetricServices} instance.
+	 * @param registry the underlying metric registry
+	 * @param reservoirFactory the factory that instantiates the {@link Reservoir} that
+	 * will be used on Timers and Histograms
+	 */
+	public DropwizardMetricServices(MetricRegistry registry,
+			ReservoirFactory reservoirFactory) {
 		this.registry = registry;
+		this.reservoirFactory = (reservoirFactory == null ? ReservoirFactory.NONE
+				: reservoirFactory);
 	}
 
 	@Override
@@ -90,14 +109,10 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 	@Override
 	public void submit(String name, double value) {
 		if (name.startsWith("histogram")) {
-			long longValue = (long) value;
-			Histogram metric = this.registry.histogram(name);
-			metric.update(longValue);
+			submitHistogram(name, value);
 		}
 		else if (name.startsWith("timer")) {
-			long longValue = (long) value;
-			Timer metric = this.registry.timer(name);
-			metric.update(longValue, TimeUnit.MILLISECONDS);
+			submitTimer(name, value);
 		}
 		else {
 			name = wrapGaugeName(name);
@@ -105,6 +120,39 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 		}
 	}
 
+	private void submitTimer(String name, double value) {
+		long longValue = (long) value;
+		Timer metric = register(name, new TimerMetricRegistrar());
+		metric.update(longValue, TimeUnit.MILLISECONDS);
+	}
+
+	private void submitHistogram(String name, double value) {
+		long longValue = (long) value;
+		Histogram metric = register(name, new HistogramMetricRegistrar());
+		metric.update(longValue);
+	}
+
+	@SuppressWarnings("unchecked")
+	private <T extends Metric> T register(String name, MetricRegistrar<T> registrar) {
+		Reservoir reservoir = this.reservoirFactory.getReservoir(name);
+		if (reservoir == null) {
+			return registrar.register(this.registry, name);
+		}
+		Metric metric = this.registry.getMetrics().get(name);
+		if (metric != null) {
+			registrar.checkExisting(metric);
+			return (T) metric;
+		}
+		try {
+			return this.registry.register(name, registrar.createForReservoir(reservoir));
+		}
+		catch (IllegalArgumentException ex) {
+			Metric added = this.registry.getMetrics().get(name);
+			registrar.checkExisting(metric);
+			return (T) added;
+		}
+	}
+
 	private void setGaugeValue(String name, double value) {
 		// NOTE: Dropwizard provides no way to do this atomically
 		SimpleGauge gauge = this.gauges.get(name);
@@ -170,4 +218,62 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 
 	}
 
+	/**
+	 * Strategy used to register metrics.
+	 */
+	private static abstract class MetricRegistrar<T extends Metric> {
+
+		private final Class<T> type;
+
+		@SuppressWarnings("unchecked")
+		MetricRegistrar() {
+			this.type = (Class<T>) ResolvableType
+					.forClass(MetricRegistrar.class, getClass()).resolveGeneric();
+		}
+
+		public void checkExisting(Metric metric) {
+			Assert.isInstanceOf(this.type, metric,
+					"Different metric type already registered");
+		}
+
+		protected abstract T register(MetricRegistry registry, String name);
+
+		protected abstract T createForReservoir(Reservoir reservoir);
+
+	}
+
+	/**
+	 * {@link MetricRegistrar} for {@link Timer} metrics.
+	 */
+	private static class TimerMetricRegistrar extends MetricRegistrar<Timer> {
+
+		@Override
+		protected Timer register(MetricRegistry registry, String name) {
+			return registry.timer(name);
+		}
+
+		@Override
+		protected Timer createForReservoir(Reservoir reservoir) {
+			return new Timer(reservoir);
+		}
+
+	}
+
+	/**
+	 * {@link MetricRegistrar} for {@link Histogram} metrics.
+	 */
+	private static class HistogramMetricRegistrar extends MetricRegistrar<Histogram> {
+
+		@Override
+		protected Histogram register(MetricRegistry registry, String name) {
+			return registry.histogram(name);
+		}
+
+		@Override
+		protected Histogram createForReservoir(Reservoir reservoir) {
+			return new Histogram(reservoir);
+		}
+
+	}
+
 }

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/dropwizard/ReservoirFactory.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2012-2016 the original author or authors.
+ *
+ * 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 org.springframework.boot.actuate.metrics.dropwizard;
+
+import com.codahale.metrics.Reservoir;
+
+/**
+ * Factory interface that can be used by {@link DropwizardMetricServices} to create a
+ * custom {@link Reservoir}.
+ *
+ * @author Lucas Saldanha
+ * @author Phillip Webb
+ * @since 1.5.0
+ */
+public interface ReservoirFactory {
+
+	/**
+	 * Default empty {@link ReservoirFactory} implementation.
+	 */
+	ReservoirFactory NONE = new ReservoirFactory() {
+
+		@Override
+		public Reservoir getReservoir(String name) {
+			return null;
+		}
+
+	};
+
+	/**
+	 * Return the {@link Reservoir} instance to use or {@code null} if a custom reservoir
+	 * is not needed.
+	 * @param name the name of the metric
+	 * @return a reservoir instance or {@code null}
+	 */
+	Reservoir getReservoir(String name);
+
+}

spring-boot-actuator/src/test/java/org/springframework/boot/actuate/autoconfigure/MetricsDropwizardAutoConfigurationTests.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2012-2016 the original author or authors.
+ *
+ * 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 org.springframework.boot.actuate.autoconfigure;
+
+import com.codahale.metrics.Reservoir;
+import com.codahale.metrics.UniformReservoir;
+import org.junit.After;
+import org.junit.Test;
+
+import org.springframework.boot.actuate.metrics.dropwizard.DropwizardMetricServices;
+import org.springframework.boot.actuate.metrics.dropwizard.ReservoirFactory;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link MetricsDropwizardAutoConfiguration}.
+ *
+ * @author Lucas Saldanha
+ */
+public class MetricsDropwizardAutoConfigurationTests {
+
+	private AnnotationConfigApplicationContext context;
+
+	@After
+	public void after() {
+		if (this.context != null) {
+			this.context.close();
+		}
+	}
+
+	@Test
+	public void dropwizardWithoutCustomReservoirConfigured() {
+		this.context = new AnnotationConfigApplicationContext(
+				MetricsDropwizardAutoConfiguration.class);
+		DropwizardMetricServices dropwizardMetricServices = this.context
+				.getBean(DropwizardMetricServices.class);
+		ReservoirFactory reservoirFactory = (ReservoirFactory) ReflectionTestUtils
+				.getField(dropwizardMetricServices, "reservoirFactory");
+		assertThat(reservoirFactory.getReservoir("test")).isNull();
+	}
+
+	@Test
+	public void dropwizardWithCustomReservoirConfigured() {
+		this.context = new AnnotationConfigApplicationContext(
+				MetricsDropwizardAutoConfiguration.class, Config.class);
+		DropwizardMetricServices dropwizardMetricServices = this.context
+				.getBean(DropwizardMetricServices.class);
+		ReservoirFactory reservoirFactory = (ReservoirFactory) ReflectionTestUtils
+				.getField(dropwizardMetricServices, "reservoirFactory");
+		assertThat(reservoirFactory.getReservoir("test"))
+				.isInstanceOf(UniformReservoir.class);
+	}
+
+	@Configuration
+	static class Config {
+
+		@Bean
+		public ReservoirFactory reservoirFactory() {
+			return new UniformReservoirFactory();
+		}
+
+	}
+
+	private static class UniformReservoirFactory implements ReservoirFactory {
+
+		@Override
+		public Reservoir getReservoir(String name) {
+			return new UniformReservoir();
+		}
+
+	}
+
+}

spring-boot-actuator/src/test/java/org/springframework/boot/actuate/metrics/dropwizard/DropwizardMetricServicesTests.java

@@ -20,22 +20,41 @@ import java.util.ArrayList;
 import java.util.List;
 
 import com.codahale.metrics.Gauge;
+import com.codahale.metrics.Histogram;
 import com.codahale.metrics.MetricRegistry;
+import com.codahale.metrics.Timer;
+import com.codahale.metrics.UniformReservoir;
+import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import org.springframework.test.util.ReflectionTestUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Matchers.anyString;
 
 /**
  * Tests for {@link DropwizardMetricServices}.
  *
  * @author Dave Syer
+ * @author Lucas Saldanha
  */
 public class DropwizardMetricServicesTests {
 
-	private final MetricRegistry registry = new MetricRegistry();
+	private MetricRegistry registry = new MetricRegistry();
+
+	@Mock
+	private ReservoirFactory reservoirFactory;
 
-	private final DropwizardMetricServices writer = new DropwizardMetricServices(
-			this.registry);
+	private DropwizardMetricServices writer;
+
+	@Before
+	public void setup() {
+		MockitoAnnotations.initMocks(this);
+		this.writer = new DropwizardMetricServices(this.registry, this.reservoirFactory);
+	}
 
 	@Test
 	public void incrementCounter() {
@@ -78,6 +97,20 @@ public class DropwizardMetricServicesTests {
 		assertThat(this.registry.timer("timer.foo").getCount()).isEqualTo(2);
 	}
 
+	@Test
+	public void setCustomReservoirTimer() {
+		given(this.reservoirFactory.getReservoir(anyString()))
+				.willReturn(new UniformReservoir());
+		this.writer.submit("timer.foo", 200);
+		this.writer.submit("timer.foo", 300);
+		assertThat(this.registry.timer("timer.foo").getCount()).isEqualTo(2);
+		Timer timer = (Timer) this.registry.getMetrics().get("timer.foo");
+		Histogram histogram = (Histogram) ReflectionTestUtils.getField(timer,
+				"histogram");
+		assertThat(ReflectionTestUtils.getField(histogram, "reservoir").getClass()
+				.equals(UniformReservoir.class)).isTrue();
+	}
+
 	@Test
 	public void setPredefinedHistogram() {
 		this.writer.submit("histogram.foo", 2.1);
@@ -85,6 +118,18 @@ public class DropwizardMetricServicesTests {
 		assertThat(this.registry.histogram("histogram.foo").getCount()).isEqualTo(2);
 	}
 
+	@Test
+	public void setCustomReservoirHistogram() {
+		given(this.reservoirFactory.getReservoir(anyString()))
+				.willReturn(new UniformReservoir());
+		this.writer.submit("histogram.foo", 2.1);
+		this.writer.submit("histogram.foo", 2.3);
+		assertThat(this.registry.histogram("histogram.foo").getCount()).isEqualTo(2);
+		assertThat(ReflectionTestUtils
+				.getField(this.registry.getMetrics().get("histogram.foo"), "reservoir")
+				.getClass().equals(UniformReservoir.class)).isTrue();
+	}
+
 	/**
 	 * Test the case where a given writer is used amongst several threads where each
 	 * thread is updating the same set of metrics. This would be an example case of the