Rewrite client middleware

Author: stuartnelson3

prometheus/promhttp/client.go

@@ -34,34 +34,29 @@ type httpClient interface {
 	PostForm(string, url.Values) (*http.Response, error)
 }
 
-type Middleware func(req *http.Request) (*http.Response, error)
+type ClientMiddleware func(req *http.Request) (*http.Response, error)
 
-type Client struct {
-	c          httpClient
-	middleware Middleware
+func (c ClientMiddleware) Do(r *http.Request) (*http.Response, error) {
+	return c(r)
 }
 
-func (c *Client) Do(r *http.Request) (*http.Response, error) {
-	return c.middleware(r)
-}
-
-func (c *Client) Get(url string) (resp *http.Response, err error) {
+func (c ClientMiddleware) Get(url string) (resp *http.Response, err error) {
 	req, err := http.NewRequest("GET", url, nil)
 	if err != nil {
 		return nil, err
 	}
 	return c.Do(req)
 }
 
-func (c *Client) Head(url string) (*http.Response, error) {
+func (c ClientMiddleware) Head(url string) (*http.Response, error) {
 	req, err := http.NewRequest("HEAD", url, nil)
 	if err != nil {
 		return nil, err
 	}
-	return c.middleware(req)
+	return c.Do(req)
 }
 
-func (c *Client) Post(url string, contentType string, body io.Reader) (*http.Response, error) {
+func (c ClientMiddleware) Post(url string, contentType string, body io.Reader) (*http.Response, error) {
 	req, err := http.NewRequest("POST", url, body)
 	if err != nil {
 		return nil, err
@@ -70,6 +65,6 @@ func (c *Client) Post(url string, contentType string, body io.Reader) (*http.Res
 	return c.Do(req)
 }
 
-func (c *Client) PostForm(url string, data url.Values) (*http.Response, error) {
+func (c ClientMiddleware) PostForm(url string, data url.Values) (*http.Response, error) {
 	return c.Post(url, "application/x-www-form-urlencoded", strings.NewReader(data.Encode()))
 }

prometheus/promhttp/client_test.go

@@ -20,89 +20,47 @@
 package promhttp
 
 import (
-	"log"
 	"net/http"
+	"net/http/httputil"
 	"testing"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 )
 
-func TestMiddlewareWrapsFirstToLast(t *testing.T) {
-	order := []int{}
-	first := func(r *http.Request, next Middleware) (*http.Response, error) {
-		order = append(order, 0)
-
-		resp, err := next(r)
-
-		order = append(order, 3)
-		return resp, err
-	}
-
-	second := func(req *http.Request, next Middleware) (*http.Response, error) {
-		order = append(order, 1)
-
-		return next(req)
-	}
-
-	third := func(req *http.Request, next Middleware) (*http.Response, error) {
-		order = append(order, 2)
-		return next(req)
-	}
-
-	promclient, err := NewClient(SetMiddleware(first, second, third))
-	if err != nil {
-		t.Fatalf("%v", err)
-	}
-
-	resp, err := promclient.Get("http://google.com")
-	if err != nil {
-		t.Fatalf("%v", err)
-	}
-	defer resp.Body.Close()
-
-	for want, got := range order {
-		if want != got {
-			t.Fatalf("wanted %d, got %d", want, got)
-		}
-	}
-}
-
-func TestMiddlewareAPI(t *testing.T) {
+func TestClientMiddlewareAPI(t *testing.T) {
 	client := *http.DefaultClient
 	client.Timeout = 300 * time.Millisecond
 
 	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{Name: "inFlight"})
-	inFlight := func(r *http.Request, next Middleware) (*http.Response, error) {
-		inFlightGauge.Inc()
-
-		resp, err := next(r)
-
-		inFlightGauge.Dec()
-
-		return resp, err
-	}
 
-	counter := prometheus.NewCounter(prometheus.CounterOpts{Name: "test_counter"})
-	addFortyTwo := func(req *http.Request, next Middleware) (*http.Response, error) {
-		counter.Add(42)
+	counter := prometheus.NewCounterVec(
+		prometheus.CounterOpts{Name: "test_counter"},
+		[]string{"code", "method"},
+	)
+
+	histVec := prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Name:    "latency",
+			Buckets: prometheus.DefBuckets,
+		},
+		[]string{"event"},
+	)
+
+	promclient := InFlightC(inFlightGauge,
+		CounterC(counter,
+			ClientTrace(histVec, &client),
+		),
+	)
 
-		return next(req)
-	}
-
-	logging := func(req *http.Request, next Middleware) (*http.Response, error) {
-		log.Println("log something interesting")
-		return next(req)
-	}
-
-	promclient, err := NewClient(SetClient(client), SetMiddleware(inFlight, addFortyTwo, logging))
+	resp, err := promclient.Get("http://google.com")
 	if err != nil {
 		t.Fatalf("%v", err)
 	}
+	defer resp.Body.Close()
 
-	resp, err := promclient.Get("http://google.com")
+	out, err := httputil.DumpResponse(resp, true)
 	if err != nil {
 		t.Fatalf("%v", err)
 	}
-	defer resp.Body.Close()
 }

prometheus/promhttp/middleware.go

@@ -33,60 +33,60 @@ import (
 func ClientTrace(obs prometheus.ObserverVec, c httpClient) httpClient {
 	// The supplied histogram NEEDS a label for the httptrace event.
 	// TODO: Using `event` for now, but any other name is acceptable.
-	wc := &Client{
-		c: c,
-	}
 
 	// TODO: Check for `event` label on histogram.
 
-	wc.middleware = func(r *http.Request) (*http.Response, error) {
+	return ClientMiddleware(func(r *http.Request) (*http.Response, error) {
 		var (
-			host  = r.URL.Host
 			start = time.Now()
 		)
 
 		trace := &httptrace.ClientTrace{
 			DNSStart: func(_ httptrace.DNSStartInfo) {
-				obs.WithLabelValues(host, "DNSStart").Observe(time.Since(start).Seconds())
+				obs.WithLabelValues("DNSStart").Observe(time.Since(start).Seconds())
 			},
 			DNSDone: func(_ httptrace.DNSDoneInfo) {
-				obs.WithLabelValues(host, "DNSDone").Observe(time.Since(start).Seconds())
+				obs.WithLabelValues("DNSDone").Observe(time.Since(start).Seconds())
 			},
 			ConnectStart: func(_, _ string) {
-				obs.WithLabelValues(host, "ConnectStart").Observe(time.Since(start).Seconds())
+				obs.WithLabelValues("ConnectStart").Observe(time.Since(start).Seconds())
 			},
 			ConnectDone: func(net, addr string, err error) {
 				if err != nil {
 					return
 				}
-				obs.WithLabelValues(host, "ConnectDone").Observe(time.Since(start).Seconds())
+				obs.WithLabelValues("ConnectDone").Observe(time.Since(start).Seconds())
 			},
 			GotConn: func(_ httptrace.GotConnInfo) {
-				obs.WithLabelValues(host, "GotConn").Observe(time.Since(start).Seconds())
+				obs.WithLabelValues("GotConn").Observe(time.Since(start).Seconds())
 			},
 			GotFirstResponseByte: func() {
-				obs.WithLabelValues(host, "GotFirstResponseByte").Observe(time.Since(start).Seconds())
+				obs.WithLabelValues("GotFirstResponseByte").Observe(time.Since(start).Seconds())
 			},
 		}
 		r = r.WithContext(httptrace.WithClientTrace(context.Background(), trace))
 
-		return wc.Do(r)
-	}
-
-	return wc
+		return c.Do(r)
+	})
 }
 
 // InFlight is middleware that instruments number of open requests partitioned
 // by http client and request host.
-var InFlight = func(gauge prometheus.Gauge, c httpClient) httpClient {
-	wc := &Client{
-		c: c,
-	}
-	wc.middleware = func(r *http.Request) (*http.Response, error) {
+func InFlightC(gauge prometheus.Gauge, c httpClient) httpClient {
+	return ClientMiddleware(func(r *http.Request) (*http.Response, error) {
 		gauge.Inc()
-		resp, err := wc.Do(r)
+		resp, err := c.Do(r)
 		gauge.Dec()
 		return resp, err
-	}
-	return wc
+	})
+}
+
+func CounterC(counter *prometheus.CounterVec, c httpClient) httpClient {
+	code, method := checkLabels(counter)
+
+	return ClientMiddleware(func(r *http.Request) (*http.Response, error) {
+		resp, err := c.Do(r)
+		counter.With(labels(code, method, r.Method, resp.StatusCode)).Inc()
+		return resp, err
+	})
 }