Prometheus: NewPrometheus constructor shall accept a config struct

prometheus/prometheus.go

@@ -29,6 +29,7 @@ import (
 	"errors"
 	"net/http"
 	"os"
+	"sort"
 	"strconv"
 	"time"
 
@@ -69,30 +70,26 @@ var reqCnt = &Metric{
 	ID:          "reqCnt",
 	Name:        "requests_total",
 	Description: "How many HTTP requests processed, partitioned by status code and HTTP method.",
-	Type:        "counter_vec",
-	Args:        []string{"code", "method", "host", "url"}}
+	Type:        "counter_vec"}
 
 var reqDur = &Metric{
 	ID:          "reqDur",
 	Name:        "request_duration_seconds",
 	Description: "The HTTP request latencies in seconds.",
-	Args:        []string{"code", "method", "host", "url"},
 	Type:        "histogram_vec",
 	Buckets:     reqDurBuckets}
 
 var resSz = &Metric{
 	ID:          "resSz",
 	Name:        "response_size_bytes",
 	Description: "The HTTP response sizes in bytes.",
-	Args:        []string{"code", "method", "host", "url"},
 	Type:        "histogram_vec",
 	Buckets:     resSzBuckets}
 
 var reqSz = &Metric{
 	ID:          "reqSz",
 	Name:        "request_size_bytes",
 	Description: "The HTTP request sizes in bytes.",
-	Args:        []string{"code", "method", "host", "url"},
 	Type:        "histogram_vec",
 	Buckets:     reqSzBuckets}
 
@@ -150,13 +147,31 @@ type Prometheus struct {
 	Subsystem   string
 	Skipper     middleware.Skipper
 
-	RequestCounterURLLabelMappingFunc  RequestCounterLabelMappingFunc
+	labelDescriminators []labelDescriminatorEntry
+	registerer          prometheus.Registerer
+	gatherer            prometheus.Gatherer
+
+	// RequestCounterURLLabelMappingFunc is deprecated, but can still be used
+	// to override how the "url" label in prometheus metrics is calculated.
+	// Moving forward, it is preferred to configure this via the "url"
+	// label descriminator.
+	RequestCounterURLLabelMappingFunc RequestCounterLabelMappingFunc
+
+	// RequestCounterHostLabelMappingFunc is deprecated, but can still be used
+	// to override how the "host" label in prometheus metrics is calculated.
+	// Moving forward, it is preferred to configure this via the "host"
+	// label descriminator.
 	RequestCounterHostLabelMappingFunc RequestCounterLabelMappingFunc
 
 	// Context string to use as a prometheus URL label
 	URLLabelFromContext string
 }
 
+type labelDescriminatorEntry struct {
+	label         string
+	descriminator RequestCounterLabelMappingFunc
+}
+
 // PushGateway contains the configuration for pushing to a Prometheus pushgateway (optional)
 type PushGateway struct {
 	// Push interval in seconds
@@ -171,45 +186,158 @@ type PushGateway struct {
 	Job string
 }
 
-// NewPrometheus generates a new set of metrics with a certain subsystem name
-func NewPrometheus(subsystem string, skipper middleware.Skipper, customMetricsList ...[]*Metric) *Prometheus {
-	var metricsList []*Metric
+// PrometheusConfig defines configuration to use when instantiating a new
+// Prometheus instance. This struct will receive additional members over time
+// so it is strongly recommended to reference fields by name rather than position.
+type PrometheusConfig struct {
+	// MetricPath overrides the default path where prometheus metrics will be
+	// served. If unspecified, the default value of "/metrics" will be used.
+	MetricsPath       string
+	CustomMetricsList []*Metric
+	Skipper           middleware.Skipper
+
+	// LabelDescriminators can be used to specify additional labels to place on
+	// the generated prometheus metrics. Each key in this map will be added
+	// as a label while the values for each of these labels will be generated by
+	// running the associated function on the request context.
+	LabelDescriminators map[string]RequestCounterLabelMappingFunc
+
+	// Registerer sets the prometheus.Registerer instance the middleware
+	// will register these metrics with. If unspecified,
+	// prometheus.DefaultRegisterer will be used.
+	Registerer prometheus.Registerer
+
+	// Gatherer sets the prometheus.Gatherer instance the middleware
+	// will use when generating the metric endpoint handler. If unspecified,
+	// prometheus.DefaultGatherer will be used.
+	Gatherer prometheus.Gatherer
+}
+
+// NewPrometheusWithConfig creates a Prometheus instance using the provided config
+// struct. It should be viewed as a superset of the older NewPrometheus call; both
+// the skipper and the customMetricsList args can be provided via the config struct.
+func NewPrometheusWithConfig(subsystem string, config *PrometheusConfig) *Prometheus {
+	skipper := config.Skipper
 	if skipper == nil {
 		skipper = middleware.DefaultSkipper
 	}
 
-	if len(customMetricsList) > 1 {
-		panic("Too many args. NewPrometheus( string, <optional []*Metric> ).")
-	} else if len(customMetricsList) == 1 {
-		metricsList = customMetricsList[0]
-	}
-
-	metricsList = append(metricsList, standardMetrics...)
+	metricsList := append(config.CustomMetricsList, standardMetrics...)
 
 	p := &Prometheus{
 		MetricsList: metricsList,
-		MetricsPath: defaultMetricPath,
+		MetricsPath: config.MetricsPath,
 		Subsystem:   defaultSubsystem,
-		Skipper:     skipper,
-		RequestCounterURLLabelMappingFunc: func(c echo.Context) string {
-			p := c.Path() // contains route path ala `/users/:id`
-			if p != "" {
-				return p
-			}
-			// as of Echo v4.10.1 path is empty for 404 cases (when router did not find any matching routes)
-			// in this case we use actual path from request to have some distinction in Prometheus
-			return c.Request().URL.Path
-		},
+		Skipper:     config.Skipper,
+		registerer:  config.Registerer,
+		gatherer:    config.Gatherer,
 		RequestCounterHostLabelMappingFunc: func(c echo.Context) string {
 			return c.Request().Host
 		},
 	}
 
+	if p.MetricsPath == "" {
+		p.MetricsPath = defaultMetricPath
+	}
+	if p.Skipper == nil {
+		p.Skipper = middleware.DefaultSkipper
+	}
+	if p.registerer == nil {
+		p.registerer = prometheus.DefaultRegisterer
+	}
+	if p.gatherer == nil {
+		p.gatherer = prometheus.DefaultGatherer
+	}
+
+	// XXX: In order to maintain backwards compatability, the default implementation
+	// of this method needs to reference the closed-over p.URLLabelFromContext. This
+	// allows the user to override thisURLLabelFromContext after construction.
+	p.RequestCounterURLLabelMappingFunc = func(c echo.Context) string {
+		if len(p.URLLabelFromContext) > 0 {
+			u := c.Get(p.URLLabelFromContext)
+			if u == nil {
+				u = "unknown"
+			}
+			return u.(string)
+		}
+		path := c.Path() // contains route path ala `/users/:id`
+		if path != "" {
+			return path
+		}
+		// as of Echo v4.10.1 path is empty for 404 cases (when router did not find any matching routes)
+		// in this case we use actual path from request to have some distinction in Prometheus
+		return c.Request().URL.Path
+	}
+
+	p.labelDescriminators = makeLabelDescriminators(config.LabelDescriminators, p)
+
 	p.registerMetrics(subsystem)
 
 	return p
 }
 
+// NewPrometheus generates a new set of metrics with a certain subsystem name
+func NewPrometheus(subsystem string, skipper middleware.Skipper, customMetricsList ...[]*Metric) *Prometheus {
+
+	var metricsList []*Metric
+	if len(customMetricsList) > 1 {
+		panic("Too many args. NewPrometheus( string, <optional []*Metric> ).")
+	} else if len(customMetricsList) == 1 {
+		metricsList = customMetricsList[0]
+	}
+
+	return NewPrometheusWithConfig(subsystem, &PrometheusConfig{
+		CustomMetricsList: metricsList,
+		Skipper:           skipper,
+	})
+}
+
+func makeLabelDescriminators(labelDescriminators map[string]RequestCounterLabelMappingFunc, p *Prometheus) []labelDescriminatorEntry {
+	// These default label descriminators are present for all metrics even if the user
+	// does not provide their own. They can be overridden if the user chooses to pass
+	// their own implementation in. Note this cannot be a global constant because
+	// we must close over `p.RequestCounterHostLabelMappingFunc` and
+	// `p.RequestCounterURLLabelMapingFunc` in order to maintain backwards compatability.
+	// Consumers should be able to override these members and see the label behavior affected.
+	mappings := map[string]RequestCounterLabelMappingFunc{
+		"code": func(c echo.Context) string {
+			return strconv.Itoa(c.Response().Status)
+		},
+		"method": func(c echo.Context) string {
+			return c.Request().Method
+		},
+		// In order to maintain backwards-compatability with the behavior that
+		// RequestCounterHostLabelMappingFunc can be overridden at any point, we
+		// implement the default implementation as a closure.
+		"host": func(c echo.Context) string {
+			return p.RequestCounterHostLabelMappingFunc(c)
+		},
+		// In order to maintain backwards-compatability with the behavior that
+		// RequestCounterURLLabelMappingFunc can be overridden at any point, we
+		// implement the default implementation as a closure.
+		"url": func(c echo.Context) string {
+			return p.RequestCounterURLLabelMappingFunc(c)
+		},
+	}
+
+	// The base mappings can be
+	for label, impl := range labelDescriminators {
+		mappings[label] = impl
+	}
+
+	descriminatorList := make([]labelDescriminatorEntry, 0, len(mappings))
+	for label, impl := range mappings {
+		descriminatorList = append(descriminatorList, labelDescriminatorEntry{
+			label:         label,
+			descriminator: impl,
+		})
+	}
+	sort.Slice(descriminatorList, func(i, j int) bool {
+		return descriminatorList[i].label > descriminatorList[j].label
+	})
+	return descriminatorList
+}
+
 // SetPushGateway sends metrics to a remote pushgateway exposed on pushGatewayURL
 // every pushInterval. Metrics are fetched from
 func (p *Prometheus) SetPushGateway(pushGatewayURL string, pushInterval time.Duration) {
@@ -244,10 +372,10 @@ func (p *Prometheus) SetPushGatewayJob(j string) {
 // SetMetricsPath set metrics paths
 func (p *Prometheus) SetMetricsPath(e *echo.Echo) {
 	if p.listenAddress != "" {
-		p.router.GET(p.MetricsPath, prometheusHandler())
+		p.router.GET(p.MetricsPath, p.prometheusHandler())
 		p.runServer()
 	} else {
-		e.GET(p.MetricsPath, prometheusHandler())
+		e.GET(p.MetricsPath, p.prometheusHandler())
 	}
 }
 
@@ -376,9 +504,15 @@ func NewMetric(m *Metric, subsystem string) prometheus.Collector {
 
 func (p *Prometheus) registerMetrics(subsystem string) {
 
+	labels := make([]string, len(p.labelDescriminators))
+	for idx, desc := range p.labelDescriminators {
+		labels[idx] = desc.label
+	}
+
 	for _, metricDef := range p.MetricsList {
+		metricDef.Args = labels
 		metric := NewMetric(metricDef, subsystem)
-		if err := prometheus.Register(metric); err != nil {
+		if err := p.registerer.Register(metric); err != nil {
 			log.Errorf("%s could not be registered in Prometheus: %v", metricDef.Name, err)
 		}
 		switch metricDef {
@@ -417,41 +551,41 @@ func (p *Prometheus) HandlerFunc(next echo.HandlerFunc) echo.HandlerFunc {
 		err := next(c)
 
 		status := c.Response().Status
+		var codeOverride string
 		if err != nil {
 			var httpError *echo.HTTPError
 			if errors.As(err, &httpError) {
-				status = httpError.Code
-			}
-			if status == 0 || status == http.StatusOK {
-				status = http.StatusInternalServerError
+				codeOverride = strconv.Itoa(httpError.Code)
+			} else if status == 0 || status == http.StatusOK {
+				codeOverride = strconv.Itoa(http.StatusInternalServerError)
 			}
 		}
 
 		elapsed := float64(time.Since(start)) / float64(time.Second)
 
-		url := p.RequestCounterURLLabelMappingFunc(c)
-		if len(p.URLLabelFromContext) > 0 {
-			u := c.Get(p.URLLabelFromContext)
-			if u == nil {
-				u = "unknown"
+		labelValues := make([]string, len(p.labelDescriminators))
+		for idx, desc := range p.labelDescriminators {
+			labelValues[idx] = desc.descriminator(c)
+			if desc.label == "code" && codeOverride != "" {
+				labelValues[idx] = codeOverride
 			}
-			url = u.(string)
 		}
 
-		statusStr := strconv.Itoa(status)
-		p.reqDur.WithLabelValues(statusStr, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Observe(elapsed)
-		p.reqCnt.WithLabelValues(statusStr, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Inc()
-		p.reqSz.WithLabelValues(statusStr, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Observe(float64(reqSz))
+		p.reqDur.WithLabelValues(labelValues...).Observe(elapsed)
+		p.reqCnt.WithLabelValues(labelValues...).Inc()
+		p.reqSz.WithLabelValues(labelValues...).Observe(float64(reqSz))
 
 		resSz := float64(c.Response().Size)
-		p.resSz.WithLabelValues(statusStr, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Observe(resSz)
+		p.resSz.WithLabelValues(labelValues...).Observe(resSz)
 
 		return err
 	}
 }
 
-func prometheusHandler() echo.HandlerFunc {
-	h := promhttp.Handler()
+func (p *Prometheus) prometheusHandler() echo.HandlerFunc {
+	h := promhttp.HandlerFor(p.gatherer, promhttp.HandlerOpts{
+		Registry: p.registerer,
+	})
 	return func(c echo.Context) error {
 		h.ServeHTTP(c.Response(), c.Request())
 		return nil