mirror of
https://github.com/logos-messaging/telemetry.git
synced 2026-01-09 09:33:11 +00:00
use shrinking map in rate limiter
This commit is contained in:
Václav Pavlín
parent
4b244c6ced
commit
296d312a57
13
go.mod
13
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
|
||||
)
|
||||
|
||||
11
go.sum
11
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=
|
||||
|
||||
@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -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())
|
||||
|
||||
}
|
||||
|
||||
@ -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,
|
||||
}
|
||||
|
||||
@ -224,12 +224,10 @@ func (s *Server) rateLimit(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
limiter := s.rateLimiter.GetLimiter(r.RemoteAddr)
|
||||
// Do stuff here
|
||||
if !limiter.Allow() {
|
||||
http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
// Call the next handler, which can be another middleware in the chain, or the final handler.
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user