From 339ba92b8f2fd9a7dec6bbff189c793402aeff19 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal.matczuk@saucelabs.com>
Date: Wed, 28 Feb 2024 14:22:48 +0100
Subject: [PATCH] middleware(prometheus): add http_requests_in_flight metric

Fixes #453
---
 http_proxy.go                                 |  5 +++++
 middleware/prometheus.go                      | 22 +++++++++++++++----
 .../testdata/TestPrometheusWrap.golden.txt    |  3 +++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/http_proxy.go b/http_proxy.go
index 0b08cd77..70463d80 100644
--- a/http_proxy.go
+++ b/http_proxy.go
@@ -374,6 +374,11 @@ func (hp *HTTPProxy) middlewareStack() (martian.RequestResponseModifier, *martia
 		stack.AddResponseModifier(p)
 
 		trace = new(martian.ProxyTrace)
+		trace.ReadRequest = func(info martian.ReadRequestInfo) {
+			if info.Req != nil {
+				p.ReadRequest(info.Req)
+			}
+		}
 		trace.WroteResponse = func(info martian.WroteResponseInfo) {
 			if info.Res != nil {
 				p.WroteResponse(info.Res)
diff --git a/middleware/prometheus.go b/middleware/prometheus.go
index f34b1d33..73917e5f 100644
--- a/middleware/prometheus.go
+++ b/middleware/prometheus.go
@@ -40,10 +40,11 @@ var sizeBuckets = []float64{ //nolint:gochecknoglobals // this is a global varia
 // Unlike the promhttp.InstrumentHandler* chaining, this middleware creates only one delegator per request.
 // It partitions the metrics by HTTP status code, HTTP method, destination host name and source IP.
 type Prometheus struct {
-	requestsTotal   *prometheus.CounterVec
-	requestDuration *prometheus.HistogramVec
-	requestSize     *prometheus.HistogramVec
-	responseSize    *prometheus.HistogramVec
+	requestsInFlight *prometheus.GaugeVec
+	requestsTotal    *prometheus.CounterVec
+	requestDuration  *prometheus.HistogramVec
+	requestSize      *prometheus.HistogramVec
+	responseSize     *prometheus.HistogramVec
 }
 
 func NewPrometheus(r prometheus.Registerer, namespace string) *Prometheus {
@@ -54,6 +55,11 @@ func NewPrometheus(r prometheus.Registerer, namespace string) *Prometheus {
 	l := []string{"code", "method"}
 
 	return &Prometheus{
+		requestsInFlight: f.NewGaugeVec(prometheus.GaugeOpts{
+			Namespace: namespace,
+			Name:      "http_requests_in_flight",
+			Help:      "Current number of HTTP requests being served.",
+		}, []string{"method"}),
 		requestsTotal: f.NewCounterVec(prometheus.CounterOpts{
 			Namespace: namespace,
 			Name:      "http_requests_total",
@@ -82,6 +88,8 @@ func NewPrometheus(r prometheus.Registerer, namespace string) *Prometheus {
 
 func (p *Prometheus) Wrap(h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		p.requestsInFlight.WithLabelValues(r.Method).Inc()
+
 		d := newDelegator(w, nil)
 
 		r.Body = bodyCounter(r.Body)
@@ -100,6 +108,7 @@ func (p *Prometheus) Wrap(h http.Handler) http.Handler {
 			reqSize = c.Count()
 		}
 
+		p.requestsInFlight.WithLabelValues(r.Method).Dec()
 		p.requestSize.WithLabelValues(lv[:]...).Observe(float64(reqSize))
 		p.responseSize.WithLabelValues(lv[:]...).Observe(float64(d.Written()))
 	})
@@ -115,12 +124,17 @@ func (p *Prometheus) ModifyResponse(res *http.Response) error {
 	return nil
 }
 
+func (p *Prometheus) ReadRequest(req *http.Request) {
+	p.requestsInFlight.WithLabelValues(req.Method).Inc()
+}
+
 func (p *Prometheus) WroteResponse(res *http.Response) {
 	elapsed := martian.ContextDuration(res.Request.Context()).Seconds()
 
 	req := res.Request
 	lv := [2]string{strconv.Itoa(res.StatusCode), req.Method}
 
+	p.requestsInFlight.WithLabelValues(req.Method).Dec()
 	p.requestsTotal.WithLabelValues(lv[:]...).Inc()
 	p.requestDuration.WithLabelValues(lv[:]...).Observe(elapsed)
 
diff --git a/middleware/testdata/TestPrometheusWrap.golden.txt b/middleware/testdata/TestPrometheusWrap.golden.txt
index 840e00d8..c8502a12 100644
--- a/middleware/testdata/TestPrometheusWrap.golden.txt
+++ b/middleware/testdata/TestPrometheusWrap.golden.txt
@@ -29,6 +29,9 @@ test_http_request_size_bytes_bucket{code="200",method="GET",le="1.048576e+07"} 4
 test_http_request_size_bytes_bucket{code="200",method="GET",le="+Inf"} 4
 test_http_request_size_bytes_sum{code="200",method="GET"} 0
 test_http_request_size_bytes_count{code="200",method="GET"} 4
+# HELP test_http_requests_in_flight Current number of HTTP requests being served.
+# TYPE test_http_requests_in_flight gauge
+test_http_requests_in_flight{method="GET"} 0
 # HELP test_http_requests_total Total number of HTTP requests processed.
 # TYPE test_http_requests_total counter
 test_http_requests_total{code="200",method="GET"} 4