From a8fcc7dfef823f58ae51744858316813a0b37959 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=A1clav=20Pavl=C3=ADn?= Date: Thu, 30 May 2024 06:59:18 +0200 Subject: [PATCH] use shrinking map in rate limiter --- go.mod | 13 ++++++++++++- go.sum | 11 +++++------ telemetry/ratelimiter.go | 30 ++++++++++++++++++++---------- telemetry/ratelimiter_test.go | 29 ++++++++++++++++++++++------- telemetry/server.go | 2 +- 5 files changed, 60 insertions(+), 25 deletions(-) diff --git a/go.mod b/go.mod index 8b80665..ea081a9 100644 --- a/go.mod +++ b/go.mod @@ -1,8 +1,9 @@ module github.com/status-im/dev-telemetry -go 1.15 +go 1.20 require ( + github.com/go-auxiliaries/shrinking-map v0.3.0 github.com/golang-migrate/migrate/v4 v4.15.2 github.com/gorilla/mux v1.8.0 // indirect github.com/lib/pq v1.10.3 @@ -11,3 +12,13 @@ require ( go.uber.org/zap v1.27.0 golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 ) + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/hashicorp/errwrap v1.1.0 // indirect + github.com/hashicorp/go-multierror v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + go.uber.org/atomic v1.7.0 // indirect + go.uber.org/multierr v1.6.0 // indirect + gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect +) diff --git a/go.sum b/go.sum index a9337e8..ae0793d 100644 --- a/go.sum +++ b/go.sum @@ -330,7 +330,6 @@ github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:ma github.com/cpuguy83/go-md2man/v2 v2.0.0/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7DoTY= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= -github.com/creack/pty v1.1.11 h1:07n33Z8lZxZ2qwegKbObQohDhXDQxiMMz1NOUGYlesw= github.com/creack/pty v1.1.11/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/cyphar/filepath-securejoin v0.2.2/go.mod h1:FpkQEhXnPnOthhzymB7CGsFk2G9VLXONKD9G7QGMM+4= github.com/cyphar/filepath-securejoin v0.2.3/go.mod h1:aPGpWjXOXUn2NCNjFvBE6aRxGGx79pTxQpKOJNYHHl4= @@ -408,6 +407,8 @@ github.com/garyburd/redigo v0.0.0-20150301180006-535138d7bcd7/go.mod h1:NR3MbYis github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= github.com/ghodss/yaml v0.0.0-20150909031657-73d445a93680/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04= github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04= +github.com/go-auxiliaries/shrinking-map v0.3.0 h1:kXiLmFY4y2s35WtOYAb02LRZ92IRnfzio+3prZn6ULs= +github.com/go-auxiliaries/shrinking-map v0.3.0/go.mod h1:UtBmTTKuUfI8wkhzaZ7G/xgHjxGxLwM2a6kf+aWmSmc= github.com/go-fonts/dejavu v0.1.0/go.mod h1:4Wt4I4OU2Nq9asgDCteaAaWZOV24E+0/Pwo0gppep4g= github.com/go-fonts/latin-modern v0.2.0/go.mod h1:rQVLdDMK+mK1xscDwsqM5J8U2jrRa3T0ecnM9pNujks= github.com/go-fonts/liberation v0.1.1/go.mod h1:K6qoJYypsmfVjWg8KOVDQhLc8UDgIK2HYqyqAO9z7GY= @@ -552,7 +553,6 @@ github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.6 h1:BKbKCqvP6I+rmFHt06ZmyQtvB8xAkWdhFyr0ZUNZcxQ= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-containerregistry v0.5.1/go.mod h1:Ct15B4yir3PLOP5jsy0GNeYVaIZs/MK/Jz5any1wFW0= github.com/google/go-github/v39 v39.2.0/go.mod h1:C1s8C5aCC9L+JXIYpJM5GYytdX52vC1bLvHEF1IhBrE= @@ -1128,12 +1128,14 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0= go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU= +go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4= go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU= go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ= go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA= go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= +go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U= go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo= go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8= go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= @@ -1439,7 +1441,6 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= @@ -1534,7 +1535,6 @@ golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gonum.org/v1/gonum v0.0.0-20180816165407-929014505bf4/go.mod h1:Y+Yx5eoAFn32cQvJDxZx5Dpnq+c3wtXuadVZAcxbbBo= gonum.org/v1/gonum v0.8.2/go.mod h1:oe/vMfY3deqTw+1EZJhuvEW2iwGF1bW9wwu7XCu0+v0= @@ -1735,6 +1735,7 @@ gopkg.in/yaml.v2 v2.2.5/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= @@ -1744,11 +1745,9 @@ gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gorm.io/driver/postgres v1.0.8/go.mod h1:4eOzrI1MUfm6ObJU/UcmbXyiHSs8jSwH95G5P5dxcAg= gorm.io/gorm v1.20.12/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw= gorm.io/gorm v1.21.4/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw= -gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo= gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw= gotest.tools/v3 v3.0.2/go.mod h1:3SzNCllyD9/Y+b5r9JIKQ474KzkZyqLqEfYqMsX94Bk= gotest.tools/v3 v3.0.3/go.mod h1:Z7Lb0S5l+klDB31fvDQX8ss/FlKDxtlFlw3Oa8Ymbl8= -gotest.tools/v3 v3.1.0 h1:rVV8Tcg/8jHUkPUorwjaMTtemIMVXfIPKiOqnhEhakk= gotest.tools/v3 v3.1.0/go.mod h1:fHy7eyTmJFO5bQbUsEGQ1v4m2J3Jz9eWL54TP2/ZuYQ= honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/telemetry/ratelimiter.go b/telemetry/ratelimiter.go index 404c8e8..363e410 100644 --- a/telemetry/ratelimiter.go +++ b/telemetry/ratelimiter.go @@ -5,6 +5,8 @@ import ( "sync" "time" + shrinkingmap "github.com/go-auxiliaries/shrinking-map/pkg/shrinking-map" + "go.uber.org/zap" "golang.org/x/time/rate" ) @@ -16,22 +18,24 @@ type Limiter struct { } type RateLimiter struct { - limiters map[string]*Limiter + limiters *shrinkingmap.Map[string, *Limiter] //map[string]*Limiter lock *sync.RWMutex r rate.Limit b int + logger *zap.Logger } -func NewRateLimiter(ctx context.Context, r rate.Limit, b int) *RateLimiter { - return NewRateLimiterWithCleanup(ctx, r, b, DEFAULT_CLEANUP_TIME) +func NewRateLimiter(ctx context.Context, r rate.Limit, b int, logger *zap.Logger) *RateLimiter { + return NewRateLimiterWithCleanup(ctx, r, b, DEFAULT_CLEANUP_TIME, logger) } -func NewRateLimiterWithCleanup(ctx context.Context, r rate.Limit, b int, cleanupTime time.Duration) *RateLimiter { +func NewRateLimiterWithCleanup(ctx context.Context, r rate.Limit, b int, cleanupTime time.Duration, logger *zap.Logger) *RateLimiter { rl := &RateLimiter{ - limiters: make(map[string]*Limiter), + limiters: shrinkingmap.New[string, *Limiter](200), lock: &sync.RWMutex{}, r: r, b: b, + logger: logger, } go rl.cleanup(ctx, cleanupTime) @@ -42,7 +46,7 @@ func NewRateLimiterWithCleanup(ctx context.Context, r rate.Limit, b int, cleanup func (rl *RateLimiter) GetLimiter(ip string) *rate.Limiter { rl.lock.Lock() - limiter, ok := rl.limiters[ip] + limiter, ok := rl.limiters.Get2(ip) if !ok { rl.lock.Unlock() return rl.AddIP(ip) @@ -59,7 +63,7 @@ func (rl *RateLimiter) AddIP(ip string) *rate.Limiter { defer rl.lock.Unlock() limiter := rate.NewLimiter(rl.r, rl.b) - rl.limiters[ip] = &Limiter{limiter: limiter, lastUsed: time.Now()} + rl.limiters.Set(ip, &Limiter{limiter: limiter, lastUsed: time.Now()}) return limiter } @@ -68,7 +72,11 @@ func (rl *RateLimiter) RemoveIP(ip string) { rl.lock.Lock() defer rl.lock.Unlock() - delete(rl.limiters, ip) + rl.limiters.Delete(ip) +} + +func (rl *RateLimiter) NumClients() int { + return len(rl.limiters.Values()) } func (rl *RateLimiter) cleanup(ctx context.Context, cleanupEvery time.Duration) { @@ -80,12 +88,14 @@ func (rl *RateLimiter) cleanup(ctx context.Context, cleanupEvery time.Duration) case <-ctx.Done(): return case now := <-t.C: - for ip := range rl.limiters { - limiter := rl.limiters[ip] + numCleaned := 0 + for ip, limiter := range rl.limiters.Values() { if limiter.lastUsed.Add(2 * time.Second).Before(now) { rl.RemoveIP(ip) + numCleaned++ } } + rl.logger.Debug("cleanup", zap.Int("removed", numCleaned)) } } } diff --git a/telemetry/ratelimiter_test.go b/telemetry/ratelimiter_test.go index 9176ffc..00c81d7 100644 --- a/telemetry/ratelimiter_test.go +++ b/telemetry/ratelimiter_test.go @@ -2,17 +2,22 @@ package telemetry import ( "context" + "fmt" "testing" "time" "github.com/stretchr/testify/require" + "go.uber.org/zap" "golang.org/x/time/rate" ) func TestRateLimit(t *testing.T) { ctx := context.Background() defer ctx.Done() - rl := NewRateLimiter(ctx, 1, 1) + + logger, err := zap.NewDevelopment() + require.NoError(t, err) + rl := NewRateLimiter(ctx, 1, 1, logger) ip1 := "1.1.1.1" @@ -37,16 +42,26 @@ func TestRateLimit(t *testing.T) { func TestRateLimitCleanup(t *testing.T) { ctx := context.Background() defer ctx.Done() - rl := NewRateLimiterWithCleanup(ctx, rate.Limit(1/time.Hour), 1, 100*time.Millisecond) - ip1 := "1.1.1.1" + logger, err := zap.NewDevelopment() + require.NoError(t, err) + rl := NewRateLimiterWithCleanup(ctx, rate.Limit(1/time.Hour), 1, 100*time.Millisecond, logger) - limiter := rl.GetLimiter(ip1) - require.True(t, limiter.Allow()) - require.False(t, limiter.Allow()) + for i := 0; i < 300; i++ { + ip := fmt.Sprintf("%d.%d.%d.%d", i, i, i, i) + limiter := rl.GetLimiter(ip) + require.True(t, limiter.Allow()) + require.False(t, limiter.Allow()) + time.Sleep(1 * time.Millisecond) + } + + require.Equal(t, 300, rl.NumClients()) time.Sleep(3 * time.Second) - limiter2 := rl.GetLimiter(ip1) + limiter2 := rl.GetLimiter("1.1.1.1") require.True(t, limiter2.Allow()) + + require.Equal(t, 1, rl.NumClients()) + } diff --git a/telemetry/server.go b/telemetry/server.go index 69f2804..7438ee2 100644 --- a/telemetry/server.go +++ b/telemetry/server.go @@ -35,7 +35,7 @@ func NewServer(db *sql.DB, logger *zap.Logger) *Server { Router: mux.NewRouter().StrictSlash(true), DB: db, logger: logger, - rateLimiter: *NewRateLimiter(ctx, RATE_LIMIT, BURST), + rateLimiter: *NewRateLimiter(ctx, RATE_LIMIT, BURST, logger), ctx: ctx, }