Adding leaky bucket and adding support to both limiters in the registry

bucket/leaky.go

@@ -0,0 +1,54 @@
+package bucket
+
+import (
+	"sync"
+	"time"
+)
+
+// LeakyLimiter implements a leaky bucket rate limiter. Requests fill the bucket,
+// which drains at a constant rate. If the bucket is full, requests are rejected.
+// Unlike TokenLimiter, this enforces a smooth output rate rather than allowing bursts.
+type LeakyLimiter struct {
+	mu sync.Mutex
+
+	capacity, level, rate float64
+	lastUpdatedAt         time.Time
+}
+
+// NewLeakyLimiter creates a new leaky bucket limiter.
+// Capacity is the maximum bucket size. Rate is how many requests drain per second.
+func NewLeakyLimiter(capacity, rate uint32) *LeakyLimiter {
+	return &LeakyLimiter{
+		capacity:      float64(capacity),
+		rate:          float64(rate),
+		lastUpdatedAt: time.Now(),
+	}
+}
+
+func (lim *LeakyLimiter) update() {
+	t := time.Now()
+	if t.Before(lim.lastUpdatedAt) {
+		return
+	}
+
+	lim.level = max(0, lim.level-t.Sub(lim.lastUpdatedAt).Seconds()*lim.rate)
+	lim.lastUpdatedAt = t
+}
+
+// Allow reports whether a request is allowed. It adds one to the bucket level
+// if there is room and returns true. If the bucket is full, it returns false
+// without blocking.
+func (lim *LeakyLimiter) Allow() bool {
+	lim.mu.Lock()
+	defer lim.mu.Unlock()
+
+	lim.update()
+
+	if lim.level+1 <= lim.capacity {
+		lim.level++
+
+		return true
+	}
+
+	return false
+}

bucket/leaky_test.go

@@ -0,0 +1,72 @@
+package bucket_test
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/serroba/rate/bucket"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLeakyLimiter_Allow(t *testing.T) {
+	type args struct {
+		capacity uint32
+		rate     uint32
+	}
+
+	tests := []struct {
+		name             string
+		args             args
+		previousAttempts int
+		want             bool
+	}{
+		{name: "Test With No Capacity", args: args{capacity: 0, rate: 0}, want: false},
+		{name: "Test With Capacity 1", args: args{capacity: 1, rate: 0}, want: true},
+		{
+			name:             "Test With Capacity 1 with previous attempt",
+			args:             args{capacity: 1, rate: 0},
+			previousAttempts: 1,
+			want:             false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			lim := bucket.NewLeakyLimiter(tt.args.capacity, tt.args.rate)
+			for range tt.previousAttempts {
+				lim.Allow()
+			}
+
+			assert.Equalf(t, tt.want, lim.Allow(), "Allow()")
+		})
+	}
+}
+
+func TestLeakyLimiter_Allow_Concurrent(t *testing.T) {
+	var (
+		allow atomic.Int64
+		deny  atomic.Int64
+		wg    sync.WaitGroup
+	)
+
+	lim := bucket.NewLeakyLimiter(10, 5)
+
+	for range 15 {
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			if lim.Allow() {
+				allow.Add(1)
+			} else {
+				deny.Add(1)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	assert.Equal(t, int64(10), allow.Load())
+	assert.Equal(t, int64(5), deny.Load())
+}

bucket/registry.go

@@ -7,24 +7,28 @@ import (
 type (
 	Identifier string
 	Registry   struct {
-		mu             sync.Mutex
-		limiters       map[Identifier]*TokenLimiter
-		capacity, rate uint32
+		mu       sync.Mutex
+		factory  LimiterFactory
+		limiters map[Identifier]Limiter
 	}
 )
 
-func NewRegistry(capacity, rate uint32, users ...Identifier) (*Registry, error) {
-	limiters := make(map[Identifier]*TokenLimiter)
+type Limiter interface {
+	Allow() bool
+}
+
+type LimiterFactory func() Limiter
+
+func NewRegistry(factory LimiterFactory, keys ...Identifier) (*Registry, error) {
+	limiters := make(map[Identifier]Limiter)
 
-	for _, user := range users {
-		limiter := NewLimiter(capacity, rate)
-		limiters[user] = limiter
+	for _, key := range keys {
+		limiters[key] = factory()
 	}
 
 	return &Registry{
 		limiters: limiters,
-		capacity: capacity,
-		rate:     rate,
+		factory:  factory,
 	}, nil
 }
 
@@ -34,7 +38,7 @@ func (r *Registry) Allow(key Identifier) bool {
 
 	lim, ok := r.limiters[key]
 	if !ok {
-		lim = NewLimiter(r.capacity, r.rate)
+		lim = r.factory()
 		r.limiters[key] = lim
 	}