Skip to content

Commit 0fa6664

Browse files
committed
Implement token bucket algo
1 parent 3255704 commit 0fa6664

File tree

1 file changed

+47
-1
lines changed

1 file changed

+47
-1
lines changed

internal/algorithm/token_bucket.go

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,47 @@
1-
// internal/algorithm/token_bucket.go
1+
package algorithm
2+
3+
import (
4+
"sync"
5+
"time"
6+
)
7+
8+
// TokenBucket kept alongside SlidingWindow to highlight the tradeoff discussion in docs/tradeoffs.md.
9+
type TokenBucket struct {
10+
mu sync.Mutex
11+
tokens float64
12+
capacity float64
13+
rate float64 // tokens per second
14+
lastFillTime time.Time
15+
}
16+
17+
func NewTokenBucet(capacity float64, ratePerSec float64) *TokenBucket {
18+
return &TokenBucket{
19+
tokens: capacity,
20+
capacity: capacity,
21+
rate: ratePerSec,
22+
lastFillTime: time.Now(),
23+
}
24+
}
25+
26+
func (tb *TokenBucket) Allow() bool {
27+
tb.mu.Lock()
28+
defer tb.mu.Unlock()
29+
30+
now := time.Now()
31+
elapsed := now.Sub(tb.lastFillTime).Seconds()
32+
tb.tokens = minF(tb.capacity, tb.tokens+elapsed*tb.rate)
33+
tb.lastFill = now
34+
35+
if tb.tokens < 1 {
36+
return false
37+
}
38+
tb.tokens--
39+
return true
40+
}
41+
42+
func minF(a, b float64) float64 {
43+
if a < b {
44+
return a
45+
}
46+
return b
47+
}

0 commit comments

Comments
 (0)