From f25df3967509f3b5959cb31363b3e2a7da232964 Mon Sep 17 00:00:00 2001
From: drpytho <filip@voiapp.io>
Date: Wed, 26 Aug 2020 15:26:26 +0200
Subject: [PATCH] Use internal connection of gatherers

---
 options.go | 11 +++++++----
 svc.go     | 16 ++++++++++++++++
 worker.go  |  7 +++++++
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/options.go b/options.go
index e0fac39..c492089 100644
--- a/options.go
+++ b/options.go
@@ -66,16 +66,16 @@ func WithHTTPServer(port string) Option {
 // WithMetrics is an option that exports metrics via prometheus.
 func WithMetrics() Option {
 	return func(s *SVC) error {
-		m := prometheus.NewGaugeFunc(
+		m := prometheus.NewGauge(
 			prometheus.GaugeOpts{
 				Name:        "svc_up",
 				Help:        "Is the service in this pod up.",
 				ConstLabels: prometheus.Labels{"version": s.Version, "name": s.Name},
 			},
-			func() float64 { return 1 },
 		)
+		m.Set(1)
 
-		if err := prometheus.Register(m); err != nil {
+		if err := s.internalRegister.Register(m); err != nil {
 			s.logger.Error("svc_up could not register", zap.Error(err))
 		}
 
@@ -87,7 +87,10 @@ func WithMetrics() Option {
 //Prometheus scraper.
 func WithMetricsHandler() Option {
 	return func(s *SVC) error {
-		s.Router.Handle("/metrics", promhttp.Handler())
+		s.Router.Handle("/metrics",
+			promhttp.InstrumentMetricHandler(
+				s.internalRegister, /* Register */
+				promhttp.HandlerFor(s.gatherers, promhttp.HandlerOpts{})))
 
 		return nil
 	}
diff --git a/svc.go b/svc.go
index 8a0716c..6ac0942 100644
--- a/svc.go
+++ b/svc.go
@@ -10,6 +10,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
 )
 
@@ -39,6 +40,9 @@ type SVC struct {
 	workers            map[string]Worker
 	workersAdded       []string
 	workersInitialized []string
+
+	gatherers        prometheus.Gatherers
+	internalRegister *prometheus.Registry
 }
 
 // New instantiates a new service by parsing configuration and initializing a
@@ -63,6 +67,9 @@ func New(name, version string, opts ...Option) (*SVC, error) {
 		return nil, err
 	}
 
+	s.internalRegister = prometheus.NewRegistry()
+	s.gatherers = []prometheus.Gatherer{s.internalRegister, prometheus.DefaultGatherer}
+
 	// Apply options
 	for _, o := range opts {
 		if err := o(s); err != nil {
@@ -82,11 +89,20 @@ func (s *SVC) AddWorker(name string, w Worker) {
 	if _, ok := w.(Healther); !ok {
 		s.logger.Warn("Worker does not implement Healther interface", zap.String("worker", name))
 	}
+	if g, ok := w.(Gatherer); ok {
+		s.AddGatherer(g.Gatherer())
+	} else {
+		s.logger.Warn("Worker does not implement Gatherer interface", zap.String("worker", name))
+	}
 	// Track workers as ordered set to initialize them in order.
 	s.workersAdded = append(s.workersAdded, name)
 	s.workers[name] = w
 }
 
+func (s *SVC) AddGatherer(gatherer prometheus.Gatherer) {
+	s.gatherers = append(s.gatherers, gatherer)
+}
+
 // Run runs the service until either receiving an interrupt or a worker
 // terminates.
 func (s *SVC) Run() {
diff --git a/worker.go b/worker.go
index 88caf58..4116c18 100644
--- a/worker.go
+++ b/worker.go
@@ -1,6 +1,7 @@
 package svc
 
 import (
+	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
 )
 
@@ -15,3 +16,9 @@ type Worker interface {
 type Healther interface {
 	Healthy() error
 }
+
+// Gatherer is a place for workers to return a prometheus.Gatherer
+// for SVC to serve on the metrics endpoint
+type Gatherer interface {
+	Gatherer() prometheus.Gatherer
+}