Sweep expired rate limiter buckets

pkg/api/server.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
@@ -87,9 +88,13 @@ func (s *Server) startGRPC() error {
 	stream = append(stream, telemetryInterceptor.Stream())
 
 	if s.Config.Authn.Enable {
+		limiter := ratelimiter.NewTokenBucketRateLimiter(s.Log)
+		// Expire buckets after 6 hours of inactivity,
+		// sweep for expired buckets every hour.
+		go limiter.Janitor(s.ctx, time.Hour, 6*time.Hour)
 		s.authorizer = NewWalletAuthorizer(&AuthnConfig{
 			AuthnOptions: s.Config.Authn,
-			Limiter:      ratelimiter.NewTokenBucketRateLimiter(s.Log),
+			Limiter:      limiter,
 			AllowLister:  s.Config.AllowLister,
 			Log:          s.Log.Named("authn"),
 		})

pkg/ratelimiter/buckets.go

@@ -0,0 +1,80 @@
+package ratelimiter
+
+import (
+	"sync"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+type Buckets struct {
+	log     *zap.Logger
+	mutex   sync.RWMutex
+	buckets map[string]*Entry
+}
+
+func NewBuckets(log *zap.Logger) *Buckets {
+	return &Buckets{
+		log:     log,
+		buckets: make(map[string]*Entry),
+		mutex:   sync.RWMutex{},
+	}
+}
+
+func (b *Buckets) getAndRefill(bucket string, limit *Limit, multiplier uint16, create bool) *Entry {
+	// The locking strategy is adapted from the following blog post: https://misfra.me/optimizing-concurrent-map-access-in-go/
+	b.mutex.RLock()
+	currentVal, exists := b.buckets[bucket]
+	b.mutex.RUnlock()
+	if !exists {
+		if !create {
+			return nil
+		}
+		b.mutex.Lock()
+		currentVal, exists = b.buckets[bucket]
+		if !exists {
+			currentVal = &Entry{
+				tokens:   uint16(limit.MaxTokens * multiplier),
+				lastSeen: time.Now(),
+				mutex:    sync.Mutex{},
+			}
+			b.buckets[bucket] = currentVal
+			b.mutex.Unlock()
+
+			return currentVal
+		}
+	}
+
+	limit.Refill(currentVal, multiplier)
+	return currentVal
+}
+
+func (b *Buckets) deleteExpired(expiresAfter time.Duration) {
+	// Use RLock to iterate over the map
+	// to allow concurrent reads
+	b.mutex.RLock()
+	var expired []string
+	for bucket, entry := range b.buckets {
+		if time.Since(entry.lastSeen) > expiresAfter {
+			expired = append(expired, bucket)
+		}
+	}
+	b.mutex.RUnlock()
+	if len(expired) == 0 {
+		return
+	}
+	b.log.Info("found expired buckets", zap.Int("count", len(expired)))
+	// Use Lock for individual deletes to avoid prolonged
+	// lockout for readers.
+	count := 0
+	for _, bucket := range expired {
+		b.mutex.Lock()
+		// check lastSeen again in case it was updated in the meantime.
+		if entry, exists := b.buckets[bucket]; exists && time.Since(entry.lastSeen) > expiresAfter {
+			delete(b.buckets, bucket)
+			count++
+		}
+		b.mutex.Unlock()
+	}
+	b.log.Info("deleted expired buckets", zap.Int("count", count))
+}

pkg/ratelimiter/rate_limiter.go

@@ -1,6 +1,7 @@
 package ratelimiter
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"time"
@@ -68,9 +69,11 @@ func (l Limit) Refill(entry *Entry, multiplier uint16) {
 
 // TokenBucketRateLimiter implements the RateLimiter interface
 type TokenBucketRateLimiter struct {
-	log                *zap.Logger
-	buckets            map[string]*Entry
-	mutex              sync.RWMutex
+	log *zap.Logger
+	// buckets that can be added to
+	buckets1 *Buckets
+	// buckets being GCed
+	buckets2           *Buckets
 	PriorityMultiplier uint16
 	Limits             map[LimitType]*Limit
 }
@@ -79,8 +82,8 @@ func NewTokenBucketRateLimiter(log *zap.Logger) *TokenBucketRateLimiter {
 	tb := new(TokenBucketRateLimiter)
 	tb.log = log.Named("ratelimiter")
 	// TODO: need to periodically clear out expired items to avoid unlimited growth of the map.
-	tb.buckets = make(map[string]*Entry)
-	tb.mutex = sync.RWMutex{}
+	tb.buckets1 = NewBuckets(log.Named("buckets1"))
+	tb.buckets2 = NewBuckets(log.Named("buckets2"))
 	tb.PriorityMultiplier = PRIORITY_MULTIPLIER
 	tb.Limits = map[LimitType]*Limit{
 		DEFAULT: {DEFAULT_MAX_TOKENS, DEFAULT_RATE_PER_MINUTE},
@@ -103,25 +106,10 @@ func (rl *TokenBucketRateLimiter) fillAndReturnEntry(limitType LimitType, bucket
 	if isPriority {
 		multiplier = rl.PriorityMultiplier
 	}
-	// The locking strategy is adapted from the following blog post: https://misfra.me/optimizing-concurrent-map-access-in-go/
-	rl.mutex.RLock()
-	currentVal, exists := rl.buckets[bucket]
-	rl.mutex.RUnlock()
-	if !exists {
-		rl.mutex.Lock()
-		currentVal = &Entry{
-			tokens:   uint16(limit.MaxTokens * multiplier),
-			lastSeen: time.Now(),
-			mutex:    sync.Mutex{},
-		}
-		rl.buckets[bucket] = currentVal
-		rl.mutex.Unlock()
-
-		return currentVal
+	if entry := rl.buckets2.getAndRefill(bucket, limit, multiplier, false); entry != nil {
+		return entry
 	}
-
-	limit.Refill(currentVal, multiplier)
-	return currentVal
+	return rl.buckets1.getAndRefill(bucket, limit, multiplier, true)
 }
 
 // The Spend function takes a bucket and a boolean asserting whether to apply the PRIORITY or the REGULAR rate limits.
@@ -140,6 +128,21 @@ func (rl *TokenBucketRateLimiter) Spend(limitType LimitType, bucket string, cost
 	return nil
 }
 
+func (rl *TokenBucketRateLimiter) Janitor(ctx context.Context, sweepInterval, expiresAfter time.Duration) {
+	ticker := time.NewTicker(sweepInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			rl.buckets2.deleteExpired(expiresAfter)
+			rl.buckets1, rl.buckets2 = rl.buckets2, rl.buckets1
+			rl.log.Info("swapped buckets")
+		}
+	}
+}
+
 func minUint16(x, y uint16) uint16 {
 	if x <= y {
 		return x

pkg/ratelimiter/rate_limiter_test.go

@@ -1,6 +1,7 @@
 package ratelimiter
 
 import (
+	"context"
 	"sync"
 	"testing"
 	"time"
@@ -14,11 +15,7 @@ const walletAddress = "0x1234"
 func TestSpend(t *testing.T) {
 	logger, _ := zap.NewDevelopment()
 	rl := NewTokenBucketRateLimiter(logger)
-	rl.buckets[walletAddress] = &Entry{
-		lastSeen: time.Now(),
-		tokens:   uint16(1),
-		mutex:    sync.Mutex{},
-	}
+	rl.buckets1.getAndRefill(walletAddress, &Limit{1, 0}, 1, true)
 
 	err1 := rl.Spend(DEFAULT, walletAddress, 1, false)
 	require.NoError(t, err1)
@@ -42,12 +39,9 @@ func TestSpendWithTime(t *testing.T) {
 	logger, _ := zap.NewDevelopment()
 	rl := NewTokenBucketRateLimiter(logger)
 	rl.Limits[DEFAULT] = &Limit{100, 1}
-	rl.buckets[walletAddress] = &Entry{
-		// Set the last seen to 1 minute ago
-		lastSeen: time.Now().Add(-1 * time.Minute),
-		tokens:   uint16(0),
-		mutex:    sync.Mutex{},
-	}
+	entry := rl.buckets1.getAndRefill(walletAddress, &Limit{0, 0}, 1, true)
+	// Set the last seen to 1 minute ago
+	entry.lastSeen = time.Now().Add(-1 * time.Minute)
 	err1 := rl.Spend(DEFAULT, walletAddress, 1, false)
 	require.NoError(t, err1)
 	err2 := rl.Spend(DEFAULT, walletAddress, 1, false)
@@ -58,41 +52,31 @@ func TestSpendWithTime(t *testing.T) {
 func TestSpendMaxBucket(t *testing.T) {
 	logger, _ := zap.NewDevelopment()
 	rl := NewTokenBucketRateLimiter(logger)
-	rl.buckets[walletAddress] = &Entry{
-		// Set last seen to 500 minutes ago
-		lastSeen: time.Now().Add(-500 * time.Minute),
-		tokens:   uint16(0),
-		mutex:    sync.Mutex{},
-	}
-	entry := rl.fillAndReturnEntry(DEFAULT, walletAddress, false)
+	entry := rl.buckets1.getAndRefill(walletAddress, &Limit{0, 0}, 1, true)
+	// Set last seen to 500 minutes ago
+	entry.lastSeen = time.Now().Add(-500 * time.Minute)
+	entry = rl.fillAndReturnEntry(DEFAULT, walletAddress, false)
 	require.Equal(t, entry.tokens, DEFAULT_MAX_TOKENS)
 }
 
 // Ensure that the allow list is being correctly applied
 func TestSpendAllowListed(t *testing.T) {
 	logger, _ := zap.NewDevelopment()
 	rl := NewTokenBucketRateLimiter(logger)
-	rl.buckets[walletAddress] = &Entry{
-		// Set last seen to 5 minutes ago
-		lastSeen: time.Now().Add(-5 * time.Minute),
-		tokens:   uint16(0),
-		mutex:    sync.Mutex{},
-	}
-	entry := rl.fillAndReturnEntry(DEFAULT, walletAddress, true)
+	entry := rl.buckets1.getAndRefill(walletAddress, &Limit{0, 0}, 1, true)
+	// Set last seen to 5 minutes ago
+	entry.lastSeen = time.Now().Add(-5 * time.Minute)
+	entry = rl.fillAndReturnEntry(DEFAULT, walletAddress, true)
 	require.Equal(t, entry.tokens, uint16(5*DEFAULT_RATE_PER_MINUTE*PRIORITY_MULTIPLIER))
 }
 
 func TestMaxUint16(t *testing.T) {
 	logger, _ := zap.NewDevelopment()
 	rl := NewTokenBucketRateLimiter(logger)
-	rl.buckets[walletAddress] = &Entry{
-		// Set last seen to 1 million minutes ago
-		lastSeen: time.Now().Add(-1000000 * time.Minute),
-		tokens:   uint16(0),
-		mutex:    sync.Mutex{},
-	}
-
-	entry := rl.fillAndReturnEntry(DEFAULT, walletAddress, true)
+	entry := rl.buckets1.getAndRefill(walletAddress, &Limit{0, 0}, 1, true)
+	// Set last seen to 1 million minutes ago
+	entry.lastSeen = time.Now().Add(-1000000 * time.Minute)
+	entry = rl.fillAndReturnEntry(DEFAULT, walletAddress, true)
 	require.Equal(t, entry.tokens, DEFAULT_MAX_TOKENS*PRIORITY_MULTIPLIER)
 }
 
@@ -114,3 +98,25 @@ func TestSpendConcurrent(t *testing.T) {
 	entry := rl.fillAndReturnEntry(PUBLISH, walletAddress, false)
 	require.Equal(t, entry.tokens, uint16(0))
 }
+
+func TestBucketExpiration(t *testing.T) {
+	logger, _ := zap.NewDevelopment()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	rl := NewTokenBucketRateLimiter(logger)
+	rl.Limits[DEFAULT] = &Limit{2, 0}
+	// expire entries after 100ms, sweep every 50ms
+	go rl.Janitor(ctx, time.Millisecond*50, time.Millisecond*100)
+	require.NoError(t, rl.Spend(DEFAULT, "ip1", 1, false))
+	require.NoError(t, rl.Spend(DEFAULT, "ip2", 1, false))
+	time.Sleep(50 * time.Millisecond)
+	require.NoError(t, rl.Spend(DEFAULT, "ip2", 1, false))
+	time.Sleep(50 * time.Millisecond)
+	require.Error(t, rl.Spend(DEFAULT, "ip2", 1, false))
+	time.Sleep(50 * time.Millisecond)
+	// ip2 has been refreshed every 50ms so it should still be out of tokens
+	require.Error(t, rl.Spend(DEFAULT, "ip2", 1, false))
+	// ip1 entry should have expired by now, so we should have 2 tokens again
+	require.NoError(t, rl.Spend(DEFAULT, "ip1", 1, false))
+	require.NoError(t, rl.Spend(DEFAULT, "ip1", 1, false))
+}