From 6efe91ea9cece66e864b9e472f48811eaf61365a Mon Sep 17 00:00:00 2001
From: Cameron Sparr <cameronsparr@gmail.com>
Date: Sun, 10 Jul 2016 14:47:47 +0100
Subject: [PATCH] prometheus_client, implement Collector interface

closes #1334
---
 CHANGELOG.md                                  |  1 +
 .../prometheus_client/prometheus_client.go    | 70 ++++++++++---------
 2 files changed, 38 insertions(+), 33 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b988508ae..9e4c9a968 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,7 @@ should now look like:
 - [#1378](https://github.com/influxdata/telegraf/issues/1378): Trim BOM from config file for Windows support.
 - [#1339](https://github.com/influxdata/telegraf/issues/1339): Prometheus client output panic on service reload.
 - [#1461](https://github.com/influxdata/telegraf/pull/1461): Prometheus parser, protobuf format header fix.
+- [#1334](https://github.com/influxdata/telegraf/issues/1334): Prometheus output, metric refresh and caching fixes.
 
 ## v1.0 beta 2 [2016-06-21]
 
diff --git a/plugins/outputs/prometheus_client/prometheus_client.go b/plugins/outputs/prometheus_client/prometheus_client.go
index 804ae1fad..790784a2b 100644
--- a/plugins/outputs/prometheus_client/prometheus_client.go
+++ b/plugins/outputs/prometheus_client/prometheus_client.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"regexp"
 	"strings"
+	"sync"
 
 	"github.com/influxdata/telegraf"
 	"github.com/influxdata/telegraf/plugins/outputs"
@@ -26,6 +27,10 @@ var (
 
 type PrometheusClient struct {
 	Listen string
+
+	metrics map[string]prometheus.Metric
+
+	sync.Mutex
 }
 
 var sampleConfig = `
@@ -34,6 +39,7 @@ var sampleConfig = `
 `
 
 func (p *PrometheusClient) Start() error {
+	prometheus.MustRegister(p)
 	defer func() {
 		if r := recover(); r != nil {
 			// recovering from panic here because there is no way to stop a
@@ -78,7 +84,27 @@ func (p *PrometheusClient) Description() string {
 	return "Configuration for the Prometheus client to spawn"
 }
 
+// Implements prometheus.Collector
+func (p *PrometheusClient) Describe(ch chan<- *prometheus.Desc) {
+	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(ch)
+}
+
+// Implements prometheus.Collector
+func (p *PrometheusClient) Collect(ch chan<- prometheus.Metric) {
+	p.Lock()
+	defer p.Unlock()
+
+	for _, m := range p.metrics {
+		ch <- m
+	}
+}
+
 func (p *PrometheusClient) Write(metrics []telegraf.Metric) error {
+	p.Lock()
+	defer p.Unlock()
+
+	p.metrics = make(map[string]prometheus.Metric)
+
 	if len(metrics) == 0 {
 		return nil
 	}
@@ -124,45 +150,23 @@ func (p *PrometheusClient) Write(metrics []telegraf.Metric) error {
 				continue
 			}
 
-			mVec := prometheus.NewUntypedVec(
-				prometheus.UntypedOpts{
-					Name: mname,
-					Help: "Telegraf collected metric",
-				},
-				labels,
-			)
-			collector, err := prometheus.RegisterOrGet(mVec)
-			if err != nil {
-				log.Printf("prometheus_client: Metric failed to register with prometheus, %s", err)
-				continue
-			}
-			mVec, ok := collector.(*prometheus.UntypedVec)
-			if !ok {
-				continue
-			}
-
+			desc := prometheus.NewDesc(mname, "Telegraf collected metric", nil, l)
+			var metric prometheus.Metric
+			var err error
 			switch val := val.(type) {
 			case int64:
-				m, err := mVec.GetMetricWith(l)
-				if err != nil {
-					log.Printf("ERROR Getting metric in Prometheus output, "+
-						"key: %s, labels: %v,\nerr: %s\n",
-						mname, l, err.Error())
-					continue
-				}
-				m.Set(float64(val))
+				metric, err = prometheus.NewConstMetric(desc, prometheus.UntypedValue, float64(val))
 			case float64:
-				m, err := mVec.GetMetricWith(l)
-				if err != nil {
-					log.Printf("ERROR Getting metric in Prometheus output, "+
-						"key: %s, labels: %v,\nerr: %s\n",
-						mname, l, err.Error())
-					continue
-				}
-				m.Set(val)
+				metric, err = prometheus.NewConstMetric(desc, prometheus.UntypedValue, val)
 			default:
 				continue
 			}
+			if err != nil {
+				log.Printf("ERROR creating prometheus metric, "+
+					"key: %s, labels: %v,\nerr: %s\n",
+					mname, l, err.Error())
+			}
+			p.metrics[desc.String()] = metric
 		}
 	}
 	return nil