From 900c8926252653cb4c241db971af522f9cc0c3b3 Mon Sep 17 00:00:00 2001
From: Roland Lammel <roland@lammel.net>
Date: Wed, 26 Feb 2020 18:52:38 +0100
Subject: [PATCH] Allow unregistration of metrics (#6)

* Allow unregistration of metrics

* Provide metrics.WritePrometheusSet to write using a dedicated metric set

* Cleanup unregister metrics patch, add generic WriteProcessMetrics
---
 metrics.go  | 21 +++++++++++++++++--
 set.go      | 25 +++++++++++++++++++++++
 set_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 2 deletions(-)

diff --git a/metrics.go b/metrics.go
index 8fc8f6e..b11e1b0 100644
--- a/metrics.go
+++ b/metrics.go
@@ -41,7 +41,24 @@ var defaultSet = NewSet()
 func WritePrometheus(w io.Writer, exposeProcessMetrics bool) {
 	defaultSet.WritePrometheus(w)
 	if exposeProcessMetrics {
-		writeGoMetrics(w)
-		writeProcessMetrics(w)
+		WriteProcessMetrics(w)
 	}
 }
+
+// WriteProcessMetrics writes additional process metrics in Prometheus format to w.
+//
+// Various `go_*` and `process_*` metrics are exposed for the currently
+// running process.
+//
+// The WriteProcessMetrics func is usually called in combination with writing Set metrics
+// inside "/metrics" handler:
+//
+//     http.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) {
+//         mySet.WritePrometheus(w)
+//         metrics.WriteProcessMetrics(w)
+//     })
+//
+func WriteProcessMetrics(w io.Writer) {
+	writeGoMetrics(w)
+	writeProcessMetrics(w)
+}
diff --git a/set.go b/set.go
index 3b37e91..18021a4 100644
--- a/set.go
+++ b/set.go
@@ -438,3 +438,28 @@ func (s *Set) registerMetric(name string, m metric) {
 		panic(fmt.Errorf("BUG: metric %q is already registered", name))
 	}
 }
+
+// UnregisterMetric will remove a registered metric
+func (s *Set) UnregisterMetric(name string) bool {
+	s.mu.Lock()
+	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:]...)
+			}
+		}
+	}
+	s.mu.Unlock()
+	return ok
+}
+
+// ListMetricNames will return a list of all registered metrics
+func (s *Set) ListMetricNames() []string {
+	var list []string
+	for name := range s.m {
+		list = append(list, name)
+	}
+	return list
+}
diff --git a/set_test.go b/set_test.go
index 3187e4c..f1952d8 100644
--- a/set_test.go
+++ b/set_test.go
@@ -34,3 +34,61 @@ func TestNewSet(t *testing.T) {
 		}
 	}
 }
+
+func TestListMetricNames(t *testing.T) {
+	s := NewSet()
+	expect := []string{"cnt1", "cnt2", "cnt3"}
+	// Initialize a few counters
+	for _, n := range expect {
+		c := s.NewCounter(n)
+		c.Inc()
+	}
+
+	list := s.ListMetricNames()
+
+	if len(list) != len(expect) {
+		t.Fatalf("Metrics count is wrong for listing")
+	}
+	for _, e := range expect {
+		found := false
+		for _, n := range list {
+			if e == n {
+				found = true
+			}
+		}
+		if !found {
+			t.Fatalf("Metric %s not found in listing", e)
+		}
+	}
+}
+
+func TestUnregisterMetric(t *testing.T) {
+	s := NewSet()
+	// Initialize a few counters
+	for i := 0; i < 5; i++ {
+		c := s.NewCounter(fmt.Sprintf("counter_%d", i))
+		c.Inc()
+	}
+	// Unregister existing counters
+	ok := s.UnregisterMetric("counter_1")
+	if !ok {
+		t.Fatalf("Metric counter_1 should return true for deregistering")
+	}
+
+	// Unregister twice must return false
+	ok = s.UnregisterMetric("counter_1")
+	if ok {
+		t.Fatalf("Metric counter_1 should not return false on unregister twice")
+	}
+
+	// Validate counters are removed
+	ok = false
+	for _, n := range s.ListMetricNames() {
+		if n == "counter_1" {
+			ok = true
+		}
+	}
+	if ok {
+		t.Fatalf("Metric counter_1 and counter_3 must not be listed anymore after unregister")
+	}
+}