refactor: encapsulate lru cache in ratelimiter package

4 files changed, 37 insertions, 20 deletions

diff --git a/app.go b/app.go
index 7d8ffebc..508b11fa 100644
--- a/app.go
+++ b/app.go
@@ -31,7 +31,6 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/handlers"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/logging"
-	"gitlab.com/gitlab-org/gitlab-pages/internal/lru"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/netutil"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/ratelimiter"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/rejectmethods"
@@ -261,13 +260,9 @@ func (a *theApp) buildHandlerPipeline() (http.Handler, error) {
 
 	if a.config.RateLimit.SourceIPLimitPerSecond > 0 {
 		rl := ratelimiter.New(
-			lru.New("source_ip",
-				// based on an avg ~4,000 unique IPs per minute
-				// https://log.gprd.gitlab.net/app/lens#/edit/f7110d00-2013-11ec-8c8e-ed83b5469915?_g=h@e78830b
-				lru.WithMaxSize(5000),
-				lru.WithCachedEntriesMetric(metrics.RateLimitSourceIPCachedEntries),
-				lru.WithCachedRequestsMetric(metrics.RateLimitSourceIPCacheRequests),
-			),
+			ratelimiter.WithSourceIPCacheMaxSize(ratelimiter.DefaultSourceIPCacheSize),
+			ratelimiter.WithSourceIPCachedEntriesMetric(metrics.RateLimitSourceIPCachedEntries),
+			ratelimiter.WithSourceIPCachedRequestsMetric(metrics.RateLimitSourceIPCacheRequests),
 			ratelimiter.WithBlockedCountMetric(metrics.RateLimitSourceIPBlockedCount),
 			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 6560956c..2cf3b3e5 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -10,7 +10,6 @@ import (
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
 
-	"gitlab.com/gitlab-org/gitlab-pages/internal/lru"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
 )
 
@@ -29,7 +28,6 @@ func TestSourceIPLimiterWithDifferentLimits(t *testing.T) {
 	for tn, tc := range sharedTestCases {
 		t.Run(tn, func(t *testing.T) {
 			rl := New(
-				lru.New("source_ip"),
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(tc.sourceIPLimit),
 				WithSourceIPBurstSize(tc.sourceIPBurstSize),
@@ -85,10 +83,8 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) {
 	for tn, tc := range tcs {
 		t.Run(tn, func(t *testing.T) {
 			rl := New(
-				lru.New("source_ip",
-					lru.WithCachedEntriesMetric(cachedEntries),
-					lru.WithCachedRequestsMetric(cacheReqs),
-				),
+				WithSourceIPCachedEntriesMetric(cachedEntries),
+				WithSourceIPCachedRequestsMetric(cacheReqs),
 				WithBlockedCountMetric(blocked),
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(1),
diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go
index 2e1304e2..37aca020 100644
--- a/internal/ratelimiter/ratelimiter.go
+++ b/internal/ratelimiter/ratelimiter.go
@@ -17,6 +17,10 @@ const (
 	// DefaultSourceIPBurstSize is the maximum burst allowed per rate limiter.
 	// E.g. The first 100 requests within 1s will succeed, but the 101st will fail.
 	DefaultSourceIPBurstSize = 100
+
+	// based on an avg ~4,000 unique IPs per minute
+	// https://log.gprd.gitlab.net/app/lens#/edit/f7110d00-2013-11ec-8c8e-ed83b5469915?_g=h@e78830b
+	DefaultSourceIPCacheSize = 5000
 )
 
 // Option function to configure a RateLimiter
@@ -33,22 +37,25 @@ type RateLimiter struct {
 	sourceIPBurstSize      int
 	sourceIPBlockedCount   *prometheus.GaugeVec
 	sourceIPCache          *lru.Cache
+
+	sourceIPCacheOptions []lru.Option
 	// TODO: add domainCache https://gitlab.com/gitlab-org/gitlab-pages/-/issues/630
 }
 
 // New creates a new RateLimiter with default values that can be configured via Option functions
-func New(c *lru.Cache, opts ...Option) *RateLimiter {
+func New(opts ...Option) *RateLimiter {
 	rl := &RateLimiter{
 		now:                    time.Now,
 		sourceIPLimitPerSecond: DefaultSourceIPLimitPerSecond,
 		sourceIPBurstSize:      DefaultSourceIPBurstSize,
-		sourceIPCache:          c,
 	}
 
 	for _, opt := range opts {
 		opt(rl)
 	}
 
+	rl.sourceIPCache = lru.New("source_ip", rl.sourceIPCacheOptions...)
+
 	return rl
 }
 
@@ -73,12 +80,35 @@ func WithSourceIPBurstSize(burst int) Option {
 	}
 }
 
+// WithBlockedCountMetric configures metric reporting how many requests were blocked based by IP
 func WithBlockedCountMetric(m *prometheus.GaugeVec) Option {
 	return func(rl *RateLimiter) {
 		rl.sourceIPBlockedCount = m
 	}
 }
 
+// WithSourceIPCacheMaxSize configures cache size for source IP ratelimiter
+func WithSourceIPCacheMaxSize(size int64) Option {
+	return func(rl *RateLimiter) {
+		rl.sourceIPCacheOptions = append(rl.sourceIPCacheOptions, lru.WithMaxSize(size))
+	}
+}
+
+// WithSourceIPCachedEntriesMetric configures metric reporting how many IPs are currently stored in
+// source-IP rate-limiter cache
+func WithSourceIPCachedEntriesMetric(m *prometheus.GaugeVec) Option {
+	return func(rl *RateLimiter) {
+		rl.sourceIPCacheOptions = append(rl.sourceIPCacheOptions, lru.WithCachedEntriesMetric(m))
+	}
+}
+
+// WithSourceIPCachedRequestsMetric configures metric for how many times we ask source IP cache
+func WithSourceIPCachedRequestsMetric(m *prometheus.CounterVec) Option {
+	return func(rl *RateLimiter) {
+		rl.sourceIPCacheOptions = append(rl.sourceIPCacheOptions, lru.WithCachedRequestsMetric(m))
+	}
+}
+
 func (rl *RateLimiter) getSourceIPLimiter(sourceIP string) *rate.Limiter {
 	limiterI, _ := rl.sourceIPCache.FindOrFetch(sourceIP, sourceIP, func() (interface{}, error) {
 		return rate.NewLimiter(rate.Limit(rl.sourceIPLimitPerSecond), rl.sourceIPBurstSize), nil
diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go
index a85663af..d6be214a 100644
--- a/internal/ratelimiter/ratelimiter_test.go
+++ b/internal/ratelimiter/ratelimiter_test.go
@@ -7,8 +7,6 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
-
-	"gitlab.com/gitlab-org/gitlab-pages/internal/lru"
 )
 
 var (
@@ -53,7 +51,6 @@ func TestSourceIPAllowed(t *testing.T) {
 	for tn, tc := range sharedTestCases {
 		t.Run(tn, func(t *testing.T) {
 			rl := New(
-				lru.New("source_ip"),
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(tc.sourceIPLimit),
 				WithSourceIPBurstSize(tc.sourceIPBurstSize),
@@ -77,7 +74,6 @@ func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) {
 	rate := 10 * time.Millisecond
 
 	rl := New(
-		lru.New("source_ip"),
 		WithSourceIPLimitPerSecond(float64(1/rate)),
 		WithSourceIPBurstSize(1),
 	)