From 65e0a968d051497c00682e72dc74aa405f88d8b6 Mon Sep 17 00:00:00 2001
From: Nicholas Wiersma <nick@wiersma.co.za>
Date: Sun, 22 May 2022 11:12:13 +0200
Subject: [PATCH] fix: stats prom buckets (#47)

---
 http/middleware/middleware.go      |  8 ++++++++
 http/middleware/middleware_test.go | 20 ++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/http/middleware/middleware.go b/http/middleware/middleware.go
index 5f2858e..745642b 100644
--- a/http/middleware/middleware.go
+++ b/http/middleware/middleware.go
@@ -7,6 +7,7 @@ import (
 	"github.com/hamba/logger/v2"
 	"github.com/hamba/logger/v2/ctx"
 	"github.com/hamba/statter/v2"
+	"github.com/hamba/statter/v2/reporter/prometheus"
 	"github.com/hamba/statter/v2/tags"
 	"github.com/hamba/timex/mono"
 )
@@ -40,6 +41,13 @@ func Recovery(log *logger.Logger) func(http.Handler) http.Handler {
 
 // WithStats collects statistics about HTTP requests.
 func WithStats(name string, s *statter.Statter, h http.Handler) http.Handler {
+	prometheus.RegisterHistogram(s,
+		"response.size",
+		[]string{"handler", "code"},
+		[]float64{200, 500, 900, 1500, 5000, 10000},
+		"The size of a response in bytes",
+	)
+
 	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		t := make([]statter.Tag, 1, 2)
 		if name == "" {
diff --git a/http/middleware/middleware_test.go b/http/middleware/middleware_test.go
index 66eee5d..be318c9 100644
--- a/http/middleware/middleware_test.go
+++ b/http/middleware/middleware_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/hamba/logger/v2"
 	"github.com/hamba/pkg/v2/http/middleware"
 	"github.com/hamba/statter/v2"
+	"github.com/hamba/statter/v2/reporter/prometheus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -115,6 +116,25 @@ func TestWithStats(t *testing.T) {
 	}
 }
 
+func TestWithStats_Prometheus(t *testing.T) {
+	reporter := prometheus.New("test")
+	s := statter.New(reporter, time.Second)
+
+	h := middleware.WithStats("test-handler", s, http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(305)
+		}),
+	)
+
+	resp := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/foobar", nil)
+
+	h.ServeHTTP(resp, req)
+
+	err := s.Close()
+	require.NoError(t, err)
+}
+
 type mockReporter struct {
 	mock.Mock
 }