From 0e0a85f2a787d4e1f5b07aec0c68fe9e0c6a424e Mon Sep 17 00:00:00 2001
From: Ferran Rodenas <rodenasf@vmware.com>
Date: Tue, 31 Oct 2017 14:19:51 +0100
Subject: [PATCH] Add metrics.UnregisterMetricAndUntrackRateLimiterUsage
 function

For testing purposes, we want to unregister a previously registered
rate limiter prometheus metric and stop the goroutine that updates
this metric.

Signed-off-by: Ferran Rodenas <rodenasf@vmware.com>
---
 pkg/util/metrics/util.go      | 44 ++++++++++++++++++++++++++++-------
 pkg/util/metrics/util_test.go | 29 +++++++++++++++++++++++
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/pkg/util/metrics/util.go b/pkg/util/metrics/util.go
index c1d92475622..3980ae818ad 100644
--- a/pkg/util/metrics/util.go
+++ b/pkg/util/metrics/util.go
@@ -34,40 +34,66 @@ const (
 
 var (
 	metricsLock        sync.Mutex
-	rateLimiterMetrics = make(map[string]prometheus.Gauge)
+	rateLimiterMetrics = make(map[string]rateLimiterMetric)
 )
 
+type rateLimiterMetric struct {
+	metric prometheus.Gauge
+	stopCh chan struct{}
+}
+
 func registerRateLimiterMetric(ownerName string) error {
 	metricsLock.Lock()
 	defer metricsLock.Unlock()
 
 	if _, ok := rateLimiterMetrics[ownerName]; ok {
-		glog.Errorf("Metric for %v already registered", ownerName)
-		return fmt.Errorf("Metric for %v already registered", ownerName)
+		return fmt.Errorf("Rate Limiter Metric for %v already registered", ownerName)
 	}
 	metric := prometheus.NewGauge(prometheus.GaugeOpts{
 		Name:      "rate_limiter_use",
 		Subsystem: ownerName,
 		Help:      fmt.Sprintf("A metric measuring the saturation of the rate limiter for %v", ownerName),
 	})
-	rateLimiterMetrics[ownerName] = metric
 	if err := prometheus.Register(metric); err != nil {
 		return fmt.Errorf("error registering rate limiter usage metric: %v", err)
 	}
+	stopCh := make(chan struct{})
+	rateLimiterMetrics[ownerName] = rateLimiterMetric{
+		metric: metric,
+		stopCh: stopCh,
+	}
 	return nil
 }
 
 // RegisterMetricAndTrackRateLimiterUsage registers a metric ownerName_rate_limiter_use in prometheus to track
 // how much used rateLimiter is and starts a goroutine that updates this metric every updatePeriod
 func RegisterMetricAndTrackRateLimiterUsage(ownerName string, rateLimiter flowcontrol.RateLimiter) error {
-	err := registerRateLimiterMetric(ownerName)
-	if err != nil {
+	if err := registerRateLimiterMetric(ownerName); err != nil {
 		return err
 	}
-	go wait.Forever(func() {
+	go wait.Until(func() {
 		metricsLock.Lock()
 		defer metricsLock.Unlock()
-		rateLimiterMetrics[ownerName].Set(rateLimiter.Saturation())
-	}, updatePeriod)
+		rateLimiterMetrics[ownerName].metric.Set(rateLimiter.Saturation())
+	}, updatePeriod, rateLimiterMetrics[ownerName].stopCh)
 	return nil
 }
+
+// UnregisterMetricAndUntrackRateLimiterUsage unregisters a metric ownerName_rate_limiter_use from prometheus and
+// stops the goroutine that updates this metric
+func UnregisterMetricAndUntrackRateLimiterUsage(ownerName string) bool {
+	metricsLock.Lock()
+	defer metricsLock.Unlock()
+
+	rlm, ok := rateLimiterMetrics[ownerName]
+	if !ok {
+		glog.Warningf("Rate Limiter Metric for %v not registered", ownerName)
+		return false
+	}
+
+	close(rlm.stopCh)
+	prometheus.Unregister(rlm.metric)
+	delete(rateLimiterMetrics, ownerName)
+
+	return true
+}
diff --git a/pkg/util/metrics/util_test.go b/pkg/util/metrics/util_test.go
index 38e14b577b4..ee83aa1b7e5 100644
--- a/pkg/util/metrics/util_test.go
+++ b/pkg/util/metrics/util_test.go
@@ -34,6 +34,11 @@ func TestRegisterMetricAndTrackRateLimiterUsage(t *testing.T) {
 			rateLimiter: flowcontrol.NewTokenBucketRateLimiter(1, 1),
 			err:         "",
 		},
+		{
+			ownerName:   "owner_name",
+			rateLimiter: flowcontrol.NewTokenBucketRateLimiter(1, 1),
+			err:         "already registered",
+		},
 		{
 			ownerName:   "invalid-owner-name",
 			rateLimiter: flowcontrol.NewTokenBucketRateLimiter(1, 1),
@@ -52,3 +57,27 @@ func TestRegisterMetricAndTrackRateLimiterUsage(t *testing.T) {
 		}
 	}
 }
+
+func TestUnregisterMetricAndUntrackRateLimiterUsage(t *testing.T) {
+	RegisterMetricAndTrackRateLimiterUsage("owner_name", flowcontrol.NewTokenBucketRateLimiter(1, 1))
+	testCases := []struct {
+		ownerName string
+		ok        bool
+	}{
+		{
+			ownerName: "owner_name",
+			ok:        true,
+		},
+		{
+			ownerName: "owner_name",
+			ok:        false,
+		},
+	}
+
+	for i, tc := range testCases {
+		ok := UnregisterMetricAndUntrackRateLimiterUsage(tc.ownerName)
+		if tc.ok != ok {
+			t.Errorf("Case[%d] Expected %v, got %v", i, tc.ok, ok)
+		}
+	}
+}