Add possibility to dynamically get label values for http instrumentation (#1066)

Signed-off-by: Quentin Devos <4972091+Okhoshi@users.noreply.github.com>

Signed-off-by: Quentin Devos <4972091+Okhoshi@users.noreply.github.com>

examples/simple/main.go

@@ -16,10 +16,11 @@ package main
 
 import (
 	"flag"
-	"github.com/prometheus/client_golang/prometheus/collectors"
 	"log"
 	"net/http"
 
+	"github.com/prometheus/client_golang/prometheus/collectors"
+
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )

prometheus/promhttp/instrument_client.go

@@ -68,16 +68,17 @@ func InstrumentRoundTripperCounter(counter *prometheus.CounterVec, next http.Rou
 		o.apply(rtOpts)
 	}
 
-	code, method := checkLabels(counter)
+	// Curry the counter with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(counter.MustCurryWith(rtOpts.emptyDynamicLabels()))
 
 	return func(r *http.Request) (*http.Response, error) {
 		resp, err := next.RoundTrip(r)
 		if err == nil {
-			addWithExemplar(
+			l := labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)
-				counter.With(labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)),
+			for label, resolve := range rtOpts.extraLabelsFromCtx {
-				1,
+				l[label] = resolve(resp.Request.Context())
-				rtOpts.getExemplarFn(r.Context()),
+			}
-			)
+			addWithExemplar(counter.With(l), 1, rtOpts.getExemplarFn(r.Context()))
 		}
 		return resp, err
 	}
@@ -110,17 +111,18 @@ func InstrumentRoundTripperDuration(obs prometheus.ObserverVec, next http.RoundT
 		o.apply(rtOpts)
 	}
 
-	code, method := checkLabels(obs)
+	// Curry the observer with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(obs.MustCurryWith(rtOpts.emptyDynamicLabels()))
 
 	return func(r *http.Request) (*http.Response, error) {
 		start := time.Now()
 		resp, err := next.RoundTrip(r)
 		if err == nil {
-			observeWithExemplar(
+			l := labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)
-				obs.With(labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)),
+			for label, resolve := range rtOpts.extraLabelsFromCtx {
-				time.Since(start).Seconds(),
+				l[label] = resolve(resp.Request.Context())
-				rtOpts.getExemplarFn(r.Context()),
+			}
-			)
+			observeWithExemplar(obs.With(l), time.Since(start).Seconds(), rtOpts.getExemplarFn(r.Context()))
 		}
 		return resp, err
 	}

prometheus/promhttp/instrument_server.go

@@ -87,7 +87,8 @@ func InstrumentHandlerDuration(obs prometheus.ObserverVec, next http.Handler, op
 		o.apply(hOpts)
 	}
 
-	code, method := checkLabels(obs)
+	// Curry the observer with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(obs.MustCurryWith(hOpts.emptyDynamicLabels()))
 
 	if code {
 		return func(w http.ResponseWriter, r *http.Request) {
@@ -95,23 +96,22 @@ func InstrumentHandlerDuration(obs prometheus.ObserverVec, next http.Handler, op
 			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 
-			observeWithExemplar(
+			l := labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)
-				obs.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+			for label, resolve := range hOpts.extraLabelsFromCtx {
-				time.Since(now).Seconds(),
+				l[label] = resolve(r.Context())
-				hOpts.getExemplarFn(r.Context()),
+			}
-			)
+			observeWithExemplar(obs.With(l), time.Since(now).Seconds(), hOpts.getExemplarFn(r.Context()))
 		}
 	}
 
 	return func(w http.ResponseWriter, r *http.Request) {
 		now := time.Now()
 		next.ServeHTTP(w, r)
-
+		l := labels(code, method, r.Method, 0, hOpts.extraMethods...)
-		observeWithExemplar(
+		for label, resolve := range hOpts.extraLabelsFromCtx {
-			obs.With(labels(code, method, r.Method, 0, hOpts.extraMethods...)),
+			l[label] = resolve(r.Context())
-			time.Since(now).Seconds(),
+		}
-			hOpts.getExemplarFn(r.Context()),
+		observeWithExemplar(obs.With(l), time.Since(now).Seconds(), hOpts.getExemplarFn(r.Context()))
-		)
 	}
 }
 
@@ -138,28 +138,30 @@ func InstrumentHandlerCounter(counter *prometheus.CounterVec, next http.Handler,
 		o.apply(hOpts)
 	}
 
-	code, method := checkLabels(counter)
+	// Curry the counter with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(counter.MustCurryWith(hOpts.emptyDynamicLabels()))
 
 	if code {
 		return func(w http.ResponseWriter, r *http.Request) {
 			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 
-			addWithExemplar(
+			l := labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)
-				counter.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+			for label, resolve := range hOpts.extraLabelsFromCtx {
-				1,
+				l[label] = resolve(r.Context())
-				hOpts.getExemplarFn(r.Context()),
+			}
-			)
+			addWithExemplar(counter.With(l), 1, hOpts.getExemplarFn(r.Context()))
 		}
 	}
 
 	return func(w http.ResponseWriter, r *http.Request) {
 		next.ServeHTTP(w, r)
-		addWithExemplar(
+
-			counter.With(labels(code, method, r.Method, 0, hOpts.extraMethods...)),
+		l := labels(code, method, r.Method, 0, hOpts.extraMethods...)
-			1,
+		for label, resolve := range hOpts.extraLabelsFromCtx {
-			hOpts.getExemplarFn(r.Context()),
+			l[label] = resolve(r.Context())
-		)
+		}
+		addWithExemplar(counter.With(l), 1, hOpts.getExemplarFn(r.Context()))
 	}
 }
 
@@ -191,16 +193,17 @@ func InstrumentHandlerTimeToWriteHeader(obs prometheus.ObserverVec, next http.Ha
 		o.apply(hOpts)
 	}
 
-	code, method := checkLabels(obs)
+	// Curry the observer with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(obs.MustCurryWith(hOpts.emptyDynamicLabels()))
 
 	return func(w http.ResponseWriter, r *http.Request) {
 		now := time.Now()
 		d := newDelegator(w, func(status int) {
-			observeWithExemplar(
+			l := labels(code, method, r.Method, status, hOpts.extraMethods...)
-				obs.With(labels(code, method, r.Method, status, hOpts.extraMethods...)),
+			for label, resolve := range hOpts.extraLabelsFromCtx {
-				time.Since(now).Seconds(),
+				l[label] = resolve(r.Context())
-				hOpts.getExemplarFn(r.Context()),
+			}
-			)
+			observeWithExemplar(obs.With(l), time.Since(now).Seconds(), hOpts.getExemplarFn(r.Context()))
 		})
 		next.ServeHTTP(d, r)
 	}
@@ -231,28 +234,32 @@ func InstrumentHandlerRequestSize(obs prometheus.ObserverVec, next http.Handler,
 		o.apply(hOpts)
 	}
 
-	code, method := checkLabels(obs)
+	// Curry the observer with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(obs.MustCurryWith(hOpts.emptyDynamicLabels()))
+
 	if code {
 		return func(w http.ResponseWriter, r *http.Request) {
 			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 			size := computeApproximateRequestSize(r)
-			observeWithExemplar(
+
-				obs.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+			l := labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)
-				float64(size),
+			for label, resolve := range hOpts.extraLabelsFromCtx {
-				hOpts.getExemplarFn(r.Context()),
+				l[label] = resolve(r.Context())
-			)
+			}
+			observeWithExemplar(obs.With(l), float64(size), hOpts.getExemplarFn(r.Context()))
 		}
 	}
 
 	return func(w http.ResponseWriter, r *http.Request) {
 		next.ServeHTTP(w, r)
 		size := computeApproximateRequestSize(r)
-		observeWithExemplar(
+
-			obs.With(labels(code, method, r.Method, 0, hOpts.extraMethods...)),
+		l := labels(code, method, r.Method, 0, hOpts.extraMethods...)
-			float64(size),
+		for label, resolve := range hOpts.extraLabelsFromCtx {
-			hOpts.getExemplarFn(r.Context()),
+			l[label] = resolve(r.Context())
-		)
+		}
+		observeWithExemplar(obs.With(l), float64(size), hOpts.getExemplarFn(r.Context()))
 	}
 }
 
@@ -281,16 +288,18 @@ func InstrumentHandlerResponseSize(obs prometheus.ObserverVec, next http.Handler
 		o.apply(hOpts)
 	}
 
-	code, method := checkLabels(obs)
+	// Curry the observer with dynamic labels before checking the remaining labels.
+	code, method := checkLabels(obs.MustCurryWith(hOpts.emptyDynamicLabels()))
 
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		d := newDelegator(w, nil)
 		next.ServeHTTP(d, r)
-		observeWithExemplar(
+
-			obs.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+		l := labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)
-			float64(d.Written()),
+		for label, resolve := range hOpts.extraLabelsFromCtx {
-			hOpts.getExemplarFn(r.Context()),
+			l[label] = resolve(r.Context())
-		)
+		}
+		observeWithExemplar(obs.With(l), float64(d.Written()), hOpts.getExemplarFn(r.Context()))
 	})
 }
 

prometheus/promhttp/instrument_server_test.go

@@ -30,6 +30,7 @@ func TestLabelCheck(t *testing.T) {
 		varLabels     []string
 		constLabels   []string
 		curriedLabels []string
+		dynamicLabels []string
 		ok            bool
 	}{
 		"empty": {
@@ -60,12 +61,14 @@ func TestLabelCheck(t *testing.T) {
 			varLabels:     []string{"code", "method"},
 			constLabels:   []string{"foo", "bar"},
 			curriedLabels: []string{"dings", "bums"},
+			dynamicLabels: []string{"dyn", "amics"},
 			ok:            true,
 		},
 		"all labels used with an invalid const label name": {
 			varLabels:     []string{"code", "method"},
 			constLabels:   []string{"in-valid", "bar"},
 			curriedLabels: []string{"dings", "bums"},
+			dynamicLabels: []string{"dyn", "amics"},
 			ok:            false,
 		},
 		"unsupported var label": {
@@ -98,6 +101,18 @@ func TestLabelCheck(t *testing.T) {
 			curriedLabels: []string{"method"},
 			ok:            true,
 		},
+		"supported label as const and dynamic": {
+			varLabels:     []string{},
+			constLabels:   []string{"code"},
+			dynamicLabels: []string{"method"},
+			ok:            true,
+		},
+		"supported label as curried and dynamic": {
+			varLabels:     []string{},
+			curriedLabels: []string{"code"},
+			dynamicLabels: []string{"method"},
+			ok:            true,
+		},
 		"supported label as const and curry with unsupported as var": {
 			varLabels:     []string{"foo"},
 			constLabels:   []string{"code"},
@@ -116,6 +131,7 @@ func TestLabelCheck(t *testing.T) {
 			varLabels:     []string{"code", "method"},
 			constLabels:   []string{"foo", "bar"},
 			curriedLabels: []string{"dings", "bums"},
+			dynamicLabels: []string{"dyn", "amics"},
 			ok:            false,
 		},
 	}
@@ -130,26 +146,39 @@ func TestLabelCheck(t *testing.T) {
 			for _, l := range sc.constLabels {
 				constLabels[l] = "dummy"
 			}
-			c := prometheus.NewCounterVec(
+			labelNames := append(append(sc.varLabels, sc.curriedLabels...), sc.dynamicLabels...)
-				prometheus.CounterOpts{
+			c := prometheus.V2.NewCounterVec(
+				prometheus.CounterVecOpts{
+					CounterOpts: prometheus.CounterOpts{
 						Name:        metricName,
 						Help:        "c help",
 						ConstLabels: constLabels,
 					},
-				append(sc.varLabels, sc.curriedLabels...),
+					VariableLabels: prometheus.UnconstrainedLabels(labelNames),
+				},
 			)
-			o := prometheus.ObserverVec(prometheus.NewHistogramVec(
+			o := prometheus.ObserverVec(prometheus.V2.NewHistogramVec(
-				prometheus.HistogramOpts{
+				prometheus.HistogramVecOpts{
+					HistogramOpts: prometheus.HistogramOpts{
 						Name:        metricName,
 						Help:        "c help",
 						ConstLabels: constLabels,
 					},
-				append(sc.varLabels, sc.curriedLabels...),
+					VariableLabels: prometheus.UnconstrainedLabels(labelNames),
+				},
 			))
 			for _, l := range sc.curriedLabels {
 				c = c.MustCurryWith(prometheus.Labels{l: "dummy"})
 				o = o.MustCurryWith(prometheus.Labels{l: "dummy"})
 			}
+			opts := []Option{}
+			for _, l := range sc.dynamicLabels {
+				opts = append(opts, WithLabelFromCtx(l,
+					func(_ context.Context) string {
+						return "foo"
+					},
+				))
+			}
 
 			func() {
 				defer func() {
@@ -161,7 +190,7 @@ func TestLabelCheck(t *testing.T) {
 						t.Error("expected panic")
 					}
 				}()
-				InstrumentHandlerCounter(c, nil)
+				InstrumentHandlerCounter(c, nil, opts...)
 			}()
 			func() {
 				defer func() {
@@ -173,7 +202,7 @@ func TestLabelCheck(t *testing.T) {
 						t.Error("expected panic")
 					}
 				}()
-				InstrumentHandlerDuration(o, nil)
+				InstrumentHandlerDuration(o, nil, opts...)
 			}()
 			if sc.ok {
 				// Test if wantCode and wantMethod were detected correctly.
@@ -186,6 +215,11 @@ func TestLabelCheck(t *testing.T) {
 						wantMethod = true
 					}
 				}
+				// Curry the dynamic labels since this is done normally behind the scenes for the check
+				for _, l := range sc.dynamicLabels {
+					c = c.MustCurryWith(prometheus.Labels{l: "dummy"})
+					o = o.MustCurryWith(prometheus.Labels{l: "dummy"})
+				}
 				gotCode, gotMethod := checkLabels(c)
 				if gotCode != wantCode {
 					t.Errorf("wanted code=%t for counter, got code=%t", wantCode, gotCode)

prometheus/promhttp/option.go

@@ -24,14 +24,32 @@ type Option interface {
 	apply(*options)
 }
 
+// LabelValueFromCtx are used to compute the label value from request context.
+// Context can be filled with values from request through middleware.
+type LabelValueFromCtx func(ctx context.Context) string
+
 // options store options for both a handler or round tripper.
 type options struct {
 	extraMethods       []string
 	getExemplarFn      func(requestCtx context.Context) prometheus.Labels
+	extraLabelsFromCtx map[string]LabelValueFromCtx
 }
 
 func defaultOptions() *options {
-	return &options{getExemplarFn: func(ctx context.Context) prometheus.Labels { return nil }}
+	return &options{
+		getExemplarFn:      func(ctx context.Context) prometheus.Labels { return nil },
+		extraLabelsFromCtx: map[string]LabelValueFromCtx{},
+	}
+}
+
+func (o *options) emptyDynamicLabels() prometheus.Labels {
+	labels := prometheus.Labels{}
+
+	for label := range o.extraLabelsFromCtx {
+		labels[label] = ""
+	}
+
+	return labels
 }
 
 type optionApplyFunc func(*options)
@@ -56,3 +74,11 @@ func WithExemplarFromContext(getExemplarFn func(requestCtx context.Context) prom
 		o.getExemplarFn = getExemplarFn
 	})
 }
+
+// WithLabelFromCtx registers a label for dynamic resolution with access to context.
+// See the example for ExampleInstrumentHandlerWithLabelResolver for example usage
+func WithLabelFromCtx(name string, valueFn LabelValueFromCtx) Option {
+	return optionApplyFunc(func(o *options) {
+		o.extraLabelsFromCtx[name] = valueFn
+	})
+}

prometheus/promhttp/option_test.go

@@ -14,12 +14,19 @@
 package promhttp
 
 import (
+	"context"
 	"log"
 	"net/http"
 
 	"github.com/prometheus/client_golang/prometheus"
 )
 
+type key int
+
+const (
+	CtxResolverKey key = iota
+)
+
 func ExampleInstrumentHandlerWithExtraMethods() {
 	counter := prometheus.NewCounterVec(
 		prometheus.CounterOpts{
@@ -62,3 +69,60 @@ func ExampleInstrumentHandlerWithExtraMethods() {
 		log.Fatal(err)
 	}
 }
+
+func ExampleInstrumentHandlerWithLabelResolver() {
+	counter := prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "api_requests_total",
+			Help: "A counter for requests to the wrapped handler.",
+		},
+		[]string{"code", "method", "myheader"},
+	)
+
+	// duration is partitioned by the HTTP method, handler and request header
+	// value. It uses custom buckets based on the expected request duration.
+	// Beware to not have too high cardinality on the values of header. You
+	// always should sanitize external inputs.
+	duration := prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Name:    "request_duration_seconds",
+			Help:    "A histogram of latencies for requests.",
+			Buckets: []float64{.25, .5, 1, 2.5, 5, 10},
+		},
+		[]string{"handler", "method", "myheader"},
+	)
+
+	// Create the handlers that will be wrapped by the middleware.
+	pullHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("Pull"))
+	})
+
+	// Specify additional HTTP methods to be added to the label allow list.
+	opts := WithLabelFromCtx("myheader",
+		func(ctx context.Context) string {
+			return ctx.Value(CtxResolverKey).(string)
+		},
+	)
+
+	// Instrument the handlers with all the metrics, injecting the "handler"
+	// label by currying.
+	pullChain := InstrumentHandlerDuration(duration.MustCurryWith(prometheus.Labels{"handler": "pull"}),
+		InstrumentHandlerCounter(counter, pullHandler, opts),
+		opts,
+	)
+
+	middleware := func(next http.HandlerFunc) http.HandlerFunc {
+		return func(w http.ResponseWriter, r *http.Request) {
+			ctx := context.WithValue(r.Context(), CtxResolverKey, r.Header.Get("x-my-header"))
+
+			next(w, r.WithContext(ctx))
+		}
+	}
+
+	http.Handle("/metrics", Handler())
+	http.Handle("/pull", middleware(pullChain))
+
+	if err := http.ListenAndServe(":3000", nil); err != nil {
+		log.Fatal(err)
+	}
+}