From e10394ba3ba1c43a4789629d7acd8c0bbb08c6a3 Mon Sep 17 00:00:00 2001
From: Cameron Sparr <cameronsparr@gmail.com>
Date: Tue, 10 Nov 2015 14:19:32 -0700
Subject: [PATCH] Use the UDP client for writing to InfluxDB

---
 outputs/influxdb/influxdb.go | 91 ++++++++++++++++++++----------------
 1 file changed, 52 insertions(+), 39 deletions(-)

diff --git a/outputs/influxdb/influxdb.go b/outputs/influxdb/influxdb.go
index 0dea2af6c..cbdd67c13 100644
--- a/outputs/influxdb/influxdb.go
+++ b/outputs/influxdb/influxdb.go
@@ -15,21 +15,23 @@ import (
 
 type InfluxDB struct {
 	// URL is only for backwards compatability
-	URL       string
-	URLs      []string `toml:"urls"`
-	Username  string
-	Password  string
-	Database  string
-	UserAgent string
-	Precision string
-	Timeout   internal.Duration
+	URL        string
+	URLs       []string `toml:"urls"`
+	Username   string
+	Password   string
+	Database   string
+	UserAgent  string
+	Precision  string
+	Timeout    internal.Duration
+	UDPPayload int
 
 	conns []client.Client
 }
 
 var sampleConfig = `
-  # The full HTTP endpoint URL for your InfluxDB instance
+  # The full HTTP or UDP endpoint URL for your InfluxDB instance
   # Multiple urls can be specified for InfluxDB cluster support.
+  # urls = ["udp://localhost:8089"] # UDP endpoint example
   urls = ["http://localhost:8086"] # required
   # The target database for metrics (telegraf will create it if not exists)
   database = "telegraf" # required
@@ -42,51 +44,62 @@ var sampleConfig = `
   # timeout = "5s"
   # username = "telegraf"
   # password = "metricsmetricsmetricsmetrics"
-  # Set the user agent for the POSTs (can be useful for log differentiation)
+  # Set the user agent for HTTP POSTs (can be useful for log differentiation)
   # user_agent = "telegraf"
+  # Set UDP payload size, defaults to InfluxDB UDP Client default (512 bytes)
+  # udp_payload = 512
 `
 
 func (i *InfluxDB) Connect() error {
-	var urls []*url.URL
-	for _, URL := range i.URLs {
-		u, err := url.Parse(URL)
-		if err != nil {
-			return err
-		}
+	var urls []string
+	for _, u := range i.URLs {
 		urls = append(urls, u)
 	}
 
 	// Backward-compatability with single Influx URL config files
 	// This could eventually be removed in favor of specifying the urls as a list
 	if i.URL != "" {
-		u, err := url.Parse(i.URL)
-		if err != nil {
-			return err
-		}
-		urls = append(urls, u)
+		urls = append(urls, i.URL)
 	}
 
 	var conns []client.Client
-	for _, parsed_url := range urls {
-		c := client.NewClient(client.Config{
-			URL:       parsed_url,
-			Username:  i.Username,
-			Password:  i.Password,
-			UserAgent: i.UserAgent,
-			Timeout:   i.Timeout.Duration,
-		})
-		conns = append(conns, c)
-	}
+	for _, u := range urls {
+		switch {
+		case strings.HasPrefix(u, "udp"):
+			if i.UDPPayload == 0 {
+				i.UDPPayload = client.UDPPayloadSize
+			}
+			c, err := client.NewUDPClient(client.UDPConfig{
+				Addr:        parsed_url.Host,
+				PayloadSize: i.UDPPayload,
+			})
+			if err != nil {
+				return err
+			}
+			conns = append(conns, c)
+		default:
+			// If URL doesn't start with "udp", assume HTTP client
+			c, err := client.NewHTTPClient(client.HTTPConfig{
+				Addr:      parsed_url.String(),
+				Username:  i.Username,
+				Password:  i.Password,
+				UserAgent: i.UserAgent,
+				Timeout:   i.Timeout.Duration,
+			})
+			if err != nil {
+				return err
+			}
 
-	for _, conn := range conns {
-		_, e := conn.Query(client.Query{
-			Command: fmt.Sprintf("CREATE DATABASE %s", i.Database),
-		})
+			// Create Database if it doesn't exist
+			_, e := c.Query(client.Query{
+				Command: fmt.Sprintf("CREATE DATABASE %s", i.Database),
+			})
 
-		if e != nil && !strings.Contains(e.Error(), "database already exists") {
-			log.Println("Database creation failed: " + e.Error())
-		} else {
-			break
+			if e != nil && !strings.Contains(e.Error(), "database already exists") {
+				log.Println("Database creation failed: " + e.Error())
+			}
+
+			conns = append(conns, c)
 		}
 	}