From d277a2975a7a00b2196102670f7dca7ffd6722d8 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Thu, 1 Nov 2018 18:23:36 +0300
Subject: [PATCH] Use mtg metrics only for prometheus endpoint

---
 stats/init.go       |  2 +-
 stats/prometheus.go | 19 +++++++++++++++----
 stats/server.go     |  5 ++---
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/stats/init.go b/stats/init.go
index ba303d7..ba3640f 100644
--- a/stats/init.go
+++ b/stats/init.go
@@ -22,7 +22,7 @@ func Init(conf *config.Config) error {
 	go prometheus.run()
 
 	go NewStats(conf).start()
-	go startServer(conf)
+	go startServer(conf, prometheus.getHTTPHandler())
 
 	return nil
 }
diff --git a/stats/prometheus.go b/stats/prometheus.go
index 00490a5..00257fa 100644
--- a/stats/prometheus.go
+++ b/stats/prometheus.go
@@ -1,10 +1,12 @@
 package stats
 
 import (
+	"net/http"
 	"time"
 
 	"github.com/juju/errors"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 
 	"github.com/9seconds/mtg/config"
 )
@@ -12,6 +14,8 @@ import (
 const prometheusPollTime = time.Second
 
 type prometheusExporter struct {
+	registry prometheus.Gatherer
+
 	connections *prometheus.GaugeVec
 	traffic     *prometheus.GaugeVec
 	speed       *prometheus.GaugeVec
@@ -36,7 +40,13 @@ func (p *prometheusExporter) run() {
 	}
 }
 
+func (p *prometheusExporter) getHTTPHandler() http.Handler {
+	return promhttp.HandlerFor(p.registry, promhttp.HandlerOpts{})
+}
+
 func newPrometheus(conf *config.Config) (*prometheusExporter, error) {
+	registry := prometheus.NewRegistry()
+
 	connections := prometheus.NewGaugeVec(prometheus.GaugeOpts{
 		Namespace: conf.Prometheus.Prefix,
 		Name:      "connections",
@@ -58,20 +68,21 @@ func newPrometheus(conf *config.Config) (*prometheusExporter, error) {
 		Help:      "How many crashes happened.",
 	})
 
-	if err := prometheus.Register(connections); err != nil {
+	if err := registry.Register(connections); err != nil {
 		return nil, errors.Annotate(err, "Cannot register connections collector")
 	}
-	if err := prometheus.Register(traffic); err != nil {
+	if err := registry.Register(traffic); err != nil {
 		return nil, errors.Annotate(err, "cannot register traffic collector")
 	}
-	if err := prometheus.Register(speed); err != nil {
+	if err := registry.Register(speed); err != nil {
 		return nil, errors.Annotate(err, "cannot register speed collector")
 	}
-	if err := prometheus.Register(crashes); err != nil {
+	if err := registry.Register(crashes); err != nil {
 		return nil, errors.Annotate(err, "cannot register crashes collector")
 	}
 
 	return &prometheusExporter{
+		registry:    registry,
 		connections: connections,
 		traffic:     traffic,
 		speed:       speed,
diff --git a/stats/server.go b/stats/server.go
index 9d96b02..90d688c 100644
--- a/stats/server.go
+++ b/stats/server.go
@@ -4,13 +4,12 @@ import (
 	"encoding/json"
 	"net/http"
 
-	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 )
 
-func startServer(conf *config.Config) {
+func startServer(conf *config.Config, prometheusHandler http.Handler) {
 	log := zap.S().Named("stats")
 
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
@@ -33,7 +32,7 @@ func startServer(conf *config.Config) {
 			log.Errorw("Cannot encode json", "error", err)
 		}
 	})
-	http.Handle("/prometheus/", promhttp.Handler())
+	http.Handle("/prometheus/", prometheusHandler)
 
 	if err := http.ListenAndServe(conf.StatAddr(), nil); err != nil {
 		log.Fatalw("Stats server has been stopped", "error", err)