From bbfc72b4e2ac06e109c151fa792bd31da8b0bd4b Mon Sep 17 00:00:00 2001
From: David Ashpole <dashpole@google.com>
Date: Fri, 5 Apr 2024 10:19:07 -0400
Subject: [PATCH] support unregistering discovery manager metrics (#13896)

Signed-off-by: David Ashpole <dashpole@google.com>
---
 discovery/manager.go      |  7 +++++++
 discovery/manager_test.go | 25 +++++++++++++++++++++++--
 discovery/metrics.go      |  9 +++++++++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/discovery/manager.go b/discovery/manager.go
index e3a263557..f14071af3 100644
--- a/discovery/manager.go
+++ b/discovery/manager.go
@@ -169,6 +169,13 @@ func (m *Manager) Providers() []*Provider {
 	return m.providers
 }
 
+// UnregisterMetrics unregisters manager metrics. It does not unregister
+// service discovery or refresh metrics, whose lifecycle is managed independent
+// of the discovery Manager.
+func (m *Manager) UnregisterMetrics() {
+	m.metrics.Unregister(m.registerer)
+}
+
 // Run starts the background processing.
 func (m *Manager) Run() error {
 	go m.sender()
diff --git a/discovery/manager_test.go b/discovery/manager_test.go
index 8f2345911..656d7c3c6 100644
--- a/discovery/manager_test.go
+++ b/discovery/manager_test.go
@@ -36,11 +36,11 @@ func TestMain(m *testing.M) {
 	testutil.TolerantVerifyLeak(m)
 }
 
-func NewTestMetrics(t *testing.T, reg prometheus.Registerer) (*RefreshMetricsManager, map[string]DiscovererMetrics) {
+func NewTestMetrics(t *testing.T, reg prometheus.Registerer) (RefreshMetricsManager, map[string]DiscovererMetrics) {
 	refreshMetrics := NewRefreshMetrics(reg)
 	sdMetrics, err := RegisterSDMetrics(reg, refreshMetrics)
 	require.NoError(t, err)
-	return &refreshMetrics, sdMetrics
+	return refreshMetrics, sdMetrics
 }
 
 // TestTargetUpdatesOrder checks that the target updates are received in the expected order.
@@ -1541,3 +1541,24 @@ func (t *testDiscoverer) update(tgs []*targetgroup.Group) {
 	<-t.ready
 	t.up <- tgs
 }
+
+func TestUnregisterMetrics(t *testing.T) {
+	reg := prometheus.NewRegistry()
+	// Check that all metrics can be unregistered, allowing a second manager to be created.
+	for i := 0; i < 2; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+
+		refreshMetrics, sdMetrics := NewTestMetrics(t, reg)
+
+		discoveryManager := NewManager(ctx, log.NewNopLogger(), reg, sdMetrics)
+		// discoveryManager will be nil if there was an error configuring metrics.
+		require.NotNil(t, discoveryManager)
+		// Unregister all metrics.
+		discoveryManager.UnregisterMetrics()
+		for _, sdMetric := range sdMetrics {
+			sdMetric.Unregister()
+		}
+		refreshMetrics.Unregister()
+		cancel()
+	}
+}
diff --git a/discovery/metrics.go b/discovery/metrics.go
index a77b86d27..e738331a1 100644
--- a/discovery/metrics.go
+++ b/discovery/metrics.go
@@ -99,3 +99,12 @@ func NewManagerMetrics(registerer prometheus.Registerer, sdManagerName string) (
 
 	return m, nil
 }
+
+// Unregister unregisters all metrics.
+func (m *Metrics) Unregister(registerer prometheus.Registerer) {
+	registerer.Unregister(m.FailedConfigs)
+	registerer.Unregister(m.DiscoveredTargets)
+	registerer.Unregister(m.ReceivedUpdates)
+	registerer.Unregister(m.DelayedUpdates)
+	registerer.Unregister(m.SentUpdates)
+}