109 lines
		
	
	
		
			3.2 KiB
		
	
	
	
		
			Go
		
	
	
	
			
		
		
	
	
			109 lines
		
	
	
		
			3.2 KiB
		
	
	
	
		
			Go
		
	
	
	
| package measured_http
 | |
| 
 | |
| import (
 | |
| 	"net/http"
 | |
| 	"strconv"
 | |
| 
 | |
| 	"github.com/jmhodges/clock"
 | |
| 	"github.com/prometheus/client_golang/prometheus"
 | |
| 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 | |
| )
 | |
| 
 | |
| // responseWriterWithStatus satisfies http.ResponseWriter, but keeps track of the
 | |
| // status code for gathering stats.
 | |
| type responseWriterWithStatus struct {
 | |
| 	http.ResponseWriter
 | |
| 	code int
 | |
| }
 | |
| 
 | |
| // WriteHeader stores a status code for generating stats.
 | |
| func (r *responseWriterWithStatus) WriteHeader(code int) {
 | |
| 	r.code = code
 | |
| 	r.ResponseWriter.WriteHeader(code)
 | |
| }
 | |
| 
 | |
| // Write writes the body and sets the status code to 200 if a status code
 | |
| // has not already been set.
 | |
| func (r *responseWriterWithStatus) Write(body []byte) (int, error) {
 | |
| 	if r.code == 0 {
 | |
| 		r.code = http.StatusOK
 | |
| 	}
 | |
| 	return r.ResponseWriter.Write(body)
 | |
| }
 | |
| 
 | |
| // serveMux is a partial interface wrapper for the method http.ServeMux
 | |
| // exposes that we use. This is needed so that we can replace the default
 | |
| // http.ServeMux in ocsp-responder where we don't want to use its path
 | |
| // canonicalization.
 | |
| type serveMux interface {
 | |
| 	Handler(*http.Request) (http.Handler, string)
 | |
| }
 | |
| 
 | |
| // MeasuredHandler wraps an http.Handler and records prometheus stats
 | |
| type MeasuredHandler struct {
 | |
| 	serveMux
 | |
| 	clk clock.Clock
 | |
| 	// Normally this is always responseTime, but we override it for testing.
 | |
| 	stat *prometheus.HistogramVec
 | |
| 	// inFlightRequestsGauge is a gauge that tracks the number of requests
 | |
| 	// currently in flight, labeled by endpoint.
 | |
| 	inFlightRequestsGauge *prometheus.GaugeVec
 | |
| }
 | |
| 
 | |
| func New(m serveMux, clk clock.Clock, stats prometheus.Registerer, opts ...otelhttp.Option) http.Handler {
 | |
| 	responseTime := prometheus.NewHistogramVec(
 | |
| 		prometheus.HistogramOpts{
 | |
| 			Name: "response_time",
 | |
| 			Help: "Time taken to respond to a request",
 | |
| 		},
 | |
| 		[]string{"endpoint", "method", "code"})
 | |
| 	stats.MustRegister(responseTime)
 | |
| 
 | |
| 	inFlightRequestsGauge := prometheus.NewGaugeVec(
 | |
| 		prometheus.GaugeOpts{
 | |
| 			Name: "in_flight_requests",
 | |
| 			Help: "Tracks the number of WFE requests currently in flight, labeled by endpoint.",
 | |
| 		},
 | |
| 		[]string{"endpoint"},
 | |
| 	)
 | |
| 	stats.MustRegister(inFlightRequestsGauge)
 | |
| 
 | |
| 	return otelhttp.NewHandler(&MeasuredHandler{
 | |
| 		serveMux:              m,
 | |
| 		clk:                   clk,
 | |
| 		stat:                  responseTime,
 | |
| 		inFlightRequestsGauge: inFlightRequestsGauge,
 | |
| 	}, "server", opts...)
 | |
| }
 | |
| 
 | |
| func (h *MeasuredHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 | |
| 	begin := h.clk.Now()
 | |
| 	rwws := &responseWriterWithStatus{w, 0}
 | |
| 
 | |
| 	subHandler, pattern := h.Handler(r)
 | |
| 	h.inFlightRequestsGauge.WithLabelValues(pattern).Inc()
 | |
| 	defer h.inFlightRequestsGauge.WithLabelValues(pattern).Dec()
 | |
| 
 | |
| 	// Use the method string only if it's a recognized HTTP method. This avoids
 | |
| 	// ballooning timeseries with invalid methods from public input.
 | |
| 	var method string
 | |
| 	switch r.Method {
 | |
| 	case http.MethodGet, http.MethodHead, http.MethodPost, http.MethodPut,
 | |
| 		http.MethodPatch, http.MethodDelete, http.MethodConnect,
 | |
| 		http.MethodOptions, http.MethodTrace:
 | |
| 		method = r.Method
 | |
| 	default:
 | |
| 		method = "unknown"
 | |
| 	}
 | |
| 
 | |
| 	defer func() {
 | |
| 		h.stat.With(prometheus.Labels{
 | |
| 			"endpoint": pattern,
 | |
| 			"method":   method,
 | |
| 			"code":     strconv.Itoa(rwws.code),
 | |
| 		}).Observe(h.clk.Since(begin).Seconds())
 | |
| 	}()
 | |
| 
 | |
| 	subHandler.ServeHTTP(rwws, r)
 | |
| }
 |