63 lines
1.0 KiB
Go
63 lines
1.0 KiB
Go
|
package ratelimit
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type TokenBucket struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
capacity int
|
||
|
|
tokens int
|
||
|
|
refillRate float64 // tokens per second
|
||
|
|
lastRefill time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewTokenBucket(capacity int, refillPerSecond float64) *TokenBucket {
|
||
|
|
return &TokenBucket{
|
||
|
|
capacity: capacity,
|
||
|
|
tokens: capacity,
|
||
|
|
refillRate: refillPerSecond,
|
||
|
|
lastRefill: time.Now(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (tb *TokenBucket) Allow() bool {
|
||
|
|
return tb.AllowN(1)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (tb *TokenBucket) AllowN(n int) bool {
|
||
|
|
tb.mu.Lock()
|
||
|
|
defer tb.mu.Unlock()
|
||
|
|
tb.refill()
|
||
|
|
if tb.tokens >= n {
|
||
|
|
tb.tokens -= n
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
func (tb *TokenBucket) Remaining() int {
|
||
|
|
tb.mu.Lock()
|
||
|
|
defer tb.mu.Unlock()
|
||
|
|
tb.refill()
|
||
|
|
return tb.tokens
|
||
|
|
}
|
||
|
|
|
||
|
|
func (tb *TokenBucket) Capacity() int {
|
||
|
|
return tb.capacity
|
||
|
|
}
|
||
|
|
|
||
|
|
func (tb *TokenBucket) refill() {
|
||
|
|
now := time.Now()
|
||
|
|
elapsed := now.Sub(tb.lastRefill).Seconds()
|
||
|
|
tb.lastRefill = now
|
||
|
|
add := int(elapsed * tb.refillRate)
|
||
|
|
if add > 0 {
|
||
|
|
tb.tokens += add
|
||
|
|
if tb.tokens > tb.capacity {
|
||
|
|
tb.tokens = tb.capacity
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|