From 5f4065a3a9f1cff0fed4f2b58e5cc6c1f40b6c59 Mon Sep 17 00:00:00 2001
From: Rob Holland <rob.holland@gmail.com>
Date: Thu, 22 Feb 2024 12:28:44 +0000
Subject: [PATCH] Rewire metrics. (#78)

* Rewire metrics.

Prior to this metrics were not being reported correctly.
---
 cmd/cmdoptions/metrics.go | 108 ++++++++++++++++++++++++++++----------
 1 file changed, 80 insertions(+), 28 deletions(-)

diff --git a/cmd/cmdoptions/metrics.go b/cmd/cmdoptions/metrics.go
index d16703f..254be7a 100644
--- a/cmd/cmdoptions/metrics.go
+++ b/cmd/cmdoptions/metrics.go
@@ -2,9 +2,10 @@ package cmdoptions
 
 import (
 	"context"
-	"errors"
+	"fmt"
 	"net"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -16,49 +17,98 @@ import (
 )
 
 type metricsHandler struct {
-	registry *prometheus.Registry
-	tags     map[string]string
+	metrics *Metrics
+	labels  []string
+	values  []string
 }
 
 var _ client.MetricsHandler = (*metricsHandler)(nil)
 
 func (h *metricsHandler) WithTags(tags map[string]string) client.MetricsHandler {
 	// Make enough space for the handlers tags which are populated first
-	mergedTags := make(map[string]string, len(h.tags))
-	for t, v := range h.tags {
-		mergedTags[t] = v
+	mergedTags := make(map[string]string, len(h.labels))
+	for i, l := range h.labels {
+		mergedTags[l] = h.values[i]
 	}
-	for t, v := range tags {
-		mergedTags[t] = v
+	for l, v := range tags {
+		mergedTags[l] = v
 	}
-	return &metricsHandler{registry: h.registry, tags: mergedTags}
-}
 
-func (h *metricsHandler) mustRegisterIgnoreDuplicate(c prometheus.Collector) {
-	err := h.registry.Register(c)
-	var alreadyRegisteredError prometheus.AlreadyRegisteredError
-	if err != nil && !errors.As(err, &alreadyRegisteredError) {
-		panic(err)
+	var labels, values []string
+	for l, v := range mergedTags {
+		labels = append(labels, l)
+		values = append(values, v)
+	}
+
+	return &metricsHandler{
+		metrics: h.metrics,
+		labels:  labels,
+		values:  values,
 	}
 }
 
 func (h *metricsHandler) Counter(name string) client.MetricsCounter {
-	ctr := prometheus.NewCounter(prometheus.CounterOpts{Name: name, ConstLabels: prometheus.Labels(h.tags)})
-	h.mustRegisterIgnoreDuplicate(ctr)
-	return metricsCounter{ctr}
+	h.metrics.mutex.Lock()
+	defer h.metrics.mutex.Unlock()
+
+	var ctr *prometheus.CounterVec
+	if c, ok := h.metrics.cache[name]; ok {
+		ctr, ok = c.(*prometheus.CounterVec)
+		if !ok {
+			panic(fmt.Errorf("duplicate metric with different type: %s", name))
+		}
+	} else {
+		ctr = prometheus.NewCounterVec(
+			prometheus.CounterOpts{Name: name},
+			h.labels,
+		)
+		h.metrics.registry.MustRegister(ctr)
+		h.metrics.cache[name] = ctr
+	}
+
+	return metricsCounter{ctr.WithLabelValues(h.values...)}
 }
 
 func (h *metricsHandler) Gauge(name string) client.MetricsGauge {
-	gauge := prometheus.NewGauge(prometheus.GaugeOpts{Name: name, ConstLabels: prometheus.Labels(h.tags)})
-	h.mustRegisterIgnoreDuplicate(gauge)
-	return metricsGauge{gauge}
+	h.metrics.mutex.Lock()
+	defer h.metrics.mutex.Unlock()
+
+	var gauge *prometheus.GaugeVec
+	if c, ok := h.metrics.cache[name]; ok {
+		gauge, ok = c.(*prometheus.GaugeVec)
+		if !ok {
+			panic(fmt.Errorf("duplicate metric with different type: %s", name))
+		}
+	} else {
+		gauge = prometheus.NewGaugeVec(
+			prometheus.GaugeOpts{Name: name},
+			h.labels,
+		)
+		h.metrics.registry.MustRegister(gauge)
+		h.metrics.cache[name] = gauge
+	}
+
+	return metricsGauge{gauge.WithLabelValues(h.values...)}
 }
 
 func (h *metricsHandler) Timer(name string) client.MetricsTimer {
-	// TODO: buckets
-	timer := prometheus.NewHistogram(prometheus.HistogramOpts{Name: name, ConstLabels: prometheus.Labels(h.tags)})
-	h.mustRegisterIgnoreDuplicate(timer)
-	return metricsTimer{timer}
+	h.metrics.mutex.Lock()
+	defer h.metrics.mutex.Unlock()
+
+	var timer *prometheus.HistogramVec
+	if c, ok := h.metrics.cache[name]; ok {
+		timer, ok = c.(*prometheus.HistogramVec)
+		if !ok {
+			panic(fmt.Errorf("duplicate metric with different type: %s", name))
+		}
+	} else {
+		// TODO: buckets
+		timer = prometheus.NewHistogramVec(prometheus.HistogramOpts{Name: name}, h.labels)
+		h.metrics.registry.MustRegister(timer)
+		h.metrics.cache[name] = timer
+	}
+
+	return metricsTimer{timer.WithLabelValues(h.values...)}
 }
 
 type metricsCounter struct {
@@ -80,7 +130,7 @@ func (m metricsGauge) Update(x float64) {
 }
 
 type metricsTimer struct {
-	prom prometheus.Histogram
+	prom prometheus.Observer
 }
 
 // Record records a duration.
@@ -102,6 +152,8 @@ type MetricsOptions struct {
 type Metrics struct {
 	server   *http.Server
 	registry *prometheus.Registry
+	cache    map[string]interface{}
+	mutex    sync.Mutex
 }
 
 // MustCreateMetrics sets up Prometheus based metrics and starts an HTTP server
@@ -116,14 +168,14 @@ func (m *MetricsOptions) MustCreateMetrics(logger *zap.SugaredLogger) *Metrics {
 	return &Metrics{
 		server:   server,
 		registry: registry,
+		cache:    make(map[string]interface{}),
 	}
 }
 
 // Handler returns a new Temporal-client-compatible metrics handler.
 func (m *Metrics) NewHandler() client.MetricsHandler {
 	return &metricsHandler{
-		registry: m.registry,
-		tags:     make(map[string]string),
+		metrics: m,
 	}
 }