From c64f97082cc101cb828b039dae5f59e140b7c9e8 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Fri, 13 Jul 2018 11:04:05 +0300
Subject: [PATCH] Propagate statsd to config

---
 Gopkg.lock       |  8 +++++++-
 Gopkg.toml       |  4 ++++
 config/config.go | 44 +++++++++++++++++++++++++++++++++++++++++++-
 main.go          | 26 +++++++++++++++++++++++++-
 4 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/Gopkg.lock b/Gopkg.lock
index 18b8f74..c656a7f 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -100,9 +100,15 @@
   revision = "947dcec5ba9c011838740e680966fd7087a71d0d"
   version = "v2.2.6"
 
+[[projects]]
+  name = "gopkg.in/alexcesaro/statsd.v2"
+  packages = ["."]
+  revision = "7fea3f0d2fab1ad973e641e51dba45443a311a90"
+  version = "v2.0.0"
+
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "f828340a30ea13c563829f9a37d0ff62974d4578411c9be02e61125dbdf98692"
+  inputs-digest = "7fad0f62feb7737b064d85cc4333a1a3e9298faec2afd864b4404f515fc7f17c"
   solver-name = "gps-cdcl"
   solver-version = 1
diff --git a/Gopkg.toml b/Gopkg.toml
index f34b7d2..e61751d 100644
--- a/Gopkg.toml
+++ b/Gopkg.toml
@@ -52,3 +52,7 @@
 [[constraint]]
   name = "github.com/beevik/ntp"
   version = "0.2.0"
+
+[[constraint]]
+  name = "gopkg.in/alexcesaro/statsd.v2"
+  version = "2.0.0"
diff --git a/config/config.go b/config/config.go
index aaab626..d595ee1 100644
--- a/config/config.go
+++ b/config/config.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/juju/errors"
+	statsd "gopkg.in/alexcesaro/statsd.v2"
 )
 
 // Buffer sizes define internal socket buffer sizes.
@@ -32,6 +33,14 @@ type Config struct {
 	PublicIPv6 net.IP
 	StatsIP    net.IP
 
+	StatsD struct {
+		Addr       net.Addr
+		Prefix     string
+		Tags       map[string]string
+		TagsFormat statsd.TagFormat
+		Enabled    bool
+	}
+
 	Secret []byte
 	AdTag  []byte
 }
@@ -109,7 +118,9 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	publicIPv4 net.IP, PublicIPv4Port uint16,
 	publicIPv6 net.IP, publicIPv6Port uint16,
 	statsIP net.IP, statsPort uint16,
-	secret, adtag string) (*Config, error) {
+	secret, adtag string,
+	statsdIP string, statsdPort uint16, statsdNetwork string, statsdPrefix string,
+	statsdTagsFormat string, statsdTags map[string]string) (*Config, error) {
 	secureMode := false
 	if strings.HasPrefix(secret, "dd") && len(secret) == 34 {
 		secureMode = true
@@ -174,5 +185,36 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		SecureMode:     secureMode,
 	}
 
+	if statsdIP != "" {
+		conf.StatsD.Enabled = true
+		conf.StatsD.Prefix = statsdPrefix
+		conf.StatsD.Tags = statsdTags
+
+		var addr net.Addr
+		hostPort := net.JoinHostPort(statsdIP, strconv.Itoa(int(statsdPort)))
+		switch statsdNetwork {
+		case "tcp":
+			addr, err = net.ResolveTCPAddr("tcp", hostPort)
+		case "udp":
+			addr, err = net.ResolveUDPAddr("udp", hostPort)
+		default:
+			err = errors.Errorf("Unknown network %s", statsdNetwork)
+		}
+		if err != nil {
+			return nil, errors.Annotate(err, "Cannot resolve statsd address")
+		}
+		conf.StatsD.Addr = addr
+
+		switch statsdTagsFormat {
+		case "datadog":
+			conf.StatsD.TagsFormat = statsd.Datadog
+		case "influxdb":
+			conf.StatsD.TagsFormat = statsd.InfluxDB
+		case "":
+		default:
+			return nil, errors.Errorf("Unknown tags format %s", statsdTagsFormat)
+		}
+	}
+
 	return conf, nil
 }
diff --git a/main.go b/main.go
index 7587613..da9b5fd 100644
--- a/main.go
+++ b/main.go
@@ -61,7 +61,7 @@ var (
 			Envar("MTG_IPV6_PORT").
 			Uint16()
 
-	statsIP = app.Flag("stats-ip", "Which IP bind stats server to").
+	statsIP = app.Flag("stats-ip", "Which IP bind stats server to.").
 		Short('t').
 		Envar("MTG_STATS_IP").
 		Default("127.0.0.1").
@@ -72,6 +72,28 @@ var (
 			Default("3129").
 			Uint16()
 
+	statsdIP = app.Flag("statsd-ip", "Which IP should we use for working with statsd.").
+			Envar("MTG_STATSD_IP").
+			String()
+	statsdPort = app.Flag("statsd-port", "Which port should we use for working with statsd.").
+			Envar("MTG_STATSD_PORT").
+			Default("8125").
+			Uint16()
+	statsdNetwork = app.Flag("statsd-network", "Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
+			Envar("MTG_STATSD_NETWORK").
+			Default("udp").
+			String()
+	statsdPrefix = app.Flag("statsd-prefix", "Which bucket prefix should we use for sending stats to statsd.").
+			Envar("MTG_STATSD_PREFIX").
+			Default("mtg").
+			String()
+	statsdTagsFormat = app.Flag("statsd-tags-format", "Which tag format should we use to send stats metrics. Valid options are 'datadog' and 'influxdb'.").
+				Envar("MTG_STATSD_TAGS_FORMAT").
+				String()
+	statsdTags = app.Flag("statsd-tags", "Tags to use for working with statsd (specified as 'key=value').").
+			Envar("MTG_STATSD_TAGS").
+			StringMap()
+
 	secret = app.Arg("secret", "Secret of this proxy.").Required().String()
 	adtag  = app.Arg("adtag", "ADTag of the proxy.").String()
 )
@@ -96,6 +118,8 @@ func main() {
 		*publicIPv6, *publicIPv6Port,
 		*statsIP, *statsPort,
 		*secret, *adtag,
+		*statsdIP, *statsdPort, *statsdNetwork, *statsdPrefix,
+		*statsdTagsFormat, *statsdTags,
 	)
 	if err != nil {
 		usage(err.Error())