From a1db7c3266ac65863e8dbf144f0f13f3b8b2ac90 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Julian=20T=C3=B6lle?= <julian.toelle@hetzner-cloud.de>
Date: Fri, 14 Jul 2023 15:30:59 +0200
Subject: [PATCH] fix(instrumentation): multiple instrumented clients cause
 panic (#289)

(cherry picked from commit c0ef9b6e6e3f36d8c2282c2b7aa9d8687141f291)
---
 hcloud/internal/instrumentation/metrics.go    | 66 +++++++++++++------
 .../internal/instrumentation/metrics_test.go  | 16 ++++-
 2 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/hcloud/internal/instrumentation/metrics.go b/hcloud/internal/instrumentation/metrics.go
index 69a7165b..aa57c710 100644
--- a/hcloud/internal/instrumentation/metrics.go
+++ b/hcloud/internal/instrumentation/metrics.go
@@ -23,29 +23,36 @@ func New(subsystemIdentifier string, instrumentationRegistry *prometheus.Registr
 
 // InstrumentedRoundTripper returns an instrumented round tripper.
 func (i *Instrumenter) InstrumentedRoundTripper() http.RoundTripper {
-	inFlightRequestsGauge := prometheus.NewGauge(prometheus.GaugeOpts{
-		Name: fmt.Sprintf("hcloud_%s_in_flight_requests", i.subsystemIdentifier),
-		Help: fmt.Sprintf("A gauge of in-flight requests to the hcloud %s.", i.subsystemIdentifier),
-	})
-
-	requestsPerEndpointCounter := prometheus.NewCounterVec(
-		prometheus.CounterOpts{
-			Name: fmt.Sprintf("hcloud_%s_requests_total", i.subsystemIdentifier),
-			Help: fmt.Sprintf("A counter for requests to the hcloud %s per endpoint.", i.subsystemIdentifier),
-		},
-		[]string{"code", "method", "api_endpoint"},
+	inFlightRequestsGauge := registerOrReuse(
+		i.instrumentationRegistry,
+		prometheus.NewGauge(prometheus.GaugeOpts{
+			Name: fmt.Sprintf("hcloud_%s_in_flight_requests", i.subsystemIdentifier),
+			Help: fmt.Sprintf("A gauge of in-flight requests to the hcloud %s.", i.subsystemIdentifier),
+		}),
 	)
 
-	requestLatencyHistogram := prometheus.NewHistogramVec(
-		prometheus.HistogramOpts{
-			Name:    fmt.Sprintf("hcloud_%s_request_duration_seconds", i.subsystemIdentifier),
-			Help:    fmt.Sprintf("A histogram of request latencies to the hcloud %s .", i.subsystemIdentifier),
-			Buckets: prometheus.DefBuckets,
-		},
-		[]string{"method"},
+	requestsPerEndpointCounter := registerOrReuse(
+		i.instrumentationRegistry,
+		prometheus.NewCounterVec(
+			prometheus.CounterOpts{
+				Name: fmt.Sprintf("hcloud_%s_requests_total", i.subsystemIdentifier),
+				Help: fmt.Sprintf("A counter for requests to the hcloud %s per endpoint.", i.subsystemIdentifier),
+			},
+			[]string{"code", "method", "api_endpoint"},
+		),
 	)
 
-	i.instrumentationRegistry.MustRegister(requestsPerEndpointCounter, requestLatencyHistogram, inFlightRequestsGauge)
+	requestLatencyHistogram := registerOrReuse(
+		i.instrumentationRegistry,
+		prometheus.NewHistogramVec(
+			prometheus.HistogramOpts{
+				Name:    fmt.Sprintf("hcloud_%s_request_duration_seconds", i.subsystemIdentifier),
+				Help:    fmt.Sprintf("A histogram of request latencies to the hcloud %s .", i.subsystemIdentifier),
+				Buckets: prometheus.DefBuckets,
+			},
+			[]string{"method"},
+		),
+	)
 
 	return promhttp.InstrumentRoundTripperInFlight(inFlightRequestsGauge,
 		promhttp.InstrumentRoundTripperDuration(requestLatencyHistogram,
@@ -74,6 +81,27 @@ func (i *Instrumenter) instrumentRoundTripperEndpoint(counter *prometheus.Counte
 	}
 }
 
+// registerOrReuse will try to register the passed Collector, but in case a conflicting collector was already registered,
+// it will instead return that collector. Make sure to always use the collector return by this method.
+// Similar to [Registry.MustRegister] it will panic if any other error occurs.
+func registerOrReuse[C prometheus.Collector](registry *prometheus.Registry, collector C) C {
+	err := registry.Register(collector)
+	if err != nil {
+		// If we get a AlreadyRegisteredError we can return the existing collector
+		if are, ok := err.(prometheus.AlreadyRegisteredError); ok {
+			if existingCollector, ok := are.ExistingCollector.(C); ok {
+				collector = existingCollector
+			} else {
+				panic("received incompatible existing collector")
+			}
+		} else {
+			panic(err)
+		}
+	}
+
+	return collector
+}
+
 func preparePathForLabel(path string) string {
 	path = strings.ToLower(path)
 
diff --git a/hcloud/internal/instrumentation/metrics_test.go b/hcloud/internal/instrumentation/metrics_test.go
index b6144a22..6b493075 100644
--- a/hcloud/internal/instrumentation/metrics_test.go
+++ b/hcloud/internal/instrumentation/metrics_test.go
@@ -1,6 +1,10 @@
 package instrumentation
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
 
 func Test_preparePath(t *testing.T) {
 	tests := []struct {
@@ -27,3 +31,13 @@ func Test_preparePath(t *testing.T) {
 		})
 	}
 }
+
+func TestMultipleInstrumentedClients(t *testing.T) {
+	reg := prometheus.NewRegistry()
+
+	t.Run("should not panic", func(t *testing.T) {
+		// Following code should run without panicking
+		New("test", reg).InstrumentedRoundTripper()
+		New("test", reg).InstrumentedRoundTripper()
+	})
+}