metrics: introduce custom registry

Signed-off-by: Mohammed Al Sahaf <msaa1990@gmail.com>

context.go

@@ -21,6 +21,7 @@ import (
 	"log"
 	"log/slog"
 	"reflect"
+	"sync"
 
 	"github.com/caddyserver/certmagic"
 	"github.com/prometheus/client_golang/prometheus"
@@ -49,7 +50,7 @@ type Context struct {
 	ancestry        []Module
 	cleanupFuncs    []func()                // invoked at every config unload
 	exitFuncs       []func(context.Context) // invoked at config unload ONLY IF the process is exiting (EXPERIMENTAL)
-	metricsRegistry *prometheus.Registry
+	metricsRegistry *registryGatherer
 }
 
 // NewContext provides a new context derived from the given
@@ -61,7 +62,8 @@ type Context struct {
 // modules which are loaded will be properly unloaded.
 // See standard library context package's documentation.
 func NewContext(ctx Context) (Context, context.CancelFunc) {
-	newCtx := Context{moduleInstances: make(map[string][]Module), cfg: ctx.cfg, metricsRegistry: prometheus.NewPedanticRegistry()}
+	r := prometheus.NewPedanticRegistry()
+	newCtx := Context{moduleInstances: make(map[string][]Module), cfg: ctx.cfg, metricsRegistry: &registryGatherer{registry: r, gatherer: r, tracker: make(map[string]*sync.Once)}}
 	c, cancel := context.WithCancel(ctx.Context)
 	wrappedCancel := func() {
 		cancel()
@@ -103,7 +105,10 @@ func (ctx *Context) FileSystems() FileSystems {
 
 // Returns the active metrics registry for the context
 // EXPERIMENTAL: This API is subject to change.
-func (ctx *Context) GetMetricsRegistry() *prometheus.Registry {
+func (ctx *Context) GetMetricsRegistry() RegistererGatherer {
+	if ctx.Module() != nil {
+		ctx.metricsRegistry.callerModule = ctx.Module().CaddyModule().String()
+	}
 	return ctx.metricsRegistry
 }
 

metrics.go

@@ -2,8 +2,10 @@ package caddy
 
 import (
 	"net/http"
+	"sync"
 
 	"github.com/prometheus/client_golang/prometheus"
+	io_prometheus_client "github.com/prometheus/client_model/go"
 
 	"github.com/caddyserver/caddy/v2/internal/metrics"
 )
@@ -82,3 +84,53 @@ func (d *delegator) WriteHeader(code int) {
 func (d *delegator) Unwrap() http.ResponseWriter {
 	return d.ResponseWriter
 }
+
+type RegistererGatherer interface {
+	prometheus.Registerer
+	prometheus.Gatherer
+}
+type registryGatherer struct {
+	registry prometheus.Registerer
+	gatherer prometheus.Gatherer
+	tracker  map[string]*sync.Once
+
+	callerModule string
+}
+
+// Gather implements prometheus.Gatherer.
+func (r *registryGatherer) Gather() ([]*io_prometheus_client.MetricFamily, error) {
+	return r.gatherer.Gather()
+}
+
+// MustRegister implements prometheus.Registerer.
+func (r *registryGatherer) MustRegister(cs ...prometheus.Collector) {
+	if _, ok := r.tracker[r.callerModule]; !ok {
+		r.tracker[r.callerModule] = &sync.Once{}
+	}
+	r.tracker[r.callerModule].Do(func() {
+		r.registry.MustRegister(cs...)
+	})
+}
+
+// Register implements prometheus.Registerer.
+func (r *registryGatherer) Register(c prometheus.Collector) error {
+	var err error
+	if _, ok := r.tracker[r.callerModule]; !ok {
+		r.tracker[r.callerModule] = &sync.Once{}
+	}
+	r.tracker[r.callerModule].Do(func() {
+		err = r.registry.Register(c)
+	})
+	return err
+}
+
+// Unregister implements prometheus.Registerer.
+func (r *registryGatherer) Unregister(c prometheus.Collector) bool {
+	delete(r.tracker, r.callerModule)
+	return r.registry.Unregister(c)
+}
+
+var (
+	_ prometheus.Registerer = (*registryGatherer)(nil)
+	_ prometheus.Gatherer   = (*registryGatherer)(nil)
+)

modules/caddyhttp/reverseproxy/metrics.go

@@ -1,7 +1,6 @@
 package reverseproxy
 
 import (
-	"errors"
 	"runtime/debug"
 	"sync"
 	"time"
@@ -19,7 +18,7 @@ var reverseProxyMetrics = struct {
 	logger           *zap.Logger
 }{}
 
-func initReverseProxyMetrics(handler *Handler, registry *prometheus.Registry) {
+func initReverseProxyMetrics(handler *Handler, registry prometheus.Registerer) {
 	const ns, sub = "caddy", "reverse_proxy"
 
 	upstreamsLabels := []string{"upstream"}
@@ -32,17 +31,7 @@ func initReverseProxyMetrics(handler *Handler, registry *prometheus.Registry) {
 		}, upstreamsLabels)
 	})
 
-	// duplicate registration could happen if multiple sites with reverse proxy are configured; so ignore the error because
-	// there's no good way to capture having multiple sites with reverse proxy. If this happens, the metrics will be
-	// registered twice, but the second registration will be ignored.
-	if err := registry.Register(reverseProxyMetrics.upstreamsHealthy); err != nil &&
-		!errors.Is(err, prometheus.AlreadyRegisteredError{
-			ExistingCollector: reverseProxyMetrics.upstreamsHealthy,
-			NewCollector:      reverseProxyMetrics.upstreamsHealthy,
-		}) {
-		panic(err)
-	}
-
+	registry.MustRegister(reverseProxyMetrics.upstreamsHealthy)
 	reverseProxyMetrics.logger = handler.logger.Named("reverse_proxy.metrics")
 }
 

modules/metrics/adminmetrics.go

@@ -18,8 +18,6 @@ import (
 	"errors"
 	"net/http"
 
-	"github.com/prometheus/client_golang/prometheus"
-
 	"github.com/caddyserver/caddy/v2"
 )
 
@@ -33,7 +31,7 @@ func init() {
 // See the Metrics module for a configurable endpoint that is usable if the
 // Admin API is disabled.
 type AdminMetrics struct {
-	registry *prometheus.Registry
+	registry caddy.RegistererGatherer
 
 	metricsHandler http.Handler
 }

modules/metrics/metrics.go

@@ -18,7 +18,6 @@ import (
 	"errors"
 	"net/http"
 
-	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"go.uber.org/zap"
 
@@ -112,7 +111,7 @@ var (
 	_ caddyfile.Unmarshaler       = (*Metrics)(nil)
 )
 
-func createMetricsHandler(logger promhttp.Logger, enableOpenMetrics bool, registry *prometheus.Registry) http.Handler {
+func createMetricsHandler(logger promhttp.Logger, enableOpenMetrics bool, registry caddy.RegistererGatherer) http.Handler {
 	return promhttp.InstrumentMetricHandler(registry,
 		promhttp.HandlerFor(registry, promhttp.HandlerOpts{
 			// will only log errors if logger is non-nil