Add time to write header handler middleware (#304)

Author: stuartnelson3

prometheus/promhttp/delegator_1_7.go

@@ -20,8 +20,11 @@ import (
 	"net/http"
 )
 
-func newDelegator(w http.ResponseWriter) delegator {
-	d := &responseWriterDelegator{ResponseWriter: w}
+func newDelegator(w http.ResponseWriter, observeWriteHeaderFunc func(int)) delegator {
+	d := &responseWriterDelegator{
+		ResponseWriter:     w,
+		observeWriteHeader: observeWriteHeaderFunc,
+	}
 
 	_, cn := w.(http.CloseNotifier)
 	_, fl := w.(http.Flusher)

prometheus/promhttp/delegator_1_8.go

@@ -22,8 +22,11 @@ import (
 
 // newDelegator handles the four different methods of upgrading a
 // http.ResponseWriter to delegator.
-func newDelegator(w http.ResponseWriter) delegator {
-	d := &responseWriterDelegator{ResponseWriter: w}
+func newDelegator(w http.ResponseWriter, observeWriteHeaderFunc func(int)) delegator {
+	d := &responseWriterDelegator{
+		ResponseWriter:     w,
+		observeWriteHeader: observeWriteHeaderFunc,
+	}
 
 	_, cn := w.(http.CloseNotifier)
 	_, fl := w.(http.Flusher)

prometheus/promhttp/instrument_server.go

@@ -63,7 +63,7 @@ func InstrumentHandlerDuration(obs prometheus.ObserverVec, next http.Handler) ht
 	if code {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			now := time.Now()
-			d := newDelegator(w)
+			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 
 			obs.With(labels(code, method, r.Method, d.Status())).Observe(time.Since(now).Seconds())
@@ -96,7 +96,7 @@ func InstrumentHandlerCounter(counter *prometheus.CounterVec, next http.Handler)
 
 	if code {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			d := newDelegator(w)
+			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 			counter.With(labels(code, method, r.Method, d.Status())).Inc()
 		})
@@ -108,6 +108,34 @@ func InstrumentHandlerCounter(counter *prometheus.CounterVec, next http.Handler)
 	})
 }
 
+// InstrumentHandlerTimeToWriteHeader is a middleware that wraps the provided
+// http.Handler to observe with the provided ObserverVec the request duration
+// until the response headers are written. The ObserverVec must have zero, one,
+// or two labels. The only allowed label names are "code" and "method". The
+// function panics if any other instance labels are provided. The Observe
+// method of the Observer in the ObserverVec is called with the request
+// duration in seconds. Partitioning happens by HTTP status code and/or HTTP
+// method if the respective instance label names are present in the
+// ObserverVec. For unpartitioned observations, use an ObserverVec with zero
+// labels. Note that partitioning of Histograms is expensive and should be used
+// judiciously.
+//
+// If the wrapped Handler panics before calling WriteHeader, no value is
+// reported.
+//
+// See the example for InstrumentHandlerDuration for example usage.
+func InstrumentHandlerTimeToWriteHeader(obs prometheus.ObserverVec, next http.Handler) http.HandlerFunc {
+	code, method := checkLabels(obs)
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		now := time.Now()
+		d := newDelegator(w, func(status int) {
+			obs.With(labels(code, method, r.Method, status)).Observe(time.Since(now).Seconds())
+		})
+		next.ServeHTTP(d, r)
+	})
+}
+
 // InstrumentHandlerRequestSize is a middleware that wraps the provided
 // http.Handler to observe the request size with the provided ObserverVec.
 // The ObserverVec must have zero, one, or two labels. The only allowed label
@@ -129,7 +157,7 @@ func InstrumentHandlerRequestSize(obs prometheus.ObserverVec, next http.Handler)
 
 	if code {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			d := newDelegator(w)
+			d := newDelegator(w, nil)
 			next.ServeHTTP(d, r)
 			size := computeApproximateRequestSize(r)
 			obs.With(labels(code, method, r.Method, d.Status())).Observe(float64(size))
@@ -162,7 +190,7 @@ func InstrumentHandlerRequestSize(obs prometheus.ObserverVec, next http.Handler)
 func InstrumentHandlerResponseSize(obs prometheus.ObserverVec, next http.Handler) http.Handler {
 	code, method := checkLabels(obs)
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		d := newDelegator(w)
+		d := newDelegator(w, nil)
 		next.ServeHTTP(d, r)
 		obs.With(labels(code, method, r.Method, d.Status())).Observe(float64(d.Written()))
 	})
@@ -418,10 +446,11 @@ type delegator interface {
 type responseWriterDelegator struct {
 	http.ResponseWriter
 
-	handler, method string
-	status          int
-	written         int64
-	wroteHeader     bool
+	handler, method    string
+	status             int
+	written            int64
+	wroteHeader        bool
+	observeWriteHeader func(int)
 }
 
 func (r *responseWriterDelegator) Status() int {
@@ -436,6 +465,9 @@ func (r *responseWriterDelegator) WriteHeader(code int) {
 	r.status = code
 	r.wroteHeader = true
 	r.ResponseWriter.WriteHeader(code)
+	if r.observeWriteHeader != nil {
+		r.observeWriteHeader(code)
+	}
 }
 
 func (r *responseWriterDelegator) Write(b []byte) (int, error) {

prometheus/promhttp/instrument_server_test.go

@@ -48,6 +48,16 @@ func TestMiddlewareAPI(t *testing.T) {
 		[]string{"method"},
 	)
 
+	writeHeaderVec := prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Name:        "write_header_duration_seconds",
+			Help:        "A histogram of time to first write latencies.",
+			Buckets:     prometheus.DefBuckets,
+			ConstLabels: prometheus.Labels{"handler": "api"},
+		},
+		[]string{},
+	)
+
 	responseSize := prometheus.NewHistogramVec(
 		prometheus.HistogramOpts{
 			Name:    "push_request_size_bytes",
@@ -61,12 +71,14 @@ func TestMiddlewareAPI(t *testing.T) {
 		w.Write([]byte("OK"))
 	})
 
-	reg.MustRegister(inFlightGauge, counter, histVec, responseSize)
+	reg.MustRegister(inFlightGauge, counter, histVec, responseSize, writeHeaderVec)
 
 	chain := InstrumentHandlerInFlight(inFlightGauge,
 		InstrumentHandlerCounter(counter,
 			InstrumentHandlerDuration(histVec,
-				InstrumentHandlerResponseSize(responseSize, handler),
+				InstrumentHandlerTimeToWriteHeader(writeHeaderVec,
+					InstrumentHandlerResponseSize(responseSize, handler),
+				),
 			),
 		),
 	)
@@ -76,6 +88,23 @@ func TestMiddlewareAPI(t *testing.T) {
 	chain.ServeHTTP(w, r)
 }
 
+func TestInstrumentTimeToFirstWrite(t *testing.T) {
+	var i int
+	dobs := &responseWriterDelegator{
+		ResponseWriter: httptest.NewRecorder(),
+		observeWriteHeader: func(status int) {
+			i = status
+		},
+	}
+	d := newDelegator(dobs, nil)
+
+	d.WriteHeader(http.StatusOK)
+
+	if i != http.StatusOK {
+		t.Fatalf("failed to execute observeWriteHeader")
+	}
+}
+
 func ExampleInstrumentHandlerDuration() {
 	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{
 		Name: "in_flight_requests",