Added exemplar support to http middleware. (#1055)

* Added exemplar support to http middlewares.

Signed-off-by: Bartlomiej Plotka <bwplotka@gmail.com>

* Small fix.

Signed-off-by: Bartlomiej Plotka <bwplotka@gmail.com>

* Fixed test.

Signed-off-by: Bartlomiej Plotka <bwplotka@gmail.com>

* Added tests and options for RT.

Signed-off-by: bwplotka <bwplotka@gmail.com>

* goimports.

Signed-off-by: bwplotka <bwplotka@gmail.com>

prometheus/promhttp/instrument_client.go

@@ -38,11 +38,11 @@ func (rt RoundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
 //
 // See the example for ExampleInstrumentRoundTripperDuration for example usage.
 func InstrumentRoundTripperInFlight(gauge prometheus.Gauge, next http.RoundTripper) RoundTripperFunc {
-	return RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
+	return func(r *http.Request) (*http.Response, error) {
 		gauge.Inc()
 		defer gauge.Dec()
 		return next.RoundTrip(r)
-	})
+	}
 }
 
 // InstrumentRoundTripperCounter is a middleware that wraps the provided
@@ -59,22 +59,29 @@ func InstrumentRoundTripperInFlight(gauge prometheus.Gauge, next http.RoundTripp
 // If the wrapped RoundTripper panics or returns a non-nil error, the Counter
 // is not incremented.
 //
+// Use with WithExemplarFromContext to instrument the exemplars on the counter of requests.
+//
 // See the example for ExampleInstrumentRoundTripperDuration for example usage.
 func InstrumentRoundTripperCounter(counter *prometheus.CounterVec, next http.RoundTripper, opts ...Option) RoundTripperFunc {
-	rtOpts := &option{}
+	rtOpts := defaultOptions()
 	for _, o := range opts {
-		o(rtOpts)
+		o.apply(rtOpts)
 	}
 
 	code, method := checkLabels(counter)
 
-	return RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
+	return func(r *http.Request) (*http.Response, error) {
 		resp, err := next.RoundTrip(r)
 		if err == nil {
+			exemplarAdd(
+				counter.With(labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)),
+				1,
+				rtOpts.getExemplarFn(r.Context()),
+			)
 			counter.With(labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)).Inc()
 		}
 		return resp, err
-	})
+	}
 }
 
 // InstrumentRoundTripperDuration is a middleware that wraps the provided
@@ -94,24 +101,30 @@ func InstrumentRoundTripperCounter(counter *prometheus.CounterVec, next http.Rou
 // If the wrapped RoundTripper panics or returns a non-nil error, no values are
 // reported.
 //
+// Use with WithExemplarFromContext to instrument the exemplars on the duration histograms.
+//
 // Note that this method is only guaranteed to never observe negative durations
 // if used with Go1.9+.
 func InstrumentRoundTripperDuration(obs prometheus.ObserverVec, next http.RoundTripper, opts ...Option) RoundTripperFunc {
-	rtOpts := &option{}
+	rtOpts := defaultOptions()
 	for _, o := range opts {
-		o(rtOpts)
+		o.apply(rtOpts)
 	}
 
 	code, method := checkLabels(obs)
 
-	return RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
+	return func(r *http.Request) (*http.Response, error) {
 		start := time.Now()
 		resp, err := next.RoundTrip(r)
 		if err == nil {
-			obs.With(labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)).Observe(time.Since(start).Seconds())
+			exemplarObserve(
+				obs.With(labels(code, method, r.Method, resp.StatusCode, rtOpts.extraMethods...)),
+				time.Since(start).Seconds(),
+				rtOpts.getExemplarFn(r.Context()),
+			)
 		}
 		return resp, err
-	})
+	}
 }
 
 // InstrumentTrace is used to offer flexibility in instrumenting the available
@@ -149,7 +162,7 @@ type InstrumentTrace struct {
 //
 // See the example for ExampleInstrumentRoundTripperDuration for example usage.
 func InstrumentRoundTripperTrace(it *InstrumentTrace, next http.RoundTripper) RoundTripperFunc {
-	return RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
+	return func(r *http.Request) (*http.Response, error) {
 		start := time.Now()
 
 		trace := &httptrace.ClientTrace{
@@ -231,5 +244,5 @@ func InstrumentRoundTripperTrace(it *InstrumentTrace, next http.RoundTripper) Ro
 		r = r.WithContext(httptrace.WithClientTrace(r.Context(), trace))
 
 		return next.RoundTrip(r)
-	})
+	}
 }

prometheus/promhttp/instrument_client_test.go

@@ -18,14 +18,19 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
+	"sort"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
+
+	dto "github.com/prometheus/client_model/go"
+	"google.golang.org/protobuf/proto"
 )
 
-func makeInstrumentedClient() (*http.Client, *prometheus.Registry) {
+func makeInstrumentedClient(opts ...Option) (*http.Client, *prometheus.Registry) {
 	client := http.DefaultClient
 	client.Timeout = 1 * time.Second
 
@@ -91,13 +96,91 @@ func makeInstrumentedClient() (*http.Client, *prometheus.Registry) {
 	client.Transport = InstrumentRoundTripperInFlight(inFlightGauge,
 		InstrumentRoundTripperCounter(counter,
 			InstrumentRoundTripperTrace(trace,
-				InstrumentRoundTripperDuration(histVec, http.DefaultTransport),
+				InstrumentRoundTripperDuration(histVec, http.DefaultTransport, opts...),
-			),
 			),
+			opts...),
 	)
 	return client, reg
 }
 
+func labelsToLabelPair(l prometheus.Labels) []*dto.LabelPair {
+	ret := make([]*dto.LabelPair, 0, len(l))
+	for k, v := range l {
+		ret = append(ret, &dto.LabelPair{Name: proto.String(k), Value: proto.String(v)})
+	}
+	sort.Slice(ret, func(i, j int) bool {
+		return *ret[i].Name < *ret[j].Name
+	})
+	return ret
+}
+
+func assetMetricAndExemplars(
+	t *testing.T,
+	reg *prometheus.Registry,
+	expectedNumMetrics int,
+	expectedExemplar []*dto.LabelPair,
+) {
+	t.Helper()
+
+	mfs, err := reg.Gather()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want, got := expectedNumMetrics, len(mfs); want != got {
+		t.Fatalf("unexpected number of metric families gathered, want %d, got %d", want, got)
+	}
+
+	for _, mf := range mfs {
+		if len(mf.Metric) == 0 {
+			t.Errorf("metric family %s must not be empty", mf.GetName())
+		}
+		for _, m := range mf.GetMetric() {
+			if c := m.GetCounter(); c != nil {
+				if len(expectedExemplar) == 0 {
+					if c.Exemplar != nil {
+						t.Errorf("expected no exemplar on the counter %v%v, got %v", mf.GetName(), m.Label, c.Exemplar.String())
+					}
+					continue
+				}
+
+				if c.Exemplar == nil {
+					t.Errorf("expected exemplar %v on the counter %v%v, got none", expectedExemplar, mf.GetName(), m.Label)
+					continue
+				}
+				if got := c.Exemplar.Label; !reflect.DeepEqual(expectedExemplar, got) {
+					t.Errorf("expected exemplar %v on the counter %v%v, got %v", expectedExemplar, mf.GetName(), m.Label, got)
+				}
+				continue
+			}
+			if h := m.GetHistogram(); h != nil {
+				found := false
+				for _, b := range h.GetBucket() {
+					if len(expectedExemplar) == 0 {
+						if b.Exemplar != nil {
+							t.Errorf("expected no exemplar on histogram %v%v bkt %v, got %v", mf.GetName(), m.Label, b.GetUpperBound(), b.Exemplar.String())
+						}
+						continue
+					}
+
+					if b.Exemplar == nil {
+						continue
+					}
+					if got := b.Exemplar.Label; !reflect.DeepEqual(expectedExemplar, got) {
+						t.Errorf("expected exemplar %v on the histogram %v%v on bkt %v, got %v", expectedExemplar, mf.GetName(), m.Label, b.GetUpperBound(), got)
+						continue
+					}
+					found = true
+					break
+				}
+
+				if len(expectedExemplar) > 0 && !found {
+					t.Errorf("expected exemplar %v on at least one bucket of the histogram %v%v, got none", expectedExemplar, mf.GetName(), m.Label)
+				}
+			}
+		}
+	}
+}
+
 func TestClientMiddlewareAPI(t *testing.T) {
 	client, reg := makeInstrumentedClient()
 	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -111,21 +194,28 @@ func TestClientMiddlewareAPI(t *testing.T) {
 	}
 	defer resp.Body.Close()
 
-	mfs, err := reg.Gather()
+	assetMetricAndExemplars(t, reg, 3, nil)
+}
+
+func TestClientMiddlewareAPI_WithExemplars(t *testing.T) {
+	exemplar := prometheus.Labels{"traceID": "example situation observed by this metric"}
+
+	client, reg := makeInstrumentedClient(WithExemplarFromContext(func(_ context.Context) prometheus.Labels { return exemplar }))
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer backend.Close()
+
+	resp, err := client.Get(backend.URL)
 	if err != nil {
 		t.Fatal(err)
 	}
-	if want, got := 3, len(mfs); want != got {
+	defer resp.Body.Close()
-		t.Fatalf("unexpected number of metric families gathered, want %d, got %d", want, got)
+
-	}
+	assetMetricAndExemplars(t, reg, 3, labelsToLabelPair(exemplar))
-	for _, mf := range mfs {
-		if len(mf.Metric) == 0 {
-			t.Errorf("metric family %s must not be empty", mf.GetName())
-		}
-	}
 }
 
-func TestClientMiddlewareAPIWithRequestContext(t *testing.T) {
+func TestClientMiddlewareAPI_WithRequestContext(t *testing.T) {
 	client, reg := makeInstrumentedClient()
 	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)

prometheus/promhttp/instrument_server.go

@@ -28,6 +28,22 @@ import (
 // magicString is used for the hacky label test in checkLabels. Remove once fixed.
 const magicString = "zZgWfBxLqvG8kc8IMv3POi2Bb0tZI3vAnBx+gBaFi9FyPzB/CzKUer1yufDa"
 
+func exemplarObserve(obs prometheus.Observer, val float64, labels map[string]string) {
+	if labels == nil {
+		obs.Observe(val)
+		return
+	}
+	obs.(prometheus.ExemplarObserver).ObserveWithExemplar(val, labels)
+}
+
+func exemplarAdd(obs prometheus.Counter, val float64, labels map[string]string) {
+	if labels == nil {
+		obs.Add(val)
+		return
+	}
+	obs.(prometheus.ExemplarAdder).AddWithExemplar(val, labels)
+}
+
 // InstrumentHandlerInFlight is a middleware that wraps the provided
 // http.Handler. It sets the provided prometheus.Gauge to the number of
 // requests currently handled by the wrapped http.Handler.
@@ -62,28 +78,37 @@ func InstrumentHandlerInFlight(g prometheus.Gauge, next http.Handler) http.Handl
 // Note that this method is only guaranteed to never observe negative durations
 // if used with Go1.9+.
 func InstrumentHandlerDuration(obs prometheus.ObserverVec, next http.Handler, opts ...Option) http.HandlerFunc {
-	mwOpts := &option{}
+	hOpts := defaultOptions()
 	for _, o := range opts {
-		o(mwOpts)
+		o.apply(hOpts)
 	}
 
 	code, method := checkLabels(obs)
 
 	if code {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		return func(w http.ResponseWriter, r *http.Request) {
 			now := time.Now()
 			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 
-			obs.With(labels(code, method, r.Method, d.Status(), mwOpts.extraMethods...)).Observe(time.Since(now).Seconds())
+			exemplarObserve(
-		})
+				obs.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+				time.Since(now).Seconds(),
+				hOpts.getExemplarFn(r.Context()),
+			)
+		}
 	}
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		now := time.Now()
 		next.ServeHTTP(w, r)
-		obs.With(labels(code, method, r.Method, 0, mwOpts.extraMethods...)).Observe(time.Since(now).Seconds())
+
-	})
+		exemplarObserve(
+			obs.With(labels(code, method, r.Method, 0, hOpts.extraMethods...)),
+			time.Since(now).Seconds(),
+			hOpts.getExemplarFn(r.Context()),
+		)
+	}
 }
 
 // InstrumentHandlerCounter is a middleware that wraps the provided http.Handler
@@ -104,25 +129,34 @@ func InstrumentHandlerDuration(obs prometheus.ObserverVec, next http.Handler, op
 //
 // See the example for InstrumentHandlerDuration for example usage.
 func InstrumentHandlerCounter(counter *prometheus.CounterVec, next http.Handler, opts ...Option) http.HandlerFunc {
-	mwOpts := &option{}
+	hOpts := defaultOptions()
 	for _, o := range opts {
-		o(mwOpts)
+		o.apply(hOpts)
 	}
 
 	code, method := checkLabels(counter)
 
 	if code {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		return func(w http.ResponseWriter, r *http.Request) {
 			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
-			counter.With(labels(code, method, r.Method, d.Status(), mwOpts.extraMethods...)).Inc()
+
-		})
+			exemplarAdd(
+				counter.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+				1,
+				hOpts.getExemplarFn(r.Context()),
+			)
+		}
 	}
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		next.ServeHTTP(w, r)
-		counter.With(labels(code, method, r.Method, 0, mwOpts.extraMethods...)).Inc()
+		exemplarAdd(
-	})
+			counter.With(labels(code, method, r.Method, 0, hOpts.extraMethods...)),
+			1,
+			hOpts.getExemplarFn(r.Context()),
+		)
+	}
 }
 
 // InstrumentHandlerTimeToWriteHeader is a middleware that wraps the provided
@@ -148,20 +182,24 @@ func InstrumentHandlerCounter(counter *prometheus.CounterVec, next http.Handler,
 //
 // See the example for InstrumentHandlerDuration for example usage.
 func InstrumentHandlerTimeToWriteHeader(obs prometheus.ObserverVec, next http.Handler, opts ...Option) http.HandlerFunc {
-	mwOpts := &option{}
+	hOpts := defaultOptions()
 	for _, o := range opts {
-		o(mwOpts)
+		o.apply(hOpts)
 	}
 
 	code, method := checkLabels(obs)
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		now := time.Now()
 		d := newDelegator(w, func(status int) {
-			obs.With(labels(code, method, r.Method, status, mwOpts.extraMethods...)).Observe(time.Since(now).Seconds())
+			exemplarObserve(
+				obs.With(labels(code, method, r.Method, status, hOpts.extraMethods...)),
+				time.Since(now).Seconds(),
+				hOpts.getExemplarFn(r.Context()),
+			)
 		})
 		next.ServeHTTP(d, r)
-	})
+	}
 }
 
 // InstrumentHandlerRequestSize is a middleware that wraps the provided
@@ -184,27 +222,34 @@ func InstrumentHandlerTimeToWriteHeader(obs prometheus.ObserverVec, next http.Ha
 //
 // See the example for InstrumentHandlerDuration for example usage.
 func InstrumentHandlerRequestSize(obs prometheus.ObserverVec, next http.Handler, opts ...Option) http.HandlerFunc {
-	mwOpts := &option{}
+	hOpts := defaultOptions()
 	for _, o := range opts {
-		o(mwOpts)
+		o.apply(hOpts)
 	}
 
 	code, method := checkLabels(obs)
-
 	if code {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		return func(w http.ResponseWriter, r *http.Request) {
 			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 			size := computeApproximateRequestSize(r)
-			obs.With(labels(code, method, r.Method, d.Status(), mwOpts.extraMethods...)).Observe(float64(size))
+			exemplarObserve(
-		})
+				obs.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+				float64(size),
+				hOpts.getExemplarFn(r.Context()),
+			)
+		}
 	}
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		next.ServeHTTP(w, r)
 		size := computeApproximateRequestSize(r)
-		obs.With(labels(code, method, r.Method, 0, mwOpts.extraMethods...)).Observe(float64(size))
+		exemplarObserve(
-	})
+			obs.With(labels(code, method, r.Method, 0, hOpts.extraMethods...)),
+			float64(size),
+			hOpts.getExemplarFn(r.Context()),
+		)
+	}
 }
 
 // InstrumentHandlerResponseSize is a middleware that wraps the provided
@@ -227,9 +272,9 @@ func InstrumentHandlerRequestSize(obs prometheus.ObserverVec, next http.Handler,
 //
 // See the example for InstrumentHandlerDuration for example usage.
 func InstrumentHandlerResponseSize(obs prometheus.ObserverVec, next http.Handler, opts ...Option) http.Handler {
-	mwOpts := &option{}
+	hOpts := defaultOptions()
 	for _, o := range opts {
-		o(mwOpts)
+		o.apply(hOpts)
 	}
 
 	code, method := checkLabels(obs)
@@ -237,7 +282,11 @@ func InstrumentHandlerResponseSize(obs prometheus.ObserverVec, next http.Handler
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		d := newDelegator(w, nil)
 		next.ServeHTTP(d, r)
-		obs.With(labels(code, method, r.Method, d.Status(), mwOpts.extraMethods...)).Observe(float64(d.Written()))
+		exemplarObserve(
+			obs.With(labels(code, method, r.Method, d.Status(), hOpts.extraMethods...)),
+			float64(d.Written()),
+			hOpts.getExemplarFn(r.Context()),
+		)
 	})
 }
 

prometheus/promhttp/instrument_server_test.go

@@ -14,6 +14,7 @@
 package promhttp
 
 import (
+	"context"
 	"io"
 	"log"
 	"net/http"
@@ -321,7 +322,7 @@ func TestLabels(t *testing.T) {
 	}
 }
 
-func TestMiddlewareAPI(t *testing.T) {
+func makeInstrumentedHandler(handler http.HandlerFunc, opts ...Option) (http.Handler, *prometheus.Registry) {
 	reg := prometheus.NewRegistry()
 
 	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{
@@ -366,25 +367,43 @@ func TestMiddlewareAPI(t *testing.T) {
 		[]string{},
 	)
 
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Write([]byte("OK"))
-	})
-
 	reg.MustRegister(inFlightGauge, counter, histVec, responseSize, writeHeaderVec)
 
-	chain := InstrumentHandlerInFlight(inFlightGauge,
+	return InstrumentHandlerInFlight(inFlightGauge,
 		InstrumentHandlerCounter(counter,
 			InstrumentHandlerDuration(histVec,
 				InstrumentHandlerTimeToWriteHeader(writeHeaderVec,
-					InstrumentHandlerResponseSize(responseSize, handler),
+					InstrumentHandlerResponseSize(responseSize, handler, opts...),
-				),
+					opts...),
-			),
+				opts...),
-		),
+			opts...),
-	)
+	), reg
+}
+
+func TestMiddlewareAPI(t *testing.T) {
+	chain, reg := makeInstrumentedHandler(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("OK"))
+	})
 
 	r, _ := http.NewRequest("GET", "www.example.com", nil)
 	w := httptest.NewRecorder()
 	chain.ServeHTTP(w, r)
+
+	assetMetricAndExemplars(t, reg, 5, nil)
+}
+
+func TestMiddlewareAPI_WithExemplars(t *testing.T) {
+	exemplar := prometheus.Labels{"traceID": "example situation observed by this metric"}
+
+	chain, reg := makeInstrumentedHandler(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write([]byte("OK"))
+	}, WithExemplarFromContext(func(_ context.Context) prometheus.Labels { return exemplar }))
+
+	r, _ := http.NewRequest("GET", "www.example.com", nil)
+	w := httptest.NewRecorder()
+	chain.ServeHTTP(w, r)
+
+	assetMetricAndExemplars(t, reg, 5, labelsToLabelPair(exemplar))
 }
 
 func TestInstrumentTimeToFirstWrite(t *testing.T) {

prometheus/promhttp/option.go

@@ -13,19 +13,46 @@
 
 package promhttp
 
-// Option are used to configure a middleware or round tripper..
+import (
-type Option func(*option)
+	"context"
 
-type option struct {
+	"github.com/prometheus/client_golang/prometheus"
-	extraMethods []string
+)
+
+// Option are used to configure both handler (middleware) or round tripper.
+type Option interface {
+	apply(*options)
 }
 
+// options store options for both a handler or round tripper.
+type options struct {
+	extraMethods  []string
+	getExemplarFn func(requestCtx context.Context) prometheus.Labels
+}
+
+func defaultOptions() *options {
+	return &options{getExemplarFn: func(ctx context.Context) prometheus.Labels { return nil }}
+}
+
+type optionApplyFunc func(*options)
+
+func (o optionApplyFunc) apply(opt *options) { o(opt) }
+
 // WithExtraMethods adds additional HTTP methods to the list of allowed methods.
 // See https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods for the default list.
 //
 // See the example for ExampleInstrumentHandlerWithExtraMethods for example usage.
 func WithExtraMethods(methods ...string) Option {
-	return func(o *option) {
+	return optionApplyFunc(func(o *options) {
 		o.extraMethods = methods
+	})
 }
+
+// WithExemplarFromContext adds allows to put a hook to all counter and histogram metrics.
+// If the hook function returns non-nil labels, exemplars will be added for that request, otherwise metric
+// will get instrumented without exemplar.
+func WithExemplarFromContext(getExemplarFn func(requestCtx context.Context) prometheus.Labels) Option {
+	return optionApplyFunc(func(o *options) {
+		o.getExemplarFn = getExemplarFn
+	})
 }