use shrinking map in rate limiter

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
+)

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=

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))
 		}
 	}
 }

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())
+
 }

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,
 	}