From 49f2c8908c0d831e6a2880c8ad659116ad70d74d Mon Sep 17 00:00:00 2001 From: Jaime Martinez Date: Mon, 20 Sep 2021 16:16:45 +1000 Subject: feat: add ratelimiter package Changelog: added --- go.mod | 1 + go.sum | 1 + internal/ratelimiter/ratelimiter.go | 106 ++++++++++++++++++++++++++++++ internal/ratelimiter/ratelimiter_test.go | 108 +++++++++++++++++++++++++++++++ 4 files changed, 216 insertions(+) create mode 100644 internal/ratelimiter/ratelimiter.go create mode 100644 internal/ratelimiter/ratelimiter_test.go diff --git a/go.mod b/go.mod index 5f17d253..0ff516ed 100644 --- a/go.mod +++ b/go.mod @@ -25,4 +25,5 @@ require ( golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 + golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 ) diff --git a/go.sum b/go.sum index f004e86f..c57b16ad 100644 --- a/go.sum +++ b/go.sum @@ -398,6 +398,7 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= +golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 h1:SvFZT6jyqRaOeXpc5h/JSfZenJ2O330aBsf7JfSUXmQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20181221001348-537d06c36207/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go new file mode 100644 index 00000000..3965f49f --- /dev/null +++ b/internal/ratelimiter/ratelimiter.go @@ -0,0 +1,106 @@ +package ratelimiter + +import ( + "time" + + "github.com/prometheus/client_golang/prometheus" + "golang.org/x/time/rate" + + "gitlab.com/gitlab-org/gitlab-pages/internal/lru" +) + +const ( + // DefaultPerDomainFrequency the maximum number of requests per second to be allowed per domain. + // The default value of 25ms equals 1 request every 25ms -> 40 rps + DefaultPerDomainFrequency = 25 * time.Millisecond + // DefaultPerDomainBurstSize is the maximum burst allowed per rate limiter + // E.g. The first 40 requests within 25ms will succeed, but the 41st will fail until the next + // refill occurs at DefaultPerDomainFrequency, allowing only 1 request per rate frequency. + DefaultPerDomainBurstSize = 40 + + // avg of ~18,000 unique domains per hour + // https://log.gprd.gitlab.net/app/lens#/edit/3c45a610-15c9-11ec-a012-eb2e5674cacf?_g=h@e78830b + defaultDomainsItems = 20000 + defaultDomainsExpirationInterval = time.Hour +) + +type cache interface { + FindOrFetch(cacheNamespace, key string, fetchFn func() (interface{}, error)) (interface{}, error) +} + +// Option function to configure a RateLimiter +type Option func(*RateLimiter) + +// RateLimiter holds a map ot domain names with counters that enable rate limiting per domain. +// It uses "golang.org/x/time/rate" as its Token Bucket rate limiter per domain entry. +// See example https://www.fatalerrors.org/a/design-and-implementation-of-time-rate-limiter-for-golang-standard-library.html +// Cleanup runs every cleanupTimer iteration over all domains and removing them if +// the time since counter.lastSeen is greater than the domainMaxTTL. +type RateLimiter struct { + now func() time.Time + perDomainFrequency time.Duration + perDomainBurstSize int + domainsCache cache + // TODO: add sourceIPCache 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(opts ...Option) *RateLimiter { + rl := &RateLimiter{ + now: time.Now, + perDomainFrequency: DefaultPerDomainFrequency, + perDomainBurstSize: DefaultPerDomainBurstSize, + domainsCache: lru.New( + "domains", + defaultDomainsItems, + defaultDomainsExpirationInterval, + // TODO: @jaime to add proper metrics in subsequent MR + prometheus.NewGaugeVec(prometheus.GaugeOpts{}, []string{"op"}), + prometheus.NewCounterVec(prometheus.CounterOpts{}, []string{"op", "cache"}), + ), + } + + for _, opt := range opts { + opt(rl) + } + + return rl +} + +// WithNow replaces the RateLimiter now function +func WithNow(now func() time.Time) Option { + return func(rl *RateLimiter) { + rl.now = now + } +} + +// WithPerDomainFrequency allows configuring perDomain frequency for the RateLimiter +func WithPerDomainFrequency(d time.Duration) Option { + return func(rl *RateLimiter) { + rl.perDomainFrequency = d + } +} + +// WithPerDomainBurstSize configures burst per domain for the RateLimiter +func WithPerDomainBurstSize(burst int) Option { + return func(rl *RateLimiter) { + rl.perDomainBurstSize = burst + } +} + +func (rl *RateLimiter) getDomainCounter(domain string) *rate.Limiter { + limiterI, _ := rl.domainsCache.FindOrFetch(domain, domain, func() (interface{}, error) { + return rate.NewLimiter(rate.Every(rl.perDomainFrequency), rl.perDomainBurstSize), nil + }) + + return limiterI.(*rate.Limiter) +} + +// DomainAllowed checks that the requested domain can be accessed within +// the maxCountPerDomain in the given domainWindow. +func (rl *RateLimiter) DomainAllowed(domain string) (res bool) { + limiter := rl.getDomainCounter(domain) + + // AllowN allows us to use the rl.now function so we can test this more easily. + return limiter.AllowN(rl.now(), 1) +} diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go new file mode 100644 index 00000000..e2352209 --- /dev/null +++ b/internal/ratelimiter/ratelimiter_test.go @@ -0,0 +1,108 @@ +package ratelimiter + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +var ( + now = "2021-09-13T15:00:00Z" + validTime, _ = time.Parse(time.RFC3339, now) +) + +func mockNow() time.Time { + return validTime +} + +func TestDomainAllowed(t *testing.T) { + t.Parallel() + + tcs := map[string]struct { + now string + domainRate time.Duration + perDomainBurstPerSecond int + domain string + reqNum int + }{ + "one_request_per_second": { + domainRate: 1, // 1 per second + perDomainBurstPerSecond: 1, + reqNum: 2, + domain: "rate.gitlab.io", + }, + "one_request_per_second_but_big_bucket": { + domainRate: 1, // 1 per second + perDomainBurstPerSecond: 10, + reqNum: 11, + domain: "rate.gitlab.io", + }, + "three_req_per_second_bucket_size_one": { + domainRate: 3, // 3 per second + perDomainBurstPerSecond: 1, // max burst 1 means 1 at a time + reqNum: 3, + domain: "rate.gitlab.io", + }, + "10_requests_per_second": { + domainRate: 10, + perDomainBurstPerSecond: 10, + reqNum: 11, + domain: "rate.gitlab.io", + }, + } + + for tn, tc := range tcs { + t.Run(tn, func(t *testing.T) { + rl := New( + WithNow(mockNow), + WithPerDomainFrequency(tc.domainRate), + WithPerDomainBurstSize(tc.perDomainBurstPerSecond), + ) + + for i := 0; i < tc.reqNum; i++ { + got := rl.DomainAllowed(tc.domain) + if i < tc.perDomainBurstPerSecond { + require.Truef(t, got, "expected true for request no. %d", i+1) + } else { + require.False(t, got, "expected false for request no. %d", i+1) + } + } + }) + } +} + +func TestDomainAllowedWitSleeps(t *testing.T) { + rate := 10 * time.Millisecond + rl := New( + WithPerDomainFrequency(rate), + WithPerDomainBurstSize(1), + ) + + domain := "test.gitlab.io" + + t.Run("one request every 10ms with burst 1", func(t *testing.T) { + // prove cache entries per domain + t.Parallel() + + for i := 0; i < 10; i++ { + got := rl.DomainAllowed(domain) + require.Truef(t, got, "expected true for request no. %d", i+1) + time.Sleep(rate) + } + }) + + t.Run("requests start failing after reaching burst", func(t *testing.T) { + // prove cache entries per domain + t.Parallel() + + for i := 0; i < 5; i++ { + got := rl.DomainAllowed(domain + ".diff") + if i < 1 { + require.Truef(t, got, "expected true for request no. %d", i) + } else { + require.False(t, got, "expected false for request no. %d", i) + } + } + }) +} -- cgit v1.2.3 From 128bf6b4cbee628cdb1fe3ee26b9442a18b85ef3 Mon Sep 17 00:00:00 2001 From: Jaime Martinez Date: Wed, 22 Sep 2021 15:14:39 +1000 Subject: test: use goroutines --- internal/ratelimiter/ratelimiter.go | 29 ++++++------- internal/ratelimiter/ratelimiter_test.go | 71 ++++++++++++++++---------------- 2 files changed, 48 insertions(+), 52 deletions(-) diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go index 3965f49f..f30722c0 100644 --- a/internal/ratelimiter/ratelimiter.go +++ b/internal/ratelimiter/ratelimiter.go @@ -10,37 +10,34 @@ import ( ) const ( - // DefaultPerDomainFrequency the maximum number of requests per second to be allowed per domain. - // The default value of 25ms equals 1 request every 25ms -> 40 rps - DefaultPerDomainFrequency = 25 * time.Millisecond - // DefaultPerDomainBurstSize is the maximum burst allowed per rate limiter - // E.g. The first 40 requests within 25ms will succeed, but the 41st will fail until the next + // DefaultPerDomainFrequency is the rate in time.Duration at which the rate.Limiter + // bucket is filled with 1 token. A token is equivalent to a request. + // The default value of 20ms, or 1 request every 20ms, equals 50 requests per second. + DefaultPerDomainFrequency = 20 * time.Millisecond + // DefaultPerDomainBurstSize is the maximum burst allowed per rate limiter. + // E.g. The first 50 requests within 20ms will succeed, but the 51st will fail until the next // refill occurs at DefaultPerDomainFrequency, allowing only 1 request per rate frequency. - DefaultPerDomainBurstSize = 40 + DefaultPerDomainBurstSize = 50 - // avg of ~18,000 unique domains per hour + // based on an avg of ~18,000 unique domains per hour // https://log.gprd.gitlab.net/app/lens#/edit/3c45a610-15c9-11ec-a012-eb2e5674cacf?_g=h@e78830b defaultDomainsItems = 20000 defaultDomainsExpirationInterval = time.Hour ) -type cache interface { - FindOrFetch(cacheNamespace, key string, fetchFn func() (interface{}, error)) (interface{}, error) -} - // Option function to configure a RateLimiter type Option func(*RateLimiter) -// RateLimiter holds a map ot domain names with counters that enable rate limiting per domain. +// RateLimiter holds an LRU cache of elements to be rate limited. Currently, it supports +// a domainsCache and each item returns a rate.Limiter. // It uses "golang.org/x/time/rate" as its Token Bucket rate limiter per domain entry. // See example https://www.fatalerrors.org/a/design-and-implementation-of-time-rate-limiter-for-golang-standard-library.html -// Cleanup runs every cleanupTimer iteration over all domains and removing them if -// the time since counter.lastSeen is greater than the domainMaxTTL. +// It also holds a now function that can be mocked in unit tests. type RateLimiter struct { now func() time.Time perDomainFrequency time.Duration perDomainBurstSize int - domainsCache cache + domainsCache *lru.Cache // TODO: add sourceIPCache https://gitlab.com/gitlab-org/gitlab-pages/-/issues/630 } @@ -101,6 +98,6 @@ func (rl *RateLimiter) getDomainCounter(domain string) *rate.Limiter { func (rl *RateLimiter) DomainAllowed(domain string) (res bool) { limiter := rl.getDomainCounter(domain) - // AllowN allows us to use the rl.now function so we can test this more easily. + // AllowN allows us to use the rl.now function, so we can test this more easily. return limiter.AllowN(rl.now(), 1) } diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go index e2352209..96f50b1b 100644 --- a/internal/ratelimiter/ratelimiter_test.go +++ b/internal/ratelimiter/ratelimiter_test.go @@ -1,6 +1,7 @@ package ratelimiter import ( + "sync" "testing" "time" @@ -23,32 +24,27 @@ func TestDomainAllowed(t *testing.T) { now string domainRate time.Duration perDomainBurstPerSecond int - domain string reqNum int }{ - "one_request_per_second": { - domainRate: 1, // 1 per second + "one_request_per_nanosecond": { + domainRate: time.Nanosecond, // 1 per nanosecond perDomainBurstPerSecond: 1, reqNum: 2, - domain: "rate.gitlab.io", }, - "one_request_per_second_but_big_bucket": { - domainRate: 1, // 1 per second + "one_request_per_nanosecond_but_big_bucket": { + domainRate: time.Nanosecond, perDomainBurstPerSecond: 10, reqNum: 11, - domain: "rate.gitlab.io", }, "three_req_per_second_bucket_size_one": { domainRate: 3, // 3 per second perDomainBurstPerSecond: 1, // max burst 1 means 1 at a time reqNum: 3, - domain: "rate.gitlab.io", }, "10_requests_per_second": { domainRate: 10, perDomainBurstPerSecond: 10, reqNum: 11, - domain: "rate.gitlab.io", }, } @@ -61,48 +57,51 @@ func TestDomainAllowed(t *testing.T) { ) for i := 0; i < tc.reqNum; i++ { - got := rl.DomainAllowed(tc.domain) + got := rl.DomainAllowed("rate.gitlab.io") if i < tc.perDomainBurstPerSecond { - require.Truef(t, got, "expected true for request no. %d", i+1) + require.Truef(t, got, "expected true for request no. %d", i) } else { - require.False(t, got, "expected false for request no. %d", i+1) + // requests should fail after reaching tc.perDomainBurstPerSecond because mockNow + // always returns the same time + require.False(t, got, "expected false for request no. %d", i) } } }) } } -func TestDomainAllowedWitSleeps(t *testing.T) { +func TestSingleRateLimiterWithMultipleDomains(t *testing.T) { rate := 10 * time.Millisecond rl := New( WithPerDomainFrequency(rate), WithPerDomainBurstSize(1), ) - domain := "test.gitlab.io" + wg := sync.WaitGroup{} + wg.Add(3) - t.Run("one request every 10ms with burst 1", func(t *testing.T) { - // prove cache entries per domain - t.Parallel() + testFn := func(domain string) func(t *testing.T) { + return func(t *testing.T) { + go func() { + defer wg.Done() - for i := 0; i < 10; i++ { - got := rl.DomainAllowed(domain) - require.Truef(t, got, "expected true for request no. %d", i+1) - time.Sleep(rate) - } - }) - - t.Run("requests start failing after reaching burst", func(t *testing.T) { - // prove cache entries per domain - t.Parallel() - - for i := 0; i < 5; i++ { - got := rl.DomainAllowed(domain + ".diff") - if i < 1 { - require.Truef(t, got, "expected true for request no. %d", i) - } else { - require.False(t, got, "expected false for request no. %d", i) - } + for i := 0; i < 5; i++ { + got := rl.DomainAllowed(domain) + require.Truef(t, got, "expected true for request no. %d", i) + time.Sleep(rate) + } + }() } - }) + } + + first := "first.gitlab.io" + t.Run(first, testFn(first)) + + second := "second.gitlab.io" + t.Run(second, testFn(second)) + + third := "third.gitlab.io" + t.Run(third, testFn(third)) + + wg.Wait() } -- cgit v1.2.3 From 59913b292288277aa1192439fdfa97b1f59d03c5 Mon Sep 17 00:00:00 2001 From: Jaime Martinez Date: Tue, 28 Sep 2021 15:38:27 +1000 Subject: refactor: change to source IP instead --- internal/ratelimiter/ratelimiter.go | 78 ++++++++++++++++---------------- internal/ratelimiter/ratelimiter_test.go | 60 ++++++++++++------------ 2 files changed, 68 insertions(+), 70 deletions(-) diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go index f30722c0..0753e187 100644 --- a/internal/ratelimiter/ratelimiter.go +++ b/internal/ratelimiter/ratelimiter.go @@ -10,47 +10,46 @@ import ( ) const ( - // DefaultPerDomainFrequency is the rate in time.Duration at which the rate.Limiter - // bucket is filled with 1 token. A token is equivalent to a request. - // The default value of 20ms, or 1 request every 20ms, equals 50 requests per second. - DefaultPerDomainFrequency = 20 * time.Millisecond - // DefaultPerDomainBurstSize is the maximum burst allowed per rate limiter. - // E.g. The first 50 requests within 20ms will succeed, but the 51st will fail until the next - // refill occurs at DefaultPerDomainFrequency, allowing only 1 request per rate frequency. - DefaultPerDomainBurstSize = 50 - - // based on an avg of ~18,000 unique domains per hour - // https://log.gprd.gitlab.net/app/lens#/edit/3c45a610-15c9-11ec-a012-eb2e5674cacf?_g=h@e78830b - defaultDomainsItems = 20000 - defaultDomainsExpirationInterval = time.Hour + // DefaultSourceIPLimitPerSecond is the limit per second that rate.Limiter + // needs to generate tokens every second. + // The default value is 20 requests per second. + DefaultSourceIPLimitPerSecond = 20.0 + // DefaultSourceIPBurstSize is the maximum burst allowed per rate limiter. + // E.g. The first 20 requests within 1s will succeed, but the 21st will fail. + DefaultSourceIPBurstSize = 20 + + // 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 type Option func(*RateLimiter) // RateLimiter holds an LRU cache of elements to be rate limited. Currently, it supports -// a domainsCache and each item returns a rate.Limiter. -// It uses "golang.org/x/time/rate" as its Token Bucket rate limiter per domain entry. +// a sourceIPCache and each item returns a rate.Limiter. +// It uses "golang.org/x/time/rate" as its Token Bucket rate limiter per source IP entry. // See example https://www.fatalerrors.org/a/design-and-implementation-of-time-rate-limiter-for-golang-standard-library.html // It also holds a now function that can be mocked in unit tests. type RateLimiter struct { - now func() time.Time - perDomainFrequency time.Duration - perDomainBurstSize int - domainsCache *lru.Cache - // TODO: add sourceIPCache https://gitlab.com/gitlab-org/gitlab-pages/-/issues/630 + now func() time.Time + sourceIPLimitPerSecond float64 + sourceIPBurstSize int + sourceIPCache *lru.Cache + // 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(opts ...Option) *RateLimiter { rl := &RateLimiter{ - now: time.Now, - perDomainFrequency: DefaultPerDomainFrequency, - perDomainBurstSize: DefaultPerDomainBurstSize, - domainsCache: lru.New( - "domains", - defaultDomainsItems, - defaultDomainsExpirationInterval, + now: time.Now, + sourceIPLimitPerSecond: DefaultSourceIPLimitPerSecond, + sourceIPBurstSize: DefaultSourceIPBurstSize, + sourceIPCache: lru.New( + "source_ip", + defaultSourceIPItems, + defaultSourceIPExpirationInterval, // TODO: @jaime to add proper metrics in subsequent MR prometheus.NewGaugeVec(prometheus.GaugeOpts{}, []string{"op"}), prometheus.NewCounterVec(prometheus.CounterOpts{}, []string{"op", "cache"}), @@ -71,32 +70,31 @@ func WithNow(now func() time.Time) Option { } } -// WithPerDomainFrequency allows configuring perDomain frequency for the RateLimiter -func WithPerDomainFrequency(d time.Duration) Option { +// WithSourceIPLimitPerSecond allows configuring per source IP limit per second for RateLimiter +func WithSourceIPLimitPerSecond(limit float64) Option { return func(rl *RateLimiter) { - rl.perDomainFrequency = d + rl.sourceIPLimitPerSecond = limit } } -// WithPerDomainBurstSize configures burst per domain for the RateLimiter -func WithPerDomainBurstSize(burst int) Option { +// WithSourceIPBurstSize configures burst per source IP for the RateLimiter +func WithSourceIPBurstSize(burst int) Option { return func(rl *RateLimiter) { - rl.perDomainBurstSize = burst + rl.sourceIPBurstSize = burst } } -func (rl *RateLimiter) getDomainCounter(domain string) *rate.Limiter { - limiterI, _ := rl.domainsCache.FindOrFetch(domain, domain, func() (interface{}, error) { - return rate.NewLimiter(rate.Every(rl.perDomainFrequency), rl.perDomainBurstSize), nil +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 }) return limiterI.(*rate.Limiter) } -// DomainAllowed checks that the requested domain can be accessed within -// the maxCountPerDomain in the given domainWindow. -func (rl *RateLimiter) DomainAllowed(domain string) (res bool) { - limiter := rl.getDomainCounter(domain) +// SourceIPAllowed checks that the real remote IP address is allowed to perform an operation +func (rl *RateLimiter) SourceIPAllowed(sourceIP string) bool { + limiter := rl.getSourceIPLimiter(sourceIP) // AllowN allows us to use the rl.now function, so we can test this more easily. return limiter.AllowN(rl.now(), 1) diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go index 96f50b1b..5b4a52fd 100644 --- a/internal/ratelimiter/ratelimiter_test.go +++ b/internal/ratelimiter/ratelimiter_test.go @@ -21,30 +21,30 @@ func TestDomainAllowed(t *testing.T) { t.Parallel() tcs := map[string]struct { - now string - domainRate time.Duration - perDomainBurstPerSecond int - reqNum int + now string + sourceIPLimit float64 + sourceIPBurstSize int + reqNum int }{ - "one_request_per_nanosecond": { - domainRate: time.Nanosecond, // 1 per nanosecond - perDomainBurstPerSecond: 1, - reqNum: 2, + "one_request_per_second": { + sourceIPLimit: 1, + sourceIPBurstSize: 1, + reqNum: 2, }, - "one_request_per_nanosecond_but_big_bucket": { - domainRate: time.Nanosecond, - perDomainBurstPerSecond: 10, - reqNum: 11, + "one_request_per_second_but_big_bucket": { + sourceIPLimit: 1, + sourceIPBurstSize: 10, + reqNum: 11, }, "three_req_per_second_bucket_size_one": { - domainRate: 3, // 3 per second - perDomainBurstPerSecond: 1, // max burst 1 means 1 at a time - reqNum: 3, + sourceIPLimit: 3, + sourceIPBurstSize: 1, // max burst 1 means 1 at a time + reqNum: 3, }, "10_requests_per_second": { - domainRate: 10, - perDomainBurstPerSecond: 10, - reqNum: 11, + sourceIPLimit: 10, + sourceIPBurstSize: 10, + reqNum: 11, }, } @@ -52,16 +52,16 @@ func TestDomainAllowed(t *testing.T) { t.Run(tn, func(t *testing.T) { rl := New( WithNow(mockNow), - WithPerDomainFrequency(tc.domainRate), - WithPerDomainBurstSize(tc.perDomainBurstPerSecond), + WithSourceIPLimitPerSecond(tc.sourceIPLimit), + WithSourceIPBurstSize(tc.sourceIPBurstSize), ) for i := 0; i < tc.reqNum; i++ { - got := rl.DomainAllowed("rate.gitlab.io") - if i < tc.perDomainBurstPerSecond { + got := rl.SourceIPAllowed("172.16.123.1") + if i < tc.sourceIPBurstSize { require.Truef(t, got, "expected true for request no. %d", i) } else { - // requests should fail after reaching tc.perDomainBurstPerSecond because mockNow + // requests should fail after reaching tc.sourceIPBurstSize because mockNow // always returns the same time require.False(t, got, "expected false for request no. %d", i) } @@ -73,20 +73,20 @@ func TestDomainAllowed(t *testing.T) { func TestSingleRateLimiterWithMultipleDomains(t *testing.T) { rate := 10 * time.Millisecond rl := New( - WithPerDomainFrequency(rate), - WithPerDomainBurstSize(1), + WithSourceIPLimitPerSecond(float64(1/rate)), + WithSourceIPBurstSize(1), ) wg := sync.WaitGroup{} - wg.Add(3) testFn := func(domain string) func(t *testing.T) { return func(t *testing.T) { + wg.Add(1) go func() { defer wg.Done() for i := 0; i < 5; i++ { - got := rl.DomainAllowed(domain) + got := rl.SourceIPAllowed(domain) require.Truef(t, got, "expected true for request no. %d", i) time.Sleep(rate) } @@ -94,13 +94,13 @@ func TestSingleRateLimiterWithMultipleDomains(t *testing.T) { } } - first := "first.gitlab.io" + first := "172.16.123.10" t.Run(first, testFn(first)) - second := "second.gitlab.io" + second := "172.16.123.20" t.Run(second, testFn(second)) - third := "third.gitlab.io" + third := "172.16.123.30" t.Run(third, testFn(third)) wg.Wait() -- cgit v1.2.3 From 25eeea495282065e82d7e72c6d8ffd01a6f79602 Mon Sep 17 00:00:00 2001 From: Jaime Martinez Date: Thu, 30 Sep 2021 13:33:53 +1000 Subject: chore: update test names and default burst --- internal/ratelimiter/ratelimiter.go | 4 ++-- internal/ratelimiter/ratelimiter_test.go | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/ratelimiter/ratelimiter.go b/internal/ratelimiter/ratelimiter.go index 0753e187..e1cf076d 100644 --- a/internal/ratelimiter/ratelimiter.go +++ b/internal/ratelimiter/ratelimiter.go @@ -15,8 +15,8 @@ const ( // The default value is 20 requests per second. DefaultSourceIPLimitPerSecond = 20.0 // DefaultSourceIPBurstSize is the maximum burst allowed per rate limiter. - // E.g. The first 20 requests within 1s will succeed, but the 21st will fail. - DefaultSourceIPBurstSize = 20 + // 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 diff --git a/internal/ratelimiter/ratelimiter_test.go b/internal/ratelimiter/ratelimiter_test.go index 5b4a52fd..cdf12fe6 100644 --- a/internal/ratelimiter/ratelimiter_test.go +++ b/internal/ratelimiter/ratelimiter_test.go @@ -17,7 +17,7 @@ func mockNow() time.Time { return validTime } -func TestDomainAllowed(t *testing.T) { +func TestSourceIPAllowed(t *testing.T) { t.Parallel() tcs := map[string]struct { @@ -70,7 +70,7 @@ func TestDomainAllowed(t *testing.T) { } } -func TestSingleRateLimiterWithMultipleDomains(t *testing.T) { +func TestSingleRateLimiterWithMultipleSourceIPs(t *testing.T) { rate := 10 * time.Millisecond rl := New( WithSourceIPLimitPerSecond(float64(1/rate)), -- cgit v1.2.3