From c6c6640e5b9633416791ed985f3bd4f83318261f Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Mon, 8 Aug 2022 17:13:46 +0300
Subject: [PATCH] add RegisterSet() and UnregisterSet() functions for
 registering the set in metrics export via global WritePrometheus() call

---
 metrics.go      | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 metrics_test.go | 26 ++++++++++++++++++++++++++
 set.go          |  2 ++
 3 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/metrics.go b/metrics.go
index 087026d..12537dd 100644
--- a/metrics.go
+++ b/metrics.go
@@ -14,6 +14,9 @@ package metrics
 
 import (
 	"io"
+	"sort"
+	"sync"
+	"unsafe"
 )
 
 type namedMetric struct {
@@ -27,7 +30,34 @@ type metric interface {
 
 var defaultSet = NewSet()
 
-// WritePrometheus writes all the registered metrics in Prometheus format to w.
+func init() {
+	RegisterSet(defaultSet)
+}
+
+var (
+	registeredSets = make(map[*Set]struct{})
+	registeredSetsLock sync.Mutex
+)
+
+// RegisterSet registers the given set s for metrics export via global WritePrometheus() call.
+//
+// See also UnregisterSet.
+func RegisterSet(s *Set) {
+	registeredSetsLock.Lock()
+	registeredSets[s] = struct{}{}
+	registeredSetsLock.Unlock()
+}
+
+// UnregisterSet stops exporting metrics for the given s via global WritePrometheus() call.
+func UnregisterSet(s *Set) {
+	registeredSetsLock.Lock()
+	delete(registeredSets, s)
+	registeredSetsLock.Unlock()
+}
+
+// WritePrometheus writes all the metrics from default set and all the registered sets in Prometheus format to w.
+//
+// Additional sets can be registered via RegisterSet() call.
 //
 // If exposeProcessMetrics is true, then various `go_*` and `process_*` metrics
 // are exposed for the current process.
@@ -39,7 +69,19 @@ var defaultSet = NewSet()
 //     })
 //
 func WritePrometheus(w io.Writer, exposeProcessMetrics bool) {
-	defaultSet.WritePrometheus(w)
+	registeredSetsLock.Lock()
+	sets := make([]*Set, 0, len(registeredSets))
+	for s := range registeredSets {
+		sets = append(sets, s)
+	}
+	registeredSetsLock.Unlock()
+
+	sort.Slice(sets, func(i, j int) bool {
+		return uintptr(unsafe.Pointer(sets[i])) < uintptr(unsafe.Pointer(sets[j]))
+	})
+	for _, s := range sets {
+		s.WritePrometheus(w)
+	}
 	if exposeProcessMetrics {
 		WriteProcessMetrics(w)
 	}
diff --git a/metrics_test.go b/metrics_test.go
index baa624c..0af0a36 100644
--- a/metrics_test.go
+++ b/metrics_test.go
@@ -3,10 +3,36 @@ package metrics
 import (
 	"bytes"
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 )
 
+func TestRegisterUnregisterSet(t *testing.T) {
+	const metricName = "metric_from_set"
+	const metricValue = 123
+	s := NewSet()
+	c := s.NewCounter(metricName)
+	c.Set(metricValue)
+
+	RegisterSet(s)
+	var bb bytes.Buffer
+	WritePrometheus(&bb, false)
+	data := bb.String()
+	expectedLine := fmt.Sprintf("%s %d\n", metricName, metricValue)
+	if !strings.Contains(data, expectedLine) {
+		t.Fatalf("missing %q in\n%s", expectedLine, data)
+	}
+
+	UnregisterSet(s)
+	bb.Reset()
+	WritePrometheus(&bb, false)
+	data = bb.String()
+	if strings.Contains(data, expectedLine) {
+		t.Fatalf("unepected %q in\n%s", expectedLine, data)
+	}
+}
+
 func TestInvalidName(t *testing.T) {
 	f := func(name string) {
 		t.Helper()
diff --git a/set.go b/set.go
index ae55bb7..8384ecc 100644
--- a/set.go
+++ b/set.go
@@ -22,6 +22,8 @@ type Set struct {
 }
 
 // NewSet creates new set of metrics.
+//
+// Pass the set to RegisterSet() function in order to export its metrics via global WritePrometheus() call.
 func NewSet() *Set {
 	return &Set{
 		m: make(map[string]*namedMetric),