File tree Expand file tree Collapse file tree 1 file changed +47
-1
lines changed
Expand file tree Collapse file tree 1 file changed +47
-1
lines changed Original file line number Diff line number Diff line change 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+ }
You can’t perform that action at this time.
0 commit comments