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),