-
Notifications
You must be signed in to change notification settings - Fork 0
/
tokenbucket.go
62 lines (52 loc) · 1.19 KB
/
tokenbucket.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
package ginratelimit
import (
"sync"
"time"
)
type TokenInfo struct {
RemainingTokens int
ExpiresAt time.Time
}
type TokenBucket struct {
threshold int
ttl time.Duration
tokens map[string]TokenInfo
mu sync.Mutex
}
func NewTokenBucket(threshold int, ttl time.Duration) *TokenBucket {
return &TokenBucket{
threshold: threshold,
ttl: ttl,
tokens: make(map[string]TokenInfo),
}
}
func (tb *TokenBucket) Allow(key string) bool {
tb.mu.Lock()
defer tb.mu.Unlock()
now := time.Now()
tokenInfo, exists := tb.tokens[key]
// Refill tokens if the record exists and has expired
if exists && now.After(tokenInfo.ExpiresAt) {
tokenInfo = TokenInfo{
RemainingTokens: tb.threshold,
ExpiresAt: now.Add(tb.ttl),
}
}
// If the key doesn't exist, initialize it
if !exists {
tokenInfo = TokenInfo{
RemainingTokens: tb.threshold - 1, // Consume a token
ExpiresAt: now.Add(tb.ttl),
}
tb.tokens[key] = tokenInfo
return true
}
// If tokens are available, decrement and allow
if tokenInfo.RemainingTokens > 0 {
tokenInfo.RemainingTokens--
tb.tokens[key] = tokenInfo
return true
}
// No tokens available, reject
return false
}