Add UnregisterAllMetrics() and GetDefaultSet() functions

While at it, make sure that ListMetricNames() returns sorted list of the registered metrics

metrics.go

@@ -22,6 +22,7 @@ import (
 type namedMetric struct {
 	name   string
 	metric metric
+	isAux  bool
 }
 
 type metric interface {
@@ -49,6 +50,8 @@ func RegisterSet(s *Set) {
 }
 
 // UnregisterSet stops exporting metrics for the given s via global WritePrometheus() call.
+//
+// Call s.UnregisterAllMetrics() after unregistering s if it is no longer used.
 func UnregisterSet(s *Set) {
 	registeredSetsLock.Lock()
 	delete(registeredSets, s)
@@ -180,11 +183,23 @@ func WriteFDMetrics(w io.Writer) {
 }
 
 // UnregisterMetric removes metric with the given name from default set.
+//
+// See also UnregisterAllMetrics.
 func UnregisterMetric(name string) bool {
 	return defaultSet.UnregisterMetric(name)
 }
 
-// ListMetricNames returns a list of all the metric names from default set.
+// UnregisterAllMetrics unregisters all the metrics from default set.
+func UnregisterAllMetrics() {
+	defaultSet.UnregisterAllMetrics()
+}
+
+// ListMetricNames returns sorted list of all the metric names from default set.
 func ListMetricNames() []string {
 	return defaultSet.ListMetricNames()
 }
+
+// GetDefaultSet returns the default metrics set.
+func GetDefaultSet() *Set {
+	return defaultSet
+}

metrics_test.go

@@ -8,6 +8,31 @@ import (
 	"time"
 )
 
+func TestGetDefaultSet(t *testing.T) {
+	s := GetDefaultSet()
+	if s != defaultSet {
+		t.Fatalf("GetDefaultSet must return defaultSet=%p, but returned %p", defaultSet, s)
+	}
+}
+
+func TestUnregisterAllMetrics(t *testing.T) {
+	for j := 0; j < 3; j++ {
+		for i := 0; i < 10; i++ {
+			_ = NewCounter(fmt.Sprintf("counter_%d", i))
+			_ = NewSummary(fmt.Sprintf("summary_%d", i))
+			_ = NewHistogram(fmt.Sprintf("histogram_%d", i))
+			_ = NewGauge(fmt.Sprintf("gauge_%d", i), func() float64 { return 0 })
+		}
+		if mns := ListMetricNames(); len(mns) == 0 {
+			t.Fatalf("unexpected empty list of metrics on iteration %d", j)
+		}
+		UnregisterAllMetrics()
+		if mns := ListMetricNames(); len(mns) != 0 {
+			t.Fatalf("unexpected metric names after UnregisterAllMetrics call on iteration %d: %q", j, mns)
+		}
+	}
+}
+
 func TestRegisterUnregisterSet(t *testing.T) {
 	const metricName = "metric_from_set"
 	const metricValue = 123

set.go

@@ -336,7 +336,7 @@ func (s *Set) NewSummaryExt(name string, window time.Duration, quantiles []float
 	// checks in tests
 	defer s.mu.Unlock()
 
-	s.mustRegisterLocked(name, sm)
+	s.mustRegisterLocked(name, sm, false)
 	registerSummaryLocked(sm)
 	s.registerSummaryQuantilesLocked(name, sm)
 	s.summaries = append(s.summaries, sm)
@@ -420,7 +420,7 @@ func (s *Set) registerSummaryQuantilesLocked(name string, sm *Summary) {
 			sm:  sm,
 			idx: i,
 		}
-		s.mustRegisterLocked(quantileValueName, qv)
+		s.mustRegisterLocked(quantileValueName, qv, true)
 	}
 }
 
@@ -432,18 +432,19 @@ func (s *Set) registerMetric(name string, m metric) {
 	// defer will unlock in case of panic
 	// checks in test
 	defer s.mu.Unlock()
-	s.mustRegisterLocked(name, m)
+	s.mustRegisterLocked(name, m, false)
 }
 
-// mustRegisterLocked registers given metric with
-// the given name. Panics if the given name was
-// already registered before.
-func (s *Set) mustRegisterLocked(name string, m metric) {
+// mustRegisterLocked registers given metric with the given name.
+//
+// Panics if the given name was already registered before.
+func (s *Set) mustRegisterLocked(name string, m metric, isAux bool) {
 	nm, ok := s.m[name]
 	if !ok {
 		nm = &namedMetric{
 			name:   name,
 			metric: m,
+			isAux:  isAux,
 		}
 		s.m[name] = nm
 		s.a = append(s.a, nm)
@@ -465,8 +466,16 @@ func (s *Set) UnregisterMetric(name string) bool {
 	if !ok {
 		return false
 	}
-	m := nm.metric
+	if nm.isAux {
+		// Do not allow deleting auxiliary metrics such as summary_metric{quantile="..."}
+		// Such metrics must be deleted via parent metric name, e.g. summary_metric .
+		return false
+	}
+	return s.unregisterMetricLocked(nm)
+}
 
+func (s *Set) unregisterMetricLocked(nm *namedMetric) bool {
+	name := nm.name
 	delete(s.m, name)
 
 	deleteFromList := func(metricName string) {
@@ -482,9 +491,9 @@ func (s *Set) UnregisterMetric(name string) bool {
 	// remove metric from s.a
 	deleteFromList(name)
 
-	sm, ok := m.(*Summary)
+	sm, ok := nm.metric.(*Summary)
 	if !ok {
-		// There is no need in cleaning up summary.
+		// There is no need in cleaning up non-summary metrics.
 		return true
 	}
 
@@ -511,13 +520,25 @@ func (s *Set) UnregisterMetric(name string) bool {
 	return true
 }
 
-// ListMetricNames returns a list of all the metrics in s.
+// UnregisterAllMetrics de-registers all metrics registered in s.
+func (s *Set) UnregisterAllMetrics() {
+	metricNames := s.ListMetricNames()
+	for _, name := range metricNames {
+		s.UnregisterMetric(name)
+	}
+}
+
+// ListMetricNames returns sorted list of all the metrics in s.
 func (s *Set) ListMetricNames() []string {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	var list []string
-	for name := range s.m {
-		list = append(list, name)
+	metricNames := make([]string, 0, len(s.m))
+	for _, nm := range s.m {
+		if nm.isAux {
+			continue
+		}
+		metricNames = append(metricNames, nm.name)
 	}
-	return list
+	sort.Strings(metricNames)
+	return metricNames
 }

set_test.go

@@ -64,6 +64,27 @@ func TestSetListMetricNames(t *testing.T) {
 	}
 }
 
+func TestSetUnregisterAllMetrics(t *testing.T) {
+	s := NewSet()
+	for j := 0; j < 3; j++ {
+		expectedMetricsCount := 0
+		for i := 0; i < 10; i++ {
+			_ = s.NewCounter(fmt.Sprintf("counter_%d", i))
+			_ = s.NewSummary(fmt.Sprintf("summary_%d", i))
+			_ = s.NewHistogram(fmt.Sprintf("histogram_%d", i))
+			_ = s.NewGauge(fmt.Sprintf("gauge_%d", i), func() float64 { return 0 })
+			expectedMetricsCount += 4
+		}
+		if mns := s.ListMetricNames(); len(mns) != expectedMetricsCount {
+			t.Fatalf("unexpected number of metric names on iteration %d; got %d; want %d;\nmetric names:\n%q", j, len(mns), expectedMetricsCount, mns)
+		}
+		s.UnregisterAllMetrics()
+		if mns := s.ListMetricNames(); len(mns) != 0 {
+			t.Fatalf("unexpected metric names after UnregisterAllMetrics call on iteration %d: %q", j, mns)
+		}
+	}
+}
+
 func TestSetUnregisterMetric(t *testing.T) {
 	s := NewSet()
 	const cName, smName = "counter_1", "summary_1"