Added tests and options for RT.

Signed-off-by: bwplotka <[email protected]>
prometheus · Aug 2, 2022 · 07608d4 · 07608d4
1 parent 4e2a055
commit 07608d4
Show file tree

Hide file tree

Showing 5 changed files with 216 additions and 74 deletions.
diff --git a/prometheus/promhttp/instrument_client.go b/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)
-	})
+	}
 }
diff --git a/prometheus/promhttp/instrument_client_test.go b/prometheus/promhttp/instrument_client_test.go
@@ -18,14 +18,18 @@ 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 +95,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 +193,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 {
-		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())
-		}
-	}
+	defer resp.Body.Close()
+
+	assetMetricAndExemplars(t, reg, 3, labelsToLabelPair(exemplar))
 }
 
-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)