From fdfd428a625d978305138e460757ca963ce67128 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Tue, 16 Jan 2024 00:06:40 +0200
Subject: [PATCH] Add Set.RegisterMetricsWriter() function for registering
 user-defined callbacks for arbitrary metrics generation in Prometheus text
 exposition format

---
 metrics.go      | 16 +++++++++++++++-
 metrics_test.go | 24 ++++++++++++++++++++++++
 set.go          | 29 +++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/metrics.go b/metrics.go
index c0efc58..6dd351d 100644
--- a/metrics.go
+++ b/metrics.go
@@ -62,9 +62,21 @@ func UnregisterSet(s *Set) {
 	registeredSetsLock.Unlock()
 }
 
-// WritePrometheus writes all the metrics from default set and all the registered sets in Prometheus format to w.
+// RegisterMetricsWriter registers writeMetrics callback for including metrics in the output generated by WritePrometheus.
+//
+// The writeMetrics callback must write metrics to w in Prometheus text exposition format without timestamps and trailing comments.
+// The last line generated by writeMetrics must end with \n.
+// See https://github.com/prometheus/docs/blob/main/content/docs/instrumenting/exposition_formats.md#text-based-format
+//
+// It is OK to register multiple writeMetrics callbacks - all of them will be called sequentially for gererating the output at WritePrometheus.
+func RegisterMetricsWriter(writeMetrics func(w io.Writer)) {
+	defaultSet.RegisterMetricsWriter(writeMetrics)
+}
+
+// WritePrometheus writes all the metrics in Prometheus format from the default set, all the added sets and metrics writers to w.
 //
 // Additional sets can be registered via RegisterSet() call.
+// Additional metric writers can be registered via RegisterMetricsWriter() call.
 //
 // If exposeProcessMetrics is true, then various `go_*` and `process_*` metrics
 // are exposed for the current process.
@@ -232,6 +244,8 @@ func UnregisterMetric(name string) bool {
 }
 
 // UnregisterAllMetrics unregisters all the metrics from default set.
+//
+// It also unregisters writeMetrics callbacks passed to RegisterMetricsWriter.
 func UnregisterAllMetrics() {
 	defaultSet.UnregisterAllMetrics()
 }
diff --git a/metrics_test.go b/metrics_test.go
index 7c33dab..7ec4947 100644
--- a/metrics_test.go
+++ b/metrics_test.go
@@ -3,6 +3,7 @@ package metrics
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"strings"
 	"testing"
 	"time"
@@ -108,6 +109,29 @@ func TestUnregisterAllMetrics(t *testing.T) {
 	}
 }
 
+func TestRegisterMetricsWriter(t *testing.T) {
+	RegisterMetricsWriter(func(w io.Writer) {
+		WriteCounterUint64(w, `counter{label="abc"}`, 1234)
+		WriteGaugeFloat64(w, `gauge{a="b",c="d"}`, -34.43)
+	})
+
+	var bb bytes.Buffer
+	WritePrometheus(&bb, false)
+	data := bb.String()
+
+	UnregisterAllMetrics()
+
+	expectedLine := fmt.Sprintf(`counter{label="abc"} 1234` + "\n")
+	if !strings.Contains(data, expectedLine) {
+		t.Fatalf("missing %q in\n%s", expectedLine, data)
+	}
+
+	expectedLine = fmt.Sprintf(`gauge{a="b",c="d"} -34.43` + "\n")
+	if !strings.Contains(data, expectedLine) {
+		t.Fatalf("missing %q in\n%s", expectedLine, data)
+	}
+}
+
 func TestRegisterUnregisterSet(t *testing.T) {
 	const metricName = "metric_from_set"
 	const metricValue = 123
diff --git a/set.go b/set.go
index 50a095b..868a01c 100644
--- a/set.go
+++ b/set.go
@@ -19,6 +19,8 @@ type Set struct {
 	a         []*namedMetric
 	m         map[string]*namedMetric
 	summaries []*Summary
+
+	metricsWriters []func(w io.Writer)
 }
 
 // NewSet creates new set of metrics.
@@ -45,6 +47,7 @@ func (s *Set) WritePrometheus(w io.Writer) {
 		sort.Slice(s.a, lessFunc)
 	}
 	sa := append([]*namedMetric(nil), s.a...)
+	metricsWriters := s.metricsWriters
 	s.mu.Unlock()
 
 	prevMetricFamily := ""
@@ -61,6 +64,10 @@ func (s *Set) WritePrometheus(w io.Writer) {
 		nm.metric.marshalTo(nm.name, &bb)
 	}
 	w.Write(bb.Bytes())
+
+	for _, writeMetrics := range metricsWriters {
+		writeMetrics(w)
+	}
 }
 
 // NewHistogram creates and returns new histogram in s with the given name.
@@ -523,14 +530,22 @@ func (s *Set) unregisterMetricLocked(nm *namedMetric) bool {
 }
 
 // UnregisterAllMetrics de-registers all metrics registered in s.
+//
+// It also de-registers writeMetrics callbacks passed to RegisterMetricsWriter.
 func (s *Set) UnregisterAllMetrics() {
 	metricNames := s.ListMetricNames()
 	for _, name := range metricNames {
 		s.UnregisterMetric(name)
 	}
+
+	s.mu.Lock()
+	s.metricsWriters = nil
+	s.mu.Unlock()
 }
 
 // ListMetricNames returns sorted list of all the metrics in s.
+//
+// The returned list doesn't include metrics generated by metricsWriter passed to RegisterMetricsWriter.
 func (s *Set) ListMetricNames() []string {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -544,3 +559,17 @@ func (s *Set) ListMetricNames() []string {
 	sort.Strings(metricNames)
 	return metricNames
 }
+
+// RegisterMetricsWriter registers writeMetrics callback for including metrics in the output generated by s.WritePrometheus.
+//
+// The writeMetrics callback must write metrics to w in Prometheus text exposition format without timestamps and trailing comments.
+// The last line generated by writeMetrics must end with \n.
+// See https://github.com/prometheus/docs/blob/main/content/docs/instrumenting/exposition_formats.md#text-based-format
+//
+// It is OK to reguster multiple writeMetrics callbacks - all of them will be called sequentially for gererating the output at s.WritePrometheus.
+func (s *Set) RegisterMetricsWriter(writeMetrics func(w io.Writer)) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.metricsWriters = append(s.metricsWriters, writeMetrics)
+}