Add WrapCollectorWith and WrapCollectorWithPrefix

I found myself in a situation where I had to instantiate multiple
instances of a third party library that registered some metrics and
never un-registered them.

In my app's lifecycle I needed to un-register them and then register
them again, but I couldn't achieve that by wrapping the Registerer.

I decided to register that library's metrics in a separate Registry and
register that Registry, but in order to handle multiple instances of the
library, I also needed to wrap it with labels, and while such
functionality existed already as part of the Registerer wrapping, it
wasn't exposed.

This PR just exposes Collector wrapping as a public function.

Signed-off-by: Oleg Zaytsev <[email protected]>

Author: colega

prometheus/examples_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/prometheus/common/expfmt"
 
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
@@ -784,3 +785,83 @@ func ExampleCollectorFunc() {
 	// Output:
 	// {"name":"http_requests_info","help":"Information about the received HTTP requests.","type":"COUNTER","metric":[{"label":[{"name":"code","value":"200"},{"name":"method","value":"GET"}],"counter":{"value":42}},{"label":[{"name":"code","value":"404"},{"name":"method","value":"POST"}],"counter":{"value":15}}]}
 }
+
+// Using WrapCollectorWith to un-register metrics registered by a third party lib.
+// newThirdPartyLibFoo illustrates a constructor from a third-party lib that does
+// not expose any way to un-register metrics.
+func ExampleWrapCollectorWith() {
+	reg := prometheus.NewRegistry()
+
+	// We want to create two instances of thirdPartyLibFoo, each one wrapped with
+	// its "instance" label.
+	firstReg := prometheus.NewRegistry()
+	_ = newThirdPartyLibFoo(firstReg)
+	firstCollector := prometheus.WrapCollectorWith(prometheus.Labels{"instance": "first"}, firstReg)
+	reg.MustRegister(firstCollector)
+
+	secondReg := prometheus.NewRegistry()
+	_ = newThirdPartyLibFoo(secondReg)
+	secondCollector := prometheus.WrapCollectorWith(prometheus.Labels{"instance": "second"}, secondReg)
+	reg.MustRegister(secondCollector)
+
+	// So far we have illustrated that we can create two instances of thirdPartyLibFoo,
+	// wrapping each one's metrics with some const label.
+	// This is something we could've achieved by doing:
+	// newThirdPartyLibFoo(prometheus.WrapRegistererWith(prometheus.Labels{"instance": "first"}, reg))
+	metricFamilies, err := reg.Gather()
+	if err != nil {
+		panic("unexpected behavior of registry")
+	}
+	fmt.Println("Both instances:")
+	fmt.Println(toNormalizedJSON(sanitizeMetricFamily(metricFamilies[0])))
+
+	// Now we want to unregister first Foo's metrics, and then register them again.
+	// This is not possible by passing a wrapped Registerer to newThirdPartyLibFoo,
+	// because we have already lost track of the registered Collectors,
+	// however since we've collected Foo's metrics in it's own Registry, and we have registered that
+	// as a specific Collector, we can now de-register them:
+	unregistered := reg.Unregister(firstCollector)
+	if !unregistered {
+		panic("unexpected behavior of registry")
+	}
+
+	metricFamilies, err = reg.Gather()
+	if err != nil {
+		panic("unexpected behavior of registry")
+	}
+	fmt.Println("First unregistered:")
+	fmt.Println(toNormalizedJSON(sanitizeMetricFamily(metricFamilies[0])))
+
+	// Now we can create another instance of Foo with {instance: "first"} label again.
+	firstRegAgain := prometheus.NewRegistry()
+	_ = newThirdPartyLibFoo(firstRegAgain)
+	firstCollectorAgain := prometheus.WrapCollectorWith(prometheus.Labels{"instance": "first"}, firstRegAgain)
+	reg.MustRegister(firstCollectorAgain)
+
+	metricFamilies, err = reg.Gather()
+	if err != nil {
+		panic("unexpected behavior of registry")
+	}
+	fmt.Println("Both again:")
+	fmt.Println(toNormalizedJSON(sanitizeMetricFamily(metricFamilies[0])))
+
+	// Output:
+	// Both instances:
+	// {"name":"foo","help":"Registered forever.","type":"GAUGE","metric":[{"label":[{"name":"instance","value":"first"}],"gauge":{"value":1}},{"label":[{"name":"instance","value":"second"}],"gauge":{"value":1}}]}
+	// First unregistered:
+	// {"name":"foo","help":"Registered forever.","type":"GAUGE","metric":[{"label":[{"name":"instance","value":"second"}],"gauge":{"value":1}}]}
+	// Both again:
+	// {"name":"foo","help":"Registered forever.","type":"GAUGE","metric":[{"label":[{"name":"instance","value":"first"}],"gauge":{"value":1}},{"label":[{"name":"instance","value":"second"}],"gauge":{"value":1}}]}
+
+}
+
+func newThirdPartyLibFoo(reg prometheus.Registerer) struct{} {
+	foo := struct{}{}
+	// Register the metrics of the third party lib.
+	c := promauto.With(reg).NewGauge(prometheus.GaugeOpts{
+		Name: "foo",
+		Help: "Registered forever.",
+	})
+	c.Set(1)
+	return foo
+}

prometheus/wrap.go

@@ -63,7 +63,7 @@ func WrapRegistererWith(labels Labels, reg Registerer) Registerer {
 // metric names that are standardized across applications, as that would break
 // horizontal monitoring, for example the metrics provided by the Go collector
 // (see NewGoCollector) and the process collector (see NewProcessCollector). (In
-// fact, those metrics are already prefixed with “go_” or “process_”,
+// fact, those metrics are already prefixed with "go_" or "process_",
 // respectively.)
 //
 // Conflicts between Collectors registered through the original Registerer with
@@ -78,6 +78,40 @@ func WrapRegistererWithPrefix(prefix string, reg Registerer) Registerer {
 	}
 }
 
+// WrapCollectorWith returns a Collector wrapping the provided Collector. The
+// wrapped Collector will add the provided Labels to all Metrics it collects (as
+// ConstLabels). The Metrics collected by the unmodified Collector must not
+// duplicate any of those labels.
+//
+// WrapCollectorWith can be useful to work with multiple instances of a third
+// party library that does not expose enough flexibility on the lifecycle of its
+// registered metrics. 
+// For example, let's say you have a foo.New(reg Registerer) constructor that 
+// registers metrics but never unregisters them, and you want to create multiple
+// instances of foo.Foo with different labels.
+// The way to achieve that, is to create a new Registry, pass it to foo.New,
+// then use WrapCollectorWith to wrap that Registry with the desired labels and
+// register that as a collector in your main Registry. 
+// Then you can un-register the wrapped collector effectively un-registering the
+// metrics registered by foo.New.
+func WrapCollectorWith(labels Labels, c Collector) Collector {
+	return &wrappingCollector{
+		wrappedCollector: c,
+		labels:           labels,
+	}
+}
+
+// WrapCollectorWithPrefix returns a Collector wrapping the provided Collector. The
+// wrapped Collector will add the provided prefix to the name of all Metrics it collects.
+//
+// See the documentation of WrapCollectorWith for more details on the use case.
+func WrapCollectorWithPrefix(prefix string, c Collector) Collector {
+	return &wrappingCollector{
+		wrappedCollector: c,
+		prefix:          prefix,
+	}
+}
+
 type wrappingRegisterer struct {
 	wrappedRegisterer Registerer
 	prefix            string

prometheus/wrap_test.go

@@ -43,7 +43,7 @@ func toMetricFamilies(cs ...Collector) []*dto.MetricFamily {
 	return out
 }
 
-func TestWrap(t *testing.T) {
+func TestWrapRegisterer(t *testing.T) {
 	now := time.Now()
 	nowFn := func() time.Time { return now }
 	simpleCnt := NewCounter(CounterOpts{
@@ -306,28 +306,34 @@ func TestWrap(t *testing.T) {
 			if !s.gatherFails && err != nil {
 				t.Fatal("gathering failed:", err)
 			}
-			if len(wantMF) != len(gotMF) {
-				t.Fatalf("Expected %d metricFamilies, got %d", len(wantMF), len(gotMF))
+			assertEqualMFs(t, wantMF, gotMF)
+		})
+	}
+}
+
+func assertEqualMFs(t *testing.T, wantMF, gotMF []*dto.MetricFamily) {
+	t.Helper()
+
+	if len(wantMF) != len(gotMF) {
+		t.Fatalf("Expected %d metricFamilies, got %d", len(wantMF), len(gotMF))
+	}
+	for i := range gotMF {
+		if !proto.Equal(gotMF[i], wantMF[i]) {
+			var want, got []string
+
+			for i, mf := range wantMF {
+				want = append(want, fmt.Sprintf("%3d: %s", i, mf))
 			}
-			for i := range gotMF {
-				if !proto.Equal(gotMF[i], wantMF[i]) {
-					var want, got []string
-
-					for i, mf := range wantMF {
-						want = append(want, fmt.Sprintf("%3d: %s", i, mf))
-					}
-					for i, mf := range gotMF {
-						got = append(got, fmt.Sprintf("%3d: %s", i, mf))
-					}
-
-					t.Fatalf(
-						"unexpected output of gathering:\n\nWANT:\n%s\n\nGOT:\n%s\n",
-						strings.Join(want, "\n"),
-						strings.Join(got, "\n"),
-					)
-				}
+			for i, mf := range gotMF {
+				got = append(got, fmt.Sprintf("%3d: %s", i, mf))
 			}
-		})
+
+			t.Fatalf(
+				"unexpected output of gathering:\n\nWANT:\n%s\n\nGOT:\n%s\n",
+				strings.Join(want, "\n"),
+				strings.Join(got, "\n"),
+			)
+		}
 	}
 }
 
@@ -339,3 +345,124 @@ func TestNil(t *testing.T) {
 		t.Fatal("registering failed:", err)
 	}
 }
+
+func TestWrapCollector(t *testing.T) {
+	t.Run("can be registered and un-registered", func(t *testing.T) {
+		inner := NewPedanticRegistry()
+		g := NewGauge(GaugeOpts{Name: "testing"})
+		g.Set(42)
+		err := inner.Register(g)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		wrappedWithLabels := WrapCollectorWith(Labels{"lbl": "1"}, inner)
+		wrappedWithPrefix := WrapCollectorWithPrefix("prefix", inner)
+		reg := NewPedanticRegistry()
+		err = reg.Register(wrappedWithLabels)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+		err = reg.Register(wrappedWithPrefix)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		gathered, err := reg.Gather()
+		if err != nil {
+			t.Fatal("gathering failed:", err)
+		}
+
+		lg := NewGauge(GaugeOpts{Name: "testing", ConstLabels: Labels{"lbl": "1"}})
+		lg.Set(42)
+		pg := NewGauge(GaugeOpts{Name: "prefixtesting"})
+		pg.Set(42)
+		expected := toMetricFamilies(lg, pg)
+		assertEqualMFs(t, expected, gathered)
+
+		if !reg.Unregister(wrappedWithLabels) {
+			t.Fatal("unregistering failed")
+		}
+		if !reg.Unregister(wrappedWithPrefix) {
+			t.Fatal("unregistering failed")
+		}
+
+		gathered, err = reg.Gather()
+		if err != nil {
+			t.Fatal("gathering failed:", err)
+		}
+		if len(gathered) != 0 {
+			t.Fatalf("expected 0 metric families, got %d", len(gathered))
+		}
+	})
+
+	t.Run("can wrap same collector twice", func(t *testing.T) {
+		inner := NewPedanticRegistry()
+		g := NewGauge(GaugeOpts{Name: "testing"})
+		g.Set(42)
+		err := inner.Register(g)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		wrapped := WrapCollectorWith(Labels{"lbl": "1"}, inner)
+		reg := NewPedanticRegistry()
+		err = reg.Register(wrapped)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		wrapped2 := WrapCollectorWith(Labels{"lbl": "2"}, inner)
+		err = reg.Register(wrapped2)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		gathered, err := reg.Gather()
+		if err != nil {
+			t.Fatal("gathering failed:", err)
+		}
+
+		lg := NewGauge(GaugeOpts{Name: "testing", ConstLabels: Labels{"lbl": "1"}})
+		lg.Set(42)
+		lg2 := NewGauge(GaugeOpts{Name: "testing", ConstLabels: Labels{"lbl": "2"}})
+		lg2.Set(42)
+		expected := toMetricFamilies(lg, lg2)
+		assertEqualMFs(t, expected, gathered)
+	})
+
+	t.Run("can be registered again after un-registering", func(t *testing.T) {
+		inner := NewPedanticRegistry()
+		g := NewGauge(GaugeOpts{Name: "testing"})
+		g.Set(42)
+		err := inner.Register(g)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		wrapped := WrapCollectorWith(Labels{"lbl": "1"}, inner)
+		reg := NewPedanticRegistry()
+		err = reg.Register(wrapped)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		if !reg.Unregister(wrapped) {
+			t.Fatal("unregistering failed")
+		}
+		err = reg.Register(wrapped)
+		if err != nil {
+			t.Fatal("registering failed:", err)
+		}
+
+		gathered, err := reg.Gather()
+		if err != nil {
+			t.Fatal("gathering failed:", err)
+		}
+
+		lg := NewGauge(GaugeOpts{Name: "testing", ConstLabels: Labels{"lbl": "1"}})
+		lg.Set(42)
+		expected := toMetricFamilies(lg)
+		assertEqualMFs(t, expected, gathered)
+	})
+}