From 2aff1c2c3f1fd64b7eb5f4da29bf3eb365eb1c97 Mon Sep 17 00:00:00 2001
From: Olivier Vielpeau <olivier.vielpeau@datadoghq.com>
Date: Mon, 21 Mar 2016 17:52:16 -0400
Subject: [PATCH] [reporter][statsd] Allow configuration of statsd host

---
 .../jmxfetch/reporter/ReporterFactory.java    | 11 ++++++-
 .../jmxfetch/reporter/StatsdReporter.java     | 10 +++++--
 .../jmxfetch/validator/ReporterValidator.java |  5 ++--
 .../jmxfetch/TestParsingJCommander.java       | 29 ++++++++++++++++++-
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/main/java/org/datadog/jmxfetch/reporter/ReporterFactory.java b/src/main/java/org/datadog/jmxfetch/reporter/ReporterFactory.java
index 934052033..d488ddb50 100644
--- a/src/main/java/org/datadog/jmxfetch/reporter/ReporterFactory.java
+++ b/src/main/java/org/datadog/jmxfetch/reporter/ReporterFactory.java
@@ -1,5 +1,8 @@
 package org.datadog.jmxfetch.reporter;
 
+import com.google.common.base.Joiner;
+import java.util.Arrays;
+
 public class ReporterFactory {
 
     public static Reporter getReporter(String type) {
@@ -9,7 +12,13 @@ public static Reporter getReporter(String type) {
         if ("console".equals(type)) {
             return new ConsoleReporter();
         } else if (type.startsWith("statsd:")) {
-            return new StatsdReporter(Integer.valueOf(type.split(":")[1]));
+            String[] typeElements = type.split(":");
+            String host = "localhost";
+            Integer port = Integer.valueOf(typeElements[typeElements.length - 1]);
+            if (typeElements.length > 2) {
+                host = Joiner.on(":").join(Arrays.copyOfRange(typeElements, 1, typeElements.length - 1));
+            }
+            return new StatsdReporter(host, port);
         } else {
             throw new IllegalArgumentException("Invalid reporter type: " + type);
         }
diff --git a/src/main/java/org/datadog/jmxfetch/reporter/StatsdReporter.java b/src/main/java/org/datadog/jmxfetch/reporter/StatsdReporter.java
index ae1b7709b..320d64a73 100644
--- a/src/main/java/org/datadog/jmxfetch/reporter/StatsdReporter.java
+++ b/src/main/java/org/datadog/jmxfetch/reporter/StatsdReporter.java
@@ -10,17 +10,19 @@
 public class StatsdReporter extends Reporter {
 
     private StatsDClient statsDClient;
+    private String statsdHost;
     private int statsdPort;
     private long initializationTime;
 
-    public StatsdReporter(int statsdPort) {
+    public StatsdReporter(String statsdHost, int statsdPort) {
+        this.statsdHost = statsdHost;
         this.statsdPort = statsdPort;
         this.init();
     }
 
     private void init() {
         initializationTime = System.currentTimeMillis();
-        statsDClient = new NonBlockingStatsDClient(null, "localhost", this.statsdPort, new String[]{});
+        statsDClient = new NonBlockingStatsDClient(null, this.statsdHost, this.statsdPort, new String[]{});
     }
 
     protected void sendMetricPoint(String metricName, double value, String[] tags) {
@@ -71,6 +73,10 @@ public void displayInstanceName(Instance instance) {
         throw new UnsupportedOperationException();
     }
 
+    public String getStatsdHost() {
+        return statsdHost;
+    }
+
     public int getStatsdPort() {
         return statsdPort;
     }
diff --git a/src/main/java/org/datadog/jmxfetch/validator/ReporterValidator.java b/src/main/java/org/datadog/jmxfetch/validator/ReporterValidator.java
index 4ff46b0b2..e33210371 100644
--- a/src/main/java/org/datadog/jmxfetch/validator/ReporterValidator.java
+++ b/src/main/java/org/datadog/jmxfetch/validator/ReporterValidator.java
@@ -10,7 +10,8 @@ public class ReporterValidator implements IParameterValidator {
 
     public void validate(String name, String value) throws ParameterException {
         if (value.startsWith(STATSD_PREFIX) && value.length() > STATSD_PREFIX.length()) {
-            String port = new String(value.split(":")[1]);
+            String[] splitValue = value.split(":");
+            String port = splitValue[splitValue.length - 1];
             try {
                 positiveIntegerValidator.validate(name, port);
             } catch (ParameterException pe) {
@@ -19,7 +20,7 @@ public void validate(String name, String value) throws ParameterException {
             return;
         }
         if (!value.equals("console")) {
-            throw new ParameterException("Parameter " + name + " should be either 'console' or 'statsd:[STATSD_PORT]'");
+            throw new ParameterException("Parameter " + name + " should be either 'console', 'statsd:[STATSD_PORT]' or 'statsd:[STATSD_HOST]:[STATSD_PORT]'");
         }
     }
 }
\ No newline at end of file
diff --git a/src/test/java/org/datadog/jmxfetch/TestParsingJCommander.java b/src/test/java/org/datadog/jmxfetch/TestParsingJCommander.java
index 2c6eb8389..d6827ae01 100644
--- a/src/test/java/org/datadog/jmxfetch/TestParsingJCommander.java
+++ b/src/test/java/org/datadog/jmxfetch/TestParsingJCommander.java
@@ -122,6 +122,33 @@ public void testParsingReporter() {
         appConfig = testCommand(params);
         assertNotNull(appConfig.getReporter());
         assertTrue(appConfig.getReporter() instanceof StatsdReporter);
+        assertEquals("localhost", ((StatsdReporter) appConfig.getReporter()).getStatsdHost());
+        assertEquals(10, ((StatsdReporter) appConfig.getReporter()).getStatsdPort());
+
+        // statsd reporter with custom ipv4 host
+        params = new String[]{
+                "--reporter", "statsd:127.0.0.1:10",
+                "--check", SINGLE_CHECK,
+                "--conf_directory", CONF_DIR,
+                AppConfig.ACTION_COLLECT
+        };
+        appConfig = testCommand(params);
+        assertNotNull(appConfig.getReporter());
+        assertTrue(appConfig.getReporter() instanceof StatsdReporter);
+        assertEquals("127.0.0.1", ((StatsdReporter) appConfig.getReporter()).getStatsdHost());
+        assertEquals(10, ((StatsdReporter) appConfig.getReporter()).getStatsdPort());
+
+        // statsd reporter with custom ipv6 host
+        params = new String[]{
+                "--reporter", "statsd:[::1]:10",
+                "--check", SINGLE_CHECK,
+                "--conf_directory", CONF_DIR,
+                AppConfig.ACTION_COLLECT
+        };
+        appConfig = testCommand(params);
+        assertNotNull(appConfig.getReporter());
+        assertTrue(appConfig.getReporter() instanceof StatsdReporter);
+        assertEquals("[::1]", ((StatsdReporter) appConfig.getReporter()).getStatsdHost());
         assertEquals(10, ((StatsdReporter) appConfig.getReporter()).getStatsdPort());
 
         // invalid reporter
@@ -135,7 +162,7 @@ public void testParsingReporter() {
             testCommand(params);
             fail("Should have failed because reporter is invalid");
         } catch (ParameterException pe) {
-            assertEquals("Parameter --reporter should be either 'console' or 'statsd:[STATSD_PORT]'", pe.getMessage());
+            assertEquals("Parameter --reporter should be either 'console', 'statsd:[STATSD_PORT]' or 'statsd:[STATSD_HOST]:[STATSD_PORT]'", pe.getMessage());
         }
 
         // invalid port