Merge pull request #3 from warmans/feature/add-labels

Add Label Feature

config.yml

@@ -2,7 +2,5 @@ server:
   bind: ":8080"
 timeout: 1000 #ms
 targets:
-  - "http://localhost:9100/metrics"
-  - "http://localhost:9100/metrics"
-  - "http://localhost:9100/metrics"
-  - "http://localhost:9100/metrics"
+  - "http://localhost/prom.txt"
+  - "http://localhost/prom2.txt"

fixture/histogram.golden.txt

@@ -0,0 +1,10 @@
+# HELP http_request_duration_seconds A histogram of the request duration.
+# TYPE http_request_duration_seconds histogram
+http_request_duration_seconds_bucket{ae_source="foo",le="0.05"} 24054
+http_request_duration_seconds_bucket{ae_source="foo",le="0.1"} 33444
+http_request_duration_seconds_bucket{ae_source="foo",le="0.2"} 100392
+http_request_duration_seconds_bucket{ae_source="foo",le="0.5"} 129389
+http_request_duration_seconds_bucket{ae_source="foo",le="1"} 133988
+http_request_duration_seconds_bucket{ae_source="foo",le="+Inf"} 144320
+http_request_duration_seconds_sum{ae_source="foo"} 53423
+http_request_duration_seconds_count{ae_source="foo"} 144320

fixture/histogram.txt

@@ -0,0 +1,9 @@
+# HELP http_requests_total The total number of HTTP requests.
+# TYPE http_requests_total counter
+http_requests_total{method="post",code="200"} 1027 1395066363000
+http_requests_total{method="post",code="400"}    3 1395066363000
+
+# HELP http_requests_total The total number of HTTP requests.
+# TYPE http_requests_total counter
+http_requests_total{method="post",code="200"} 1027 1395066363000
+http_requests_total{method="post",code="400"}    3 1395066363000

glide.yaml

@@ -1,3 +1,4 @@
 package: github.com/warmans/prometheus-aggregate-exporter
 import:
 - package: gopkg.in/yaml.v2
+- package: github.com/prometheus/common

main.go

@@ -13,16 +13,21 @@ import (
 	"net/http"
 	"strconv"
 	"time"
+
+	"github.com/prometheus/client_model/go"
+	"github.com/prometheus/common/expfmt"
 )
 
 var (
-	Version    = "unknown"
+	Version = "unknown"
 )
 
 var (
 	configPathFlag = flag.String("config", "config.yml", "Path to config YAML file.")
-	verboseFlag = flag.Bool("verbose", false, "Log more information")
-	versionFlag = flag.Bool("version", false, "Show version and exit")
+	verboseFlag    = flag.Bool("verbose", false, "Log more information")
+	versionFlag    = flag.Bool("version", false, "Show version and exit")
+	appendLabel    = flag.Bool("label", false, "Add a label to metrics to show their origin target")
+	labelName      = flag.String("label.name", "ae_source", "Label name to use if a target name label is appended to metrics")
 )
 
 type Config struct {
@@ -85,7 +90,7 @@ func main() {
 type Result struct {
 	URL          string
 	SecondsTaken float64
-	Payload      io.ReadCloser
+	MetricFamily map[string]*io_prometheus_client.MetricFamily
 	Error        error
 }
 
@@ -102,10 +107,14 @@ func (f *Aggregator) Aggregate(targets []string, output io.Writer) {
 	}
 
 	func(numTargets int, resultChan chan *Result) {
+
 		numResuts := 0
+
+		allFamilies := make(map[string]*io_prometheus_client.MetricFamily)
+
 		for {
 			if numTargets == numResuts {
-				return
+				break
 			}
 			select {
 			case result := <-resultChan:
@@ -116,34 +125,59 @@ func (f *Aggregator) Aggregate(targets []string, output io.Writer) {
 					continue
 				}
 
-				_, err := io.Copy(output, result.Payload)
-				if err != nil {
-					log.Printf("Copy error: %s", err.Error())
+				for mfName, mf := range result.MetricFamily {
+					if *appendLabel {
+						for _, m := range mf.Metric {
+							m.Label = append(m.Label, &io_prometheus_client.LabelPair{Name: labelName, Value: &result.URL})
+						}
+					}
+					if existingMf, ok := allFamilies[mfName]; ok {
+						for _, m := range mf.Metric {
+							existingMf.Metric = append(existingMf.Metric, m)
+						}
+					} else {
+						allFamilies[*mf.Name] = mf
+					}
 				}
-
-				err = result.Payload.Close()
-				if err != nil {
-					log.Printf("Result body close error: %s", err.Error())
-				}
-
 				if *verboseFlag {
 					log.Printf("OK: %s was refreshed in %.3f seconds", result.URL, result.SecondsTaken)
 				}
 			}
 		}
+
+		encoder := expfmt.NewEncoder(output, expfmt.FmtText)
+		for _, f := range allFamilies {
+			encoder.Encode(f)
+		}
+
 	}(len(targets), resultChan)
 }
 
 func (f *Aggregator) fetch(target string, resultChan chan *Result) {
 
 	startTime := time.Now()
 	res, err := f.HTTP.Get(target)
+
 	result := &Result{URL: target, SecondsTaken: time.Since(startTime).Seconds(), Error: nil}
 	if res != nil {
-		result.Payload = res.Body
+		result.MetricFamily, err = getMetricFamilies(res.Body)
+		if err != nil {
+			result.Error = fmt.Errorf("failed to add labels to target %s metrics: %s", target, err.Error())
+			resultChan <- result
+			return
+		}
 	}
 	if err != nil {
-		result.Error = fmt.Errorf("Failed to fetch URL %s due to error: %s", target, err.Error())
+		result.Error = fmt.Errorf("failed to fetch URL %s due to error: %s", target, err.Error())
 	}
 	resultChan <- result
 }
+
+func getMetricFamilies(sourceData io.Reader) (map[string]*io_prometheus_client.MetricFamily, error) {
+	parser := expfmt.TextParser{}
+	metricFamiles, err := parser.TextToMetricFamilies(sourceData)
+	if err != nil {
+		return nil, err
+	}
+	return metricFamiles, nil
+}

main_test.go

@@ -0,0 +1,36 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"testing"
+)
+
+var updateGoldenFile = flag.Bool("update.golden", false, "update golden files")
+
+func TestMain(m *testing.M) {
+
+	flag.Parse()
+	os.Exit(m.Run())
+}
+
+func mustOpenFile(name string, flag int) *os.File {
+	file, err := os.OpenFile(fmt.Sprintf("fixture/%s", name), flag, 0666)
+	if err != nil {
+		panic(err)
+	}
+	return file
+}
+
+func mustReadAll(r io.Reader) string {
+	b, err := ioutil.ReadAll(r)
+	if err != nil {
+		panic("failed to readall reader: " + err.Error())
+	}
+	return string(b[:])
+}
+
+//todo: re-implement tests