refactor: make internal/lru and internal/ratelimiter metrics optional

6 files changed, 91 insertions, 40 deletions

diff --git a/internal/lru/lru.go b/internal/lru/lru.go
index f36d2a10..c30573ae 100644
--- a/internal/lru/lru.go
+++ b/internal/lru/lru.go
@@ -17,6 +17,14 @@ const getsPerPromote = 64
 // needs to be pruned on OOM, this prunes 1/16 of items
 const itemsToPruneDiv = 16
 
+// 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
+const DefaultSourceIPItems = 5000
+const DefaultSourceIPExpirationInterval = time.Minute
+
+// Option function to configure a Cache
+type Option func(*Cache)
+
 // Cache wraps a ccache and allows setting custom metrics for hits/misses.
 type Cache struct {
 	op                  string
@@ -27,22 +35,29 @@ type Cache struct {
 }
 
 // New creates an LRU cache
-func New(op string, maxEntries int64, duration time.Duration, cachedEntriesMetric *prometheus.GaugeVec, cacheRequestsMetric *prometheus.CounterVec) *Cache {
+func New(op string, maxEntries int64, duration time.Duration, opts ...Option) *Cache {
+	c := &Cache{
+		op:       op,
+		duration: duration,
+	}
+
 	configuration := ccache.Configure()
 	configuration.MaxSize(maxEntries)
 	configuration.ItemsToPrune(uint32(maxEntries) / itemsToPruneDiv)
 	configuration.GetsPerPromote(getsPerPromote) // if item gets requested frequently promote it
 	configuration.OnDelete(func(*ccache.Item) {
-		cachedEntriesMetric.WithLabelValues(op).Dec()
+		if c.metricCachedEntries != nil {
+			c.metricCachedEntries.WithLabelValues(op).Dec()
+		}
 	})
 
-	return &Cache{
-		op:                  op,
-		cache:               ccache.New(configuration),
-		duration:            duration,
-		metricCachedEntries: cachedEntriesMetric,
-		metricCacheRequests: cacheRequestsMetric,
+	c.cache = ccache.New(configuration)
+
+	for _, opt := range opts {
+		opt(c)
 	}
+
+	return c
 }
 
 // FindOrFetch will try to get the item from the cache if exists and is not expired.
@@ -51,20 +66,40 @@ func (c *Cache) FindOrFetch(cacheNamespace, key string, fetchFn func() (interfac
 	item := c.cache.Get(cacheNamespace + key)
 
 	if item != nil && !item.Expired() {
-		c.metricCacheRequests.WithLabelValues(c.op, "hit").Inc()
+		if c.metricCacheRequests != nil {
+			c.metricCacheRequests.WithLabelValues(c.op, "hit").Inc()
+		}
 		return item.Value(), nil
 	}
 
 	value, err := fetchFn()
 	if err != nil {
-		c.metricCacheRequests.WithLabelValues(c.op, "error").Inc()
+		if c.metricCacheRequests != nil {
+			c.metricCacheRequests.WithLabelValues(c.op, "error").Inc()
+		}
 		return nil, err
 	}
 
-	c.metricCacheRequests.WithLabelValues(c.op, "miss").Inc()
-	c.metricCachedEntries.WithLabelValues(c.op).Inc()
+	if c.metricCacheRequests != nil {
+		c.metricCacheRequests.WithLabelValues(c.op, "miss").Inc()
+	}
+	if c.metricCachedEntries != nil {
+		c.metricCachedEntries.WithLabelValues(c.op).Inc()
+	}
 
 	c.cache.Set(cacheNamespace+key, value, c.duration)
 
 	return value, nil
 }
+
+func WithCachedEntriesMetric(m *prometheus.GaugeVec) Option {
+	return func(c *Cache) {
+		c.metricCachedEntries = m
+	}
+}
+
+func WithCachedRequestsMetric(m *prometheus.CounterVec) Option {
+	return func(c *Cache) {
+		c.metricCacheRequests = m
+	}
+}
diff --git a/internal/ratelimiter/middleware.go b/internal/ratelimiter/middleware.go
index 0cd5b81e..f26cb2e0 100644
--- a/internal/ratelimiter/middleware.go
+++ b/internal/ratelimiter/middleware.go
@@ -28,12 +28,16 @@ func (rl *RateLimiter) SourceIPLimiter(handler http.Handler) http.Handler {
 			// Only drop requests once FF_ENABLE_RATE_LIMITER is enabled
 			// https://gitlab.com/gitlab-org/gitlab-pages/-/issues/629
 			if rateLimiterEnabled() {
-				rl.sourceIPBlockedCount.WithLabelValues("true").Inc()
+				if rl.sourceIPBlockedCount != nil {
+					rl.sourceIPBlockedCount.WithLabelValues("true").Inc()
+				}
 				httperrors.Serve429(w)
 				return
 			}
 
-			rl.sourceIPBlockedCount.WithLabelValues("false").Inc()
+			if rl.sourceIPBlockedCount != nil {
+				rl.sourceIPBlockedCount.WithLabelValues("false").Inc()
+			}
 		}
 
 		handler.ServeHTTP(w, r)
diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go
index 2e51fcad..132b24b9 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -10,6 +10,7 @@ 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"
 )
 
@@ -24,11 +25,14 @@ 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(blocked, cachedEntries, cacheReqs,
+			rl := New(
+				lru.New("source_ip",
+					lru.DefaultSourceIPItems,
+					lru.DefaultSourceIPExpirationInterval,
+				),
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(tc.sourceIPLimit),
 				WithSourceIPBurstSize(tc.sourceIPBurstSize),
@@ -83,7 +87,14 @@ func TestSourceIPLimiterDenyRequestsAfterBurst(t *testing.T) {
 
 	for tn, tc := range tcs {
 		t.Run(tn, func(t *testing.T) {
-			rl := New(blocked, cachedEntries, cacheReqs,
+			rl := New(
+				lru.New("source_ip",
+					lru.DefaultSourceIPItems,
+					lru.DefaultSourceIPExpirationInterval,
+					lru.WithCachedEntriesMetric(cachedEntries),
+					lru.WithCachedRequestsMetric(cacheReqs),
+				),
+				WithBlockedCountMetric(blocked),
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(1),
 				WithSourceIPBurstSize(1),
diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go
index 1359b19c..2e1304e2 100644
--- a/internal/ratelimiter/ratelimiter.go
+++ b/internal/ratelimiter/ratelimiter.go
@@ -17,11 +17,6 @@ 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
-	defaultSourceIPItems              = 5000
-	defaultSourceIPExpirationInterval = time.Minute
 )
 
 // Option function to configure a RateLimiter
@@ -42,19 +37,12 @@ type RateLimiter struct {
 }
 
 // New creates a new RateLimiter with default values that can be configured via Option functions
-func New(blockCountMetric, cachedEntriesMetric *prometheus.GaugeVec, cacheRequestsMetric *prometheus.CounterVec, opts ...Option) *RateLimiter {
+func New(c *lru.Cache, opts ...Option) *RateLimiter {
 	rl := &RateLimiter{
 		now:                    time.Now,
 		sourceIPLimitPerSecond: DefaultSourceIPLimitPerSecond,
 		sourceIPBurstSize:      DefaultSourceIPBurstSize,
-		sourceIPBlockedCount:   blockCountMetric,
-		sourceIPCache: lru.New(
-			"source_ip",
-			defaultSourceIPItems,
-			defaultSourceIPExpirationInterval,
-			cachedEntriesMetric,
-			cacheRequestsMetric,
-		),
+		sourceIPCache:          c,
 	}
 
 	for _, opt := range opts {
@@ -85,6 +73,12 @@ func WithSourceIPBurstSize(burst int) Option {
 	}
 }
 
+func WithBlockedCountMetric(m *prometheus.GaugeVec) Option {
+	return func(rl *RateLimiter) {
+		rl.sourceIPBlockedCount = 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 03393fb0..5fb9ef2b 100644
--- a/internal/ratelimiter/ratelimiter_test.go
+++ b/internal/ratelimiter/ratelimiter_test.go
@@ -7,6 +7,8 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/lru"
 )
 
 var (
@@ -48,11 +50,13 @@ 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(blocked, cachedEntries, cacheReqs,
+			rl := New(
+				lru.New("source_ip",
+					lru.DefaultSourceIPItems,
+					lru.DefaultSourceIPExpirationInterval,
+				),
 				WithNow(mockNow),
 				WithSourceIPLimitPerSecond(tc.sourceIPLimit),
 				WithSourceIPBurstSize(tc.sourceIPBurstSize),
@@ -74,9 +78,12 @@ func TestSourceIPAllowed(t *testing.T) {
 
 func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) {
 	rate := 10 * time.Millisecond
-	blocked, cachedEntries, cacheReqs := newTestMetrics(t)
 
-	rl := New(blocked, cachedEntries, cacheReqs,
+	rl := New(
+		lru.New("source_ip",
+			lru.DefaultSourceIPItems,
+			lru.DefaultSourceIPExpirationInterval,
+		),
 		WithSourceIPLimitPerSecond(float64(1/rate)),
 		WithSourceIPBurstSize(1),
 	)
diff --git a/internal/vfs/zip/vfs.go b/internal/vfs/zip/vfs.go
index c90d7614..67fe94dc 100644
--- a/internal/vfs/zip/vfs.go
+++ b/internal/vfs/zip/vfs.go
@@ -88,15 +88,15 @@ func New(cfg *config.ZipServing) vfs.VFS {
 		"data-offset",
 		defaultDataOffsetItems,
 		defaultDataOffsetExpirationInterval,
-		metrics.ZipCachedEntries,
-		metrics.ZipCacheRequests,
+		lru.WithCachedEntriesMetric(metrics.ZipCachedEntries),
+		lru.WithCachedRequestsMetric(metrics.ZipCacheRequests),
 	)
 	zipVFS.readlinkCache = lru.New(
 		"readlink",
 		defaultReadlinkItems,
 		defaultReadlinkExpirationInterval,
-		metrics.ZipCachedEntries,
-		metrics.ZipCacheRequests,
+		lru.WithCachedEntriesMetric(metrics.ZipCachedEntries),
+		lru.WithCachedRequestsMetric(metrics.ZipCacheRequests),
 	)
 
 	return zipVFS