From fe7c2061a25eaef54ccf4eccf32f89d5cf96aed3 Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 3 Aug 2022 11:00:14 -0400
Subject: [PATCH] prometheus: implement Collector interface for Registry

This change allows Registries to be used as Collectors.

This enables new instances of Registry to be passed to ephemeral
subroutines for collecting metrics from subroutines which are still
running:

```go
package main

import (
  "fmt"

  "github.com/prometheus/client_golang/prometheus"
)

func main() {
  globalReg := prometheus.NewRegistry()

  for i := 0; i < 100; i++ {
    workerReg := prometheus.WrapRegistererWith(prometheus.Labels{
      // Add an ID label so registered metrics from workers don't
      // collide.
      "worker_id": fmt.Sprintf("%d", i),
    }, prometheus.NewRegistry()

    globalReg.MustRegister(workerReg)

    go func(i int) {
      runWorker(workerReg)

      // Unregister any metrics the worker may have created.
      globalReg.Unregister(workerReg)
    }(i)
  }
}

// runWorker runs a worker, registering worker-specific metrics.
func runWorker(reg *prometheus.Registry) {
  // ... register metrics ...
  // ... do work ...
}
```

This change makes it easier to avoid leaking metrics from subroutines
which do not consistently properly unregister metrics.
---
 prometheus/registry.go | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/prometheus/registry.go b/prometheus/registry.go
index 325f665ff..ee61ebbe4 100644
--- a/prometheus/registry.go
+++ b/prometheus/registry.go
@@ -252,9 +252,9 @@ func (errs MultiError) MaybeUnwrap() error {
 }
 
 // Registry registers Prometheus collectors, collects their metrics, and gathers
-// them into MetricFamilies for exposition. It implements both Registerer and
-// Gatherer. The zero value is not usable. Create instances with NewRegistry or
-// NewPedanticRegistry.
+// them into MetricFamilies for exposition. It implements both Registerer,
+// Gatherer, and Collector. The zero value is not usable. Create instances with
+// NewRegistry or NewPedanticRegistry.
 type Registry struct {
 	mtx                   sync.RWMutex
 	collectorsByID        map[uint64]Collector // ID is a hash of the descIDs.
@@ -556,6 +556,31 @@ func (r *Registry) Gather() ([]*dto.MetricFamily, error) {
 	return internal.NormalizeMetricFamilies(metricFamiliesByName), errs.MaybeUnwrap()
 }
 
+// Describe implements Collector.
+func (r *Registry) Describe(ch chan<- *Desc) {
+	r.mtx.RLock()
+	defer r.mtx.RUnlock()
+
+	// Only report the checked Collectors; unchecked collectors don't report any
+	// Desc.
+	for _, c := range r.collectorsByID {
+		c.Describe(ch)
+	}
+}
+
+// Collect implements Collector.
+func (r *Registry) Collect(ch chan<- Metric) {
+	r.mtx.RLock()
+	defer r.mtx.RUnlock()
+
+	for _, c := range r.collectorsByID {
+		c.Collect(ch)
+	}
+	for _, c := range r.uncheckedCollectors {
+		c.Collect(ch)
+	}
+}
+
 // WriteToTextfile calls Gather on the provided Gatherer, encodes the result in the
 // Prometheus text format, and writes it to a temporary file. Upon success, the
 // temporary file is renamed to the provided filename.