refactor: pass metrics as params

4 files changed, 58 insertions, 17 deletions

diff --git a/app.go b/app.go
index af166e72..b97ec0bc 100644
--- a/app.go
+++ b/app.go
@@ -265,6 +265,9 @@ func (a *theApp) buildHandlerPipeline() (http.Handler, error) {
 
 	if a.config.RateLimit.SourceIPLimitPerSecond > 0 {
 		rl := ratelimiter.New(
+			metrics.RateLimitSourceIPBlockedCount,
+			metrics.RateLimitSourceIPCachedEntries,
+			metrics.RateLimitSourceIPCacheRequests,
 			ratelimiter.WithSourceIPLimitPerSecond(a.config.RateLimit.SourceIPLimitPerSecond),
 			ratelimiter.WithSourceIPBurstSize(a.config.RateLimit.SourceIPBurst),
 		)
diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go
index b0134b29..2e51fcad 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/prometheus/client_golang/prometheus/testutil"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
 
@@ -23,10 +24,11 @@ var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 func TestSourceIPLimiterWithDifferentLimits(t *testing.T) {
 	hook := testlog.NewGlobal()
 	testhelpers.SetEnvironmentVariable(t, testhelpers.FFEnableRateLimiter, "true")
+	blocked, cachedEntries, cacheReqs := newTestMetrics(t)
 
 	for tn, tc := range sharedTestCases {
 		t.Run(tn, func(t *testing.T) {
-			rl := New(
+			rl := New(blocked, cachedEntries, cacheReqs,
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(tc.sourceIPLimit),
 				WithSourceIPBurstSize(tc.sourceIPBurstSize),
@@ -63,6 +65,7 @@ func TestSourceIPLimiterWithDifferentLimits(t *testing.T) {
 
 func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) {
 	hook := testlog.NewGlobal()
+	blocked, cachedEntries, cacheReqs := newTestMetrics(t)
 
 	tcs := map[string]struct {
 		enabled        bool
@@ -72,23 +75,15 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) {
 			enabled:        false,
 			expectedStatus: http.StatusNoContent,
 		},
-		"disabled_rate_limit_https": {
-			enabled:        false,
-			expectedStatus: http.StatusNoContent,
-		},
 		"enabled_rate_limit_http_blocks": {
 			enabled:        true,
 			expectedStatus: http.StatusTooManyRequests,
 		},
-		"enabled_rate_limit_https_blocks": {
-			enabled:        true,
-			expectedStatus: http.StatusTooManyRequests,
-		},
 	}
 
 	for tn, tc := range tcs {
 		t.Run(tn, func(t *testing.T) {
-			rl := New(
+			rl := New(blocked, cachedEntries, cacheReqs,
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(1),
 				WithSourceIPBurstSize(1),
@@ -120,6 +115,24 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) {
 				require.Equal(t, tc.expectedStatus, res.StatusCode)
 				assertSourceIPLog(t, remoteAddr, hook)
 			}
+
+			blockedCount := testutil.ToFloat64(blocked.WithLabelValues("true"))
+			if tc.enabled {
+				require.Equal(t, float64(4), blockedCount, "blocked count")
+			} else {
+				require.Equal(t, float64(0), blockedCount, "blocked count")
+			}
+			blocked.Reset()
+
+			cachedCount := testutil.ToFloat64(cachedEntries.WithLabelValues("source_ip"))
+			require.Equal(t, float64(1), cachedCount, "cached count")
+			cachedEntries.Reset()
+
+			cacheReqMiss := testutil.ToFloat64(cacheReqs.WithLabelValues("source_ip", "miss"))
+			require.Equal(t, float64(1), cacheReqMiss, "miss count")
+			cacheReqHit := testutil.ToFloat64(cacheReqs.WithLabelValues("source_ip", "hit"))
+			require.Equal(t, float64(4), cacheReqHit, "hit count")
+			cacheReqs.Reset()
 		})
 	}
 }
diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go
index 1a622c6f..1359b19c 100644
--- a/internal/ratelimiter/ratelimiter.go
+++ b/internal/ratelimiter/ratelimiter.go
@@ -7,7 +7,6 @@ import (
 	"golang.org/x/time/rate"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/lru"
-	"gitlab.com/gitlab-org/gitlab-pages/metrics"
 )
 
 const (
@@ -43,18 +42,18 @@ type RateLimiter struct {
 }
 
 // New creates a new RateLimiter with default values that can be configured via Option functions
-func New(opts ...Option) *RateLimiter {
+func New(blockCountMetric, cachedEntriesMetric *prometheus.GaugeVec, cacheRequestsMetric *prometheus.CounterVec, opts ...Option) *RateLimiter {
 	rl := &RateLimiter{
 		now:                    time.Now,
 		sourceIPLimitPerSecond: DefaultSourceIPLimitPerSecond,
 		sourceIPBurstSize:      DefaultSourceIPBurstSize,
-		sourceIPBlockedCount:   metrics.RateLimitSourceIPBlockedCount,
+		sourceIPBlockedCount:   blockCountMetric,
 		sourceIPCache: lru.New(
 			"source_ip",
 			defaultSourceIPItems,
 			defaultSourceIPExpirationInterval,
-			metrics.RateLimitSourceIPCachedEntries,
-			metrics.RateLimitSourceIPCacheRequests,
+			cachedEntriesMetric,
+			cacheRequestsMetric,
 		),
 	}
 
diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go
index 77da8e81..03393fb0 100644
--- a/internal/ratelimiter/ratelimiter_test.go
+++ b/internal/ratelimiter/ratelimiter_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 )
 
@@ -47,9 +48,11 @@ var sharedTestCases = map[string]struct {
 func TestSourceIPAllowed(t *testing.T) {
 	t.Parallel()
 
+	blocked, cachedEntries, cacheReqs := newTestMetrics(t)
+
 	for tn, tc := range sharedTestCases {
 		t.Run(tn, func(t *testing.T) {
-			rl := New(
+			rl := New(blocked, cachedEntries, cacheReqs,
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(tc.sourceIPLimit),
 				WithSourceIPBurstSize(tc.sourceIPBurstSize),
@@ -71,7 +74,9 @@ func TestSourceIPAllowed(t *testing.T) {
 
 func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) {
 	rate := 10 * time.Millisecond
-	rl := New(
+	blocked, cachedEntries, cacheReqs := newTestMetrics(t)
+
+	rl := New(blocked, cachedEntries, cacheReqs,
 		WithSourceIPLimitPerSecond(float64(1/rate)),
 		WithSourceIPBurstSize(1),
 	)
@@ -104,3 +109,24 @@ func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) {
 
 	wg.Wait()
 }
+
+func newTestMetrics(t *testing.T) (*prometheus.GaugeVec, *prometheus.GaugeVec, *prometheus.CounterVec) {
+	t.Helper()
+
+	blockedGauge := prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Name: t.Name(),
+		},
+		[]string{"enforced"},
+	)
+
+	cachedEntries := prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Name: t.Name(),
+	}, []string{"op"})
+
+	cacheReqs := prometheus.NewCounterVec(prometheus.CounterOpts{
+		Name: t.Name(),
+	}, []string{"op", "cache"})
+
+	return blockedGauge, cachedEntries, cacheReqs
+}