Merge pull request #3 from augustus281/DEV

[algo]: leaky bucket

Author: augustus281

go.mod

@@ -1,3 +1,8 @@
 module github.com/augustus281/ratelimiter.git
 
 go 1.22.2
+
+require (
+	github.com/edwingeng/deque v1.0.3 // indirect
+	github.com/gammazero/deque v0.2.1 // indirect
+)

go.sum

@@ -0,0 +1,4 @@
+github.com/edwingeng/deque v1.0.3 h1:gx+5OnQK8qMcYNUxcD/M76BT/LujVNVAZEP/MGF8WNw=
+github.com/edwingeng/deque v1.0.3/go.mod h1:3Ys1pJhyVaB6iWigv4o2r6Ug1GZmfDWqvqmO6bjojg0=
+github.com/gammazero/deque v0.2.1 h1:qSdsbG6pgp6nL7A0+K/B7s12mcCY/5l5SIUpMOl+dC0=
+github.com/gammazero/deque v0.2.1/go.mod h1:LFroj8x4cMYCukHJDbxFCkT+r9AndaJnFMuZDV34tuU=

leakybucket.go

@@ -0,0 +1,45 @@
+package ratelimiter
+
+import (
+	"sync"
+	"time"
+)
+
+type LeakyBucket struct {
+	capacity int         // maximum number of requests in the bucket
+	leakRate int         // rate at which requests leak (requests/second)
+	bucket   []time.Time // hold request timestamps
+	lastLeak time.Time   // last time we leaked from the bucket
+	sync.Mutex
+}
+
+func NewLeakyBucket(capacity, leakRate int) *LeakyBucket {
+	return &LeakyBucket{
+		capacity: capacity,
+		leakRate: leakRate,
+		bucket:   []time.Time{},
+		lastLeak: time.Now(),
+	}
+}
+
+func (lb *LeakyBucket) AllowRequest() bool {
+	now := time.Now()
+	leakedTime := now.Sub(lb.lastLeak).Seconds()
+
+	leaked := int(leakedTime) * lb.leakRate
+	if leaked > 0 {
+		if leaked > len(lb.bucket) {
+			lb.bucket = []time.Time{}
+		} else {
+			lb.bucket = lb.bucket[leaked:]
+		}
+		lb.lastLeak = now
+	}
+
+	if len(lb.bucket) < lb.capacity {
+		lb.bucket = append(lb.bucket, now)
+		return true
+	}
+
+	return false
+}

leakybucket_test.go

@@ -0,0 +1,60 @@
+package ratelimiter
+
+import (
+	"testing"
+	"time"
+)
+
+func TestLeakyBucket(t *testing.T) {
+	tests := []struct {
+		name           string
+		capacity       int
+		leakRate       int
+		initRequests   int
+		sleepDuration  time.Duration
+		expectedResult bool
+	}{
+		{
+			name:           "Allow request when bucket is not full",
+			capacity:       5,
+			leakRate:       1,
+			initRequests:   4,
+			sleepDuration:  0,
+			expectedResult: true,
+		},
+		{
+			name:           "Deny request when bucket is full",
+			capacity:       5,
+			leakRate:       1,
+			initRequests:   5,
+			sleepDuration:  0,
+			expectedResult: false,
+		},
+		{
+			name:           "Allow request after bucket leaks",
+			capacity:       5,
+			leakRate:       1,
+			initRequests:   5,
+			sleepDuration:  2 * time.Second,
+			expectedResult: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			lb := NewLeakyBucket(tt.capacity, tt.leakRate)
+
+			for i := 0; i < tt.initRequests; i++ {
+				lb.AllowRequest()
+			}
+
+			if tt.sleepDuration > 0 {
+				time.Sleep(tt.sleepDuration)
+			}
+
+			if lb.AllowRequest() != tt.expectedResult {
+				t.Errorf("AllowRequest() = %v, want %v", lb.AllowRequest(), tt.expectedResult)
+			}
+		})
+	}
+}

tokenbucket.go

@@ -1,4 +1,4 @@
-package tokenbucket
+package ratelimiter
 
 import (
 	"sync"

tokenbucket_test.go

@@ -1,4 +1,4 @@
-package tokenbucket
+package ratelimiter
 
 import (
 	"sync"