Enable custom Reservoir with Dropwizard metrics

Uses the ReservoirFactory to customize the implementation of
the Reservoir that will be used when creating Timer and Histogram
in the DropwizardMetricServices.

Fixes gh-5199
Closes gh-7105

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.annotation.Autowired;
 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,9 @@ import org.springframework.context.annotation.Configuration;
 @AutoConfigureBefore(MetricRepositoryAutoConfiguration.class)
 public class MetricsDropwizardAutoConfiguration {
 
+	@Autowired(required = false)
+	private ReservoirFactory reservoirFactory;
+
 	@Bean
 	@ConditionalOnMissingBean
 	public MetricRegistry metricRegistry() {
@@ -52,7 +57,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,7 +24,9 @@ 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;
@@ -53,6 +55,8 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 
 	private final MetricRegistry registry;
 
+	private ReservoirFactory reservoirFactory;
+
 	private final ConcurrentMap<String, SimpleGauge> gauges = new ConcurrentHashMap<String, SimpleGauge>();
 
 	private final ConcurrentHashMap<String, String> names = new ConcurrentHashMap<String, String>();
@@ -65,6 +69,18 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 		this.registry = registry;
 	}
 
+	/**
+	 * 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;
+	}
+
 	@Override
 	public void increment(String name) {
 		incrementInternal(name, 1L);
@@ -91,12 +107,12 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 	public void submit(String name, double value) {
 		if (name.startsWith("histogram")) {
 			long longValue = (long) value;
-			Histogram metric = this.registry.histogram(name);
+			Histogram metric = registerHistogram(name);
 			metric.update(longValue);
 		}
 		else if (name.startsWith("timer")) {
 			long longValue = (long) value;
-			Timer metric = this.registry.timer(name);
+			Timer metric = registerTimer(name);
 			metric.update(longValue, TimeUnit.MILLISECONDS);
 		}
 		else {
@@ -105,6 +121,43 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 		}
 	}
 
+	private Histogram registerHistogram(String name) {
+		if (this.reservoirFactory == null) {
+			return this.registry.histogram(name);
+		}
+		else {
+			Histogram histogram = new Histogram(this.reservoirFactory.getObject());
+			return getOrAddMetric(name, histogram);
+		}
+	}
+
+	private Timer registerTimer(String name) {
+		if (this.reservoirFactory == null) {
+			return this.registry.timer(name);
+		}
+		else {
+			Timer timer = new Timer(this.reservoirFactory.getObject());
+			return getOrAddMetric(name, timer);
+		}
+	}
+
+	@SuppressWarnings("unchecked")
+	private <T extends Metric> T getOrAddMetric(String name, T newMetric) {
+		Metric metric = this.registry.getMetrics().get(name);
+		if (metric == null) {
+			return this.registry.register(name, newMetric);
+		}
+		else {
+			if (metric.getClass().equals(newMetric.getClass())) {
+				return (T) metric;
+			}
+			else {
+				throw new IllegalArgumentException(
+						name + " is already used for a different type of metric");
+			}
+		}
+	}
+
 	private void setGaugeValue(String name, double value) {
 		// NOTE: Dropwizard provides no way to do this atomically
 		SimpleGauge gauge = this.gauges.get(name);
@@ -148,6 +201,10 @@ public class DropwizardMetricServices implements CounterService, GaugeService {
 		this.registry.remove(name);
 	}
 
+	void setReservoirFactory(ReservoirFactory reservoirFactory) {
+		this.reservoirFactory = reservoirFactory;
+	}
+
 	/**
 	 * Simple {@link Gauge} implementation to {@literal double} value.
 	 */

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

+/*
+ * 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;
+
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.ObjectFactory;
+
+/**
+ * A {@link Reservoir} factory to instantiate the Reservoir that will be set as default
+ * for the {@link DropwizardMetricServices}.
+ * The Reservoir instances can't be shared across {@link com.codahale.metrics.Metric}.
+ *
+ * @author Lucas Saldanha
+ */
+public abstract class ReservoirFactory implements ObjectFactory<Reservoir> {
+
+	protected abstract Reservoir defaultReservoir();
+
+	@Override
+	public Reservoir getObject() throws BeansException {
+		return defaultReservoir();
+	}
+}

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

+/*
+ * 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);
+
+		assertThat(ReflectionTestUtils.getField(dropwizardMetricServices, "reservoirFactory"))
+				.isNull();
+	}
+
+	@Test
+	public void dropwizardWithCustomReservoirConfigured() {
+		this.context = new AnnotationConfigApplicationContext(
+				MetricsDropwizardAutoConfiguration.class, Config.class);
+
+		DropwizardMetricServices dropwizardMetricServices = this.context
+				.getBean(DropwizardMetricServices.class);
+
+		assertThat(ReflectionTestUtils.getField(dropwizardMetricServices, "reservoirFactory"))
+				.isNotNull();
+	}
+
+	@Configuration
+	static class Config {
+
+		@Bean
+		public ReservoirFactory reservoirFactory() {
+			return new ReservoirFactory() {
+				@Override
+				protected Reservoir defaultReservoir() {
+					return new UniformReservoir();
+				}
+			};
+		}
+	}
+
+}

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

@@ -20,15 +20,22 @@ 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.Reservoir;
+import com.codahale.metrics.Timer;
+import com.codahale.metrics.UniformReservoir;
 import org.junit.Test;
 
+import org.springframework.test.util.ReflectionTestUtils;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * Tests for {@link DropwizardMetricServices}.
  *
  * @author Dave Syer
+ * @author Lucas Saldanha
  */
 public class DropwizardMetricServicesTests {
 
@@ -78,6 +85,26 @@ public class DropwizardMetricServicesTests {
 		assertThat(this.registry.timer("timer.foo").getCount()).isEqualTo(2);
 	}
 
+	@Test
+	public void setCustomReservoirTimer() {
+		this.writer.setReservoirFactory(new ReservoirFactory() {
+			@Override
+			protected Reservoir defaultReservoir() {
+				return 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 +112,23 @@ public class DropwizardMetricServicesTests {
 		assertThat(this.registry.histogram("histogram.foo").getCount()).isEqualTo(2);
 	}
 
+	@Test
+	public void setCustomReservoirHistogram() {
+		this.writer.setReservoirFactory(new ReservoirFactory() {
+			@Override
+			protected Reservoir defaultReservoir() {
+				return 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