Add customized labels based on HTTP header

metrics/metrics.go

@@ -16,6 +16,8 @@ type HTTPReqProperties struct {
 	Method string
 	// Code is the response of the request.
 	Code string
+	// CustomizedLabels is the customized header
+	Labels map[string]string
 }
 
 // HTTPProperties are the metric properties for the global server metrics.

metrics/prometheus/prometheus.go

@@ -2,6 +2,7 @@ package prometheus
 
 import (
 	"context"
+	"sort"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -30,6 +31,8 @@ type Config struct {
 	MethodLabel string
 	// ServiceLabel is the name that will be set to the service label, by default is `service`.
 	ServiceLabel string
+	// CustomizedLabels is the labels from the header
+	CustomizedLabels []string
 }
 
 func (c *Config) defaults() {
@@ -73,22 +76,29 @@ type recorder struct {
 func NewRecorder(cfg Config) metrics.Recorder {
 	cfg.defaults()
 
+	customLabels := []string{}
+	if cfg.CustomizedLabels != nil && len(cfg.CustomizedLabels) > 0 {
+		customLabels = cfg.CustomizedLabels
+		sort.Strings(customLabels)
+	}
+
+	labels := append([]string{cfg.ServiceLabel, cfg.HandlerIDLabel, cfg.MethodLabel, cfg.StatusCodeLabel}, customLabels...)
 	r := &recorder{
 		httpRequestDurHistogram: prometheus.NewHistogramVec(prometheus.HistogramOpts{
 			Namespace: cfg.Prefix,
 			Subsystem: "http",
 			Name:      "request_duration_seconds",
 			Help:      "The latency of the HTTP requests.",
 			Buckets:   cfg.DurationBuckets,
-		}, []string{cfg.ServiceLabel, cfg.HandlerIDLabel, cfg.MethodLabel, cfg.StatusCodeLabel}),
+		}, labels),
 
 		httpResponseSizeHistogram: prometheus.NewHistogramVec(prometheus.HistogramOpts{
 			Namespace: cfg.Prefix,
 			Subsystem: "http",
 			Name:      "response_size_bytes",
 			Help:      "The size of the HTTP responses.",
 			Buckets:   cfg.SizeBuckets,
-		}, []string{cfg.ServiceLabel, cfg.HandlerIDLabel, cfg.MethodLabel, cfg.StatusCodeLabel}),
+		}, labels),
 
 		httpRequestsInflight: prometheus.NewGaugeVec(prometheus.GaugeOpts{
 			Namespace: cfg.Prefix,
@@ -107,12 +117,27 @@ func NewRecorder(cfg Config) metrics.Recorder {
 	return r
 }
 
+func (r recorder) getLabelValues(p metrics.HTTPReqProperties) []string {
+	keys := []string{}
+	for key, _ := range p.Labels {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	labels := []string{p.Service, p.ID, p.Method, p.Code}
+	for _, key := range keys {
+		labels = append(labels, p.Labels[key])
+	}
+	return labels
+}
+
 func (r recorder) ObserveHTTPRequestDuration(_ context.Context, p metrics.HTTPReqProperties, duration time.Duration) {
-	r.httpRequestDurHistogram.WithLabelValues(p.Service, p.ID, p.Method, p.Code).Observe(duration.Seconds())
+	labels := r.getLabelValues(p)
+	r.httpRequestDurHistogram.WithLabelValues(labels...).Observe(duration.Seconds())
 }
 
 func (r recorder) ObserveHTTPResponseSize(_ context.Context, p metrics.HTTPReqProperties, sizeBytes int64) {
-	r.httpResponseSizeHistogram.WithLabelValues(p.Service, p.ID, p.Method, p.Code).Observe(float64(sizeBytes))
+	labels := r.getLabelValues(p)
+	r.httpResponseSizeHistogram.WithLabelValues(labels...).Observe(float64(sizeBytes))
 }
 
 func (r recorder) AddInflightRequests(_ context.Context, p metrics.HTTPProperties, quantity int) {

middleware/echo/echo.go

@@ -35,3 +35,5 @@ func (r *reporter) URLPath() string { return r.c.Request().URL.Path }
 func (r *reporter) StatusCode() int { return r.c.Response().Status }
 
 func (r *reporter) BytesWritten() int64 { return r.c.Response().Size }
+
+func (r *reporter) CustomHeaders() map[string]string { return make(map[string]string) }

middleware/fasthttp/fasthttp.go

@@ -40,3 +40,5 @@ func (r reporter) StatusCode() int {
 func (r reporter) BytesWritten() int64 {
 	return int64(len(r.c.Response.Body()))
 }
+
+func (r reporter) CustomHeaders() map[string]string { return make(map[string]string) }

middleware/gin/gin.go

@@ -32,3 +32,5 @@ func (r *reporter) URLPath() string { return r.c.Request.URL.Path }
 func (r *reporter) StatusCode() int { return r.c.Writer.Status() }
 
 func (r *reporter) BytesWritten() int64 { return int64(r.c.Writer.Size()) }
+
+func (r *reporter) CustomHeaders() map[string]string { return make(map[string]string) }

middleware/gorestful/gorestful.go

@@ -33,3 +33,5 @@ func (r *reporter) URLPath() string { return r.req.Request.URL.Path }
 func (r *reporter) StatusCode() int { return r.resp.StatusCode() }
 
 func (r *reporter) BytesWritten() int64 { return int64(r.resp.ContentLength()) }
+
+func (r *reporter) CustomHeaders() map[string]string { return make(map[string]string) }

middleware/middleware.go

@@ -99,11 +99,16 @@ func (m Middleware) Measure(handlerID string, reporter Reporter, next func()) {
 			code = strconv.Itoa(reporter.StatusCode())
 		}
 
+		labels := make(map[string]string)
+		for k, v := range reporter.CustomHeaders() {
+			labels[k] = v
+		}
 		props := metrics.HTTPReqProperties{
 			Service: m.cfg.Service,
 			ID:      hid,
 			Method:  reporter.Method(),
 			Code:    code,
+			Labels:  labels,
 		}
 		m.cfg.Recorder.ObserveHTTPRequestDuration(ctx, props, duration)
 
@@ -124,5 +129,6 @@ type Reporter interface {
 	Context() context.Context
 	URLPath() string
 	StatusCode() int
+	CustomHeaders() map[string]string
 	BytesWritten() int64
 }

middleware/std/std.go

@@ -52,6 +52,8 @@ func (s *stdReporter) StatusCode() int { return s.w.statusCode }
 
 func (s *stdReporter) BytesWritten() int64 { return int64(s.w.bytesWritten) }
 
+func (s *stdReporter) CustomHeaders() map[string]string { return make(map[string]string) }
+
 // responseWriterInterceptor is a simple wrapper to intercept set data on a
 // ResponseWriter.
 type responseWriterInterceptor struct {