From 53bee6fbecdc2a9d5884fa7259bc211ecd6d54fc Mon Sep 17 00:00:00 2001
From: DarthSim <darthsim@gmail.com>
Date: Sat, 4 Nov 2023 20:12:16 +0300
Subject: [PATCH] Add status_codes_total counter to Prometheus metrics

---
 CHANGELOG.md                     |  1 +
 metrics/metrics.go               |  2 +-
 metrics/prometheus/prometheus.go | 30 +++++++++++++++++++++++++-----
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f561cc53..285cd7d8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## [Unreleased]
 ### Add
+- Add `status_codes_total` counter to Prometheus metrics.
 - (pro) Add the `IMGPROXY_VIDEO_THUMBNAIL_KEYFRAMES` config and the [video_thumbnail_keyframes](https://docs.imgproxy.net/latest/generating_the_url?id=video-thumbnail-keyframes) processing option.
 - (pro) Add the [video_thumbnail_tile](https://docs.imgproxy.net/latest/generating_the_url?id=video-thumbnail-tile) processing option.
 
diff --git a/metrics/metrics.go b/metrics/metrics.go
index 5e768e31..18b341eb 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -47,7 +47,7 @@ func Enabled() bool {
 }
 
 func StartRequest(ctx context.Context, rw http.ResponseWriter, r *http.Request) (context.Context, context.CancelFunc, http.ResponseWriter) {
-	promCancel := prometheus.StartRequest()
+	promCancel, rw := prometheus.StartRequest(rw)
 	ctx, nrCancel, rw := newrelic.StartTransaction(ctx, rw, r)
 	ctx, ddCancel, rw := datadog.StartRootSpan(ctx, rw, r)
 	ctx, otelCancel, rw := otel.StartRootSpan(ctx, rw, r)
diff --git a/metrics/prometheus/prometheus.go b/metrics/prometheus/prometheus.go
index f3bb1e90..6ad2c9fb 100644
--- a/metrics/prometheus/prometheus.go
+++ b/metrics/prometheus/prometheus.go
@@ -4,8 +4,10 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strconv"
 	"time"
 
+	"github.com/felixge/httpsnoop"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	log "github.com/sirupsen/logrus"
@@ -18,8 +20,9 @@ import (
 var (
 	enabled = false
 
-	requestsTotal prometheus.Counter
-	errorsTotal   *prometheus.CounterVec
+	requestsTotal    prometheus.Counter
+	statusCodesTotal *prometheus.CounterVec
+	errorsTotal      *prometheus.CounterVec
 
 	requestDuration     prometheus.Histogram
 	requestSpanDuration *prometheus.HistogramVec
@@ -45,6 +48,12 @@ func Init() {
 		Help:      "A counter of the total number of HTTP requests imgproxy processed.",
 	})
 
+	statusCodesTotal = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Namespace: config.PrometheusNamespace,
+		Name:      "status_codes_total",
+		Help:      "A counter of the response status codes.",
+	}, []string{"status"})
+
 	errorsTotal = prometheus.NewCounterVec(prometheus.CounterOpts{
 		Namespace: config.PrometheusNamespace,
 		Name:      "errors_total",
@@ -108,6 +117,7 @@ func Init() {
 
 	prometheus.MustRegister(
 		requestsTotal,
+		statusCodesTotal,
 		errorsTotal,
 		requestDuration,
 		requestSpanDuration,
@@ -150,13 +160,23 @@ func StartServer(cancel context.CancelFunc) error {
 	return nil
 }
 
-func StartRequest() context.CancelFunc {
+func StartRequest(rw http.ResponseWriter) (context.CancelFunc, http.ResponseWriter) {
 	if !enabled {
-		return func() {}
+		return func() {}, rw
 	}
 
 	requestsTotal.Inc()
-	return startDuration(requestDuration)
+
+	newRw := httpsnoop.Wrap(rw, httpsnoop.Hooks{
+		WriteHeader: func(next httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return func(statusCode int) {
+				statusCodesTotal.With(prometheus.Labels{"status": strconv.Itoa(statusCode)}).Inc()
+				next(statusCode)
+			}
+		},
+	})
+
+	return startDuration(requestDuration), newRw
 }
 
 func StartQueueSegment() context.CancelFunc {