From 63085fb441a031ded799e4291c908d9baf916f59 Mon Sep 17 00:00:00 2001
From: Nick Pillitteri <nickp@smartertravelmedia.com>
Date: Mon, 8 Aug 2016 14:13:04 -0400
Subject: [PATCH] Allow injection of StatsDClient into StatsdMetricWriter

Allow an instance of StatsDClient to be injected into the StatsdMetricWriter
which is used for exporting metrics to a Statsd server. This new constructor
allows the client to be injected but does not change the default behavior of
the writer.
---
 .../metrics/statsd/StatsdMetricWriter.java    | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/statsd/StatsdMetricWriter.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/statsd/StatsdMetricWriter.java
index 78578a2d98..ae09545170 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/statsd/StatsdMetricWriter.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/statsd/StatsdMetricWriter.java
@@ -19,6 +19,7 @@ package org.springframework.boot.actuate.metrics.statsd;
 import java.io.Closeable;
 
 import com.timgroup.statsd.NonBlockingStatsDClient;
+import com.timgroup.statsd.StatsDClient;
 import com.timgroup.statsd.StatsDClientErrorHandler;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -26,6 +27,7 @@ import org.apache.commons.logging.LogFactory;
 import org.springframework.boot.actuate.metrics.Metric;
 import org.springframework.boot.actuate.metrics.writer.Delta;
 import org.springframework.boot.actuate.metrics.writer.MetricWriter;
+import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 
 /**
@@ -43,7 +45,7 @@ public class StatsdMetricWriter implements MetricWriter, Closeable {
 
 	private static final Log logger = LogFactory.getLog(StatsdMetricWriter.class);
 
-	private final NonBlockingStatsDClient client;
+	private final StatsDClient client;
 
 	/**
 	 * Create a new writer instance with the given parameters.
@@ -61,12 +63,26 @@ public class StatsdMetricWriter implements MetricWriter, Closeable {
 	 * @param port the port for the statsd server
 	 */
 	public StatsdMetricWriter(String prefix, String host, int port) {
-		prefix = StringUtils.hasText(prefix) ? prefix : null;
-		while (prefix != null && prefix.endsWith(".")) {
-			prefix = prefix.substring(0, prefix.length() - 1);
+		this(new NonBlockingStatsDClient(trimPrefix(prefix), host, port,
+			new LoggingStatsdErrorHandler()));
+	}
+
+	/**
+	 * Create a new writer with the given client.
+	 * @param client StatsD client to write metrics with
+	 */
+	public StatsdMetricWriter(StatsDClient client) {
+		Assert.notNull(client);
+		this.client = client;
+	}
+
+	private static String trimPrefix(String prefix) {
+		String trimmedPrefix = StringUtils.hasText(prefix) ? prefix : null;
+		while (trimmedPrefix != null && trimmedPrefix.endsWith(".")) {
+			trimmedPrefix = trimmedPrefix.substring(0, trimmedPrefix.length() - 1);
 		}
-		this.client = new NonBlockingStatsDClient(prefix, host, port,
-				new LoggingStatsdErrorHandler());
+
+		return trimmedPrefix;
 	}
 
 	@Override