From a494df34627707a5afa0d06c98be43eb7edd55d6 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 26 Feb 2020 20:37:47 +0200
Subject: [PATCH] Properly unregister Summary metric in Set.UnregisterMetric

---
 set.go      | 54 +++++++++++++++++++++++++++++++++++++++++++----------
 set_test.go | 34 +++++++++++++++++++--------------
 summary.go  | 19 +++++++++++++++++++
 3 files changed, 83 insertions(+), 24 deletions(-)

diff --git a/set.go b/set.go
index 18021a4..8195d8b 100644
--- a/set.go
+++ b/set.go
@@ -439,23 +439,57 @@ func (s *Set) registerMetric(name string, m metric) {
 	}
 }
 
-// UnregisterMetric will remove a registered metric
+// UnregisterMetric removes metric with the given name from s.
+//
+// True is returned if the metric has been removed.
+// False is returned if the given metric is missing in s.
 func (s *Set) UnregisterMetric(name string) bool {
 	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	nm, ok := s.m[name]
-	if ok {
-		delete(s.m, name)
-		for i, a := range s.a {
-			if a == nm {
-				s.a = append(s.a[:i], s.a[i+1:]...)
-			}
+	if !ok {
+		return false
+	}
+	m := nm.metric
+
+	delete(s.m, name)
+
+	// remove metric from s.a
+	found := false
+	for i, nm := range s.a {
+		if nm.name == name {
+			s.a = append(s.a[:i], s.a[i+1:]...)
+			found = true
+			break
 		}
 	}
-	s.mu.Unlock()
-	return ok
+	if !found {
+		panic(fmt.Errorf("BUG: cannot find metric %q in the list of registered metrics", name))
+	}
+	sm, ok := m.(*Summary)
+	if !ok {
+		// There is no need in cleaning up s.summaries.
+		return true
+	}
+
+	// Remove sm from s.summaries
+	found = false
+	for i, xsm := range s.summaries {
+		if xsm == sm {
+			s.summaries = append(s.summaries[:i], s.summaries[i+1:]...)
+			found = true
+			break
+		}
+	}
+	if !found {
+		panic(fmt.Errorf("BUG: cannot find summary %q in the list of registered summaries", name))
+	}
+	unregisterSummary(sm)
+	return true
 }
 
-// ListMetricNames will return a list of all registered metrics
+// ListMetricNames returns a list of all the metrics in s.
 func (s *Set) ListMetricNames() []string {
 	var list []string
 	for name := range s.m {
diff --git a/set_test.go b/set_test.go
index f1952d8..4c96353 100644
--- a/set_test.go
+++ b/set_test.go
@@ -35,7 +35,7 @@ func TestNewSet(t *testing.T) {
 	}
 }
 
-func TestListMetricNames(t *testing.T) {
+func TestSetListMetricNames(t *testing.T) {
 	s := NewSet()
 	expect := []string{"cnt1", "cnt2", "cnt3"}
 	// Initialize a few counters
@@ -62,33 +62,39 @@ func TestListMetricNames(t *testing.T) {
 	}
 }
 
-func TestUnregisterMetric(t *testing.T) {
+func TestSetUnregisterMetric(t *testing.T) {
 	s := NewSet()
-	// Initialize a few counters
+	// Initialize a few metrics
 	for i := 0; i < 5; i++ {
 		c := s.NewCounter(fmt.Sprintf("counter_%d", i))
 		c.Inc()
+		sm := s.NewSummary(fmt.Sprintf("summary_%d", i))
+		sm.Update(float64(i))
 	}
-	// Unregister existing counters
-	ok := s.UnregisterMetric("counter_1")
-	if !ok {
-		t.Fatalf("Metric counter_1 should return true for deregistering")
+	// Unregister existing metrics
+	if !s.UnregisterMetric("counter_1") {
+		t.Fatalf("UnregisterMetric(counter_1) must return true")
+	}
+	if !s.UnregisterMetric("summary_1") {
+		t.Fatalf("UnregisterMetric(summary_1) must return true")
 	}
 
 	// Unregister twice must return false
-	ok = s.UnregisterMetric("counter_1")
-	if ok {
-		t.Fatalf("Metric counter_1 should not return false on unregister twice")
+	if s.UnregisterMetric("counter_1") {
+		t.Fatalf("UnregisterMetric(counter_1) must return false on unregistered metric")
+	}
+	if s.UnregisterMetric("summary_1") {
+		t.Fatalf("UnregisterMetric(summary_1) must return false on unregistered metric")
 	}
 
-	// Validate counters are removed
-	ok = false
+	// Validate metrics are removed
+	ok := false
 	for _, n := range s.ListMetricNames() {
-		if n == "counter_1" {
+		if n == "counter_1" || n == "summary_1" {
 			ok = true
 		}
 	}
 	if ok {
-		t.Fatalf("Metric counter_1 and counter_3 must not be listed anymore after unregister")
+		t.Fatalf("Metric counter_1 and summary_1 must not be listed anymore after unregister")
 	}
 }
diff --git a/summary.go b/summary.go
index 2c14bea..6cd278c 100644
--- a/summary.go
+++ b/summary.go
@@ -213,6 +213,25 @@ func registerSummary(sm *Summary) {
 	summariesLock.Unlock()
 }
 
+func unregisterSummary(sm *Summary) {
+	window := sm.window
+	summariesLock.Lock()
+	sms := summaries[window]
+	found := false
+	for i, xsm := range sms {
+		if xsm == sm {
+			sms = append(sms[:i], sms[i+1:]...)
+			found = true
+			break
+		}
+	}
+	if !found {
+		panic(fmt.Errorf("BUG: cannot find registered summary %p", sm))
+	}
+	summaries[window] = sms
+	summariesLock.Unlock()
+}
+
 func summariesSwapCron(window time.Duration) {
 	for {
 		time.Sleep(window / 2)