diff options
Diffstat (limited to 'internal')
-rw-r--r-- | internal/lru/lru.go | 79 | ||||
-rw-r--r-- | internal/ratelimiter/middleware.go | 8 | ||||
-rw-r--r-- | internal/ratelimiter/middleware_test.go | 12 | ||||
-rw-r--r-- | internal/ratelimiter/ratelimiter.go | 22 | ||||
-rw-r--r-- | internal/ratelimiter/ratelimiter_test.go | 11 | ||||
-rw-r--r-- | internal/vfs/zip/vfs.go | 16 |
6 files changed, 101 insertions, 47 deletions
diff --git a/internal/lru/lru.go b/internal/lru/lru.go index f36d2a10..dafb5347 100644 --- a/internal/lru/lru.go +++ b/internal/lru/lru.go @@ -17,32 +17,49 @@ const getsPerPromote = 64 // needs to be pruned on OOM, this prunes 1/16 of items const itemsToPruneDiv = 16 +const defaultCacheMaxSize = 1000 +const defaultCacheExpirationInterval = time.Minute + +// Option function to configure a Cache +type Option func(*Cache) + // Cache wraps a ccache and allows setting custom metrics for hits/misses. +// duration and maxSize are initialized to their default values but should +// be configured using WithExpirationInterval and WithMaxSize options. type Cache struct { op string duration time.Duration + maxSize int64 cache *ccache.Cache metricCachedEntries *prometheus.GaugeVec metricCacheRequests *prometheus.CounterVec } // 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, opts ...Option) *Cache { + c := &Cache{ + op: op, + duration: defaultCacheExpirationInterval, + maxSize: defaultCacheMaxSize, + } + + for _, opt := range opts { + opt(c) + } + configuration := ccache.Configure() - configuration.MaxSize(maxEntries) - configuration.ItemsToPrune(uint32(maxEntries) / itemsToPruneDiv) + configuration.MaxSize(c.maxSize) + configuration.ItemsToPrune(uint32(c.maxSize) / 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) + + return c } // FindOrFetch will try to get the item from the cache if exists and is not expired. @@ -51,20 +68,52 @@ 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 + } +} + +func WithExpirationInterval(t time.Duration) Option { + return func(c *Cache) { + c.duration = t + } +} + +func WithMaxSize(i int64) Option { + return func(c *Cache) { + c.maxSize = i + } +} 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..6560956c 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,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(blocked, cachedEntries, cacheReqs, + rl := New( + lru.New("source_ip"), WithNow(mockNow), WithSourceIPLimitPerSecond(tc.sourceIPLimit), WithSourceIPBurstSize(tc.sourceIPBurstSize), @@ -83,7 +84,12 @@ 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.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..a85663af 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,10 @@ 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"), WithNow(mockNow), WithSourceIPLimitPerSecond(tc.sourceIPLimit), WithSourceIPBurstSize(tc.sourceIPBurstSize), @@ -74,9 +75,9 @@ 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"), WithSourceIPLimitPerSecond(float64(1/rate)), WithSourceIPBurstSize(1), ) diff --git a/internal/vfs/zip/vfs.go b/internal/vfs/zip/vfs.go index 07afa691..a9825190 100644 --- a/internal/vfs/zip/vfs.go +++ b/internal/vfs/zip/vfs.go @@ -87,17 +87,17 @@ func New(cfg *config.ZipServing) vfs.VFS { // TODO: To be removed with https://gitlab.com/gitlab-org/gitlab-pages/-/issues/480 zipVFS.dataOffsetCache = lru.New( "data-offset", - defaultDataOffsetItems, - defaultDataOffsetExpirationInterval, - metrics.ZipCachedEntries, - metrics.ZipCacheRequests, + lru.WithMaxSize(defaultDataOffsetItems), + lru.WithExpirationInterval(defaultDataOffsetExpirationInterval), + lru.WithCachedEntriesMetric(metrics.ZipCachedEntries), + lru.WithCachedRequestsMetric(metrics.ZipCacheRequests), ) zipVFS.readlinkCache = lru.New( "readlink", - defaultReadlinkItems, - defaultReadlinkExpirationInterval, - metrics.ZipCachedEntries, - metrics.ZipCacheRequests, + lru.WithMaxSize(defaultReadlinkItems), + lru.WithExpirationInterval(defaultReadlinkExpirationInterval), + lru.WithCachedEntriesMetric(metrics.ZipCachedEntries), + lru.WithCachedRequestsMetric(metrics.ZipCacheRequests), ) return zipVFS |