Rate Limiter : Token Bucket Algorithm

package mainimport (
"fmt"
"sync"
"time"
)
func min(a, b int) int {
if a < b {
return a
}
return b
}
type TockenBucket struct {
MaxBucketSize int
RefillRate int
CurBucketSize int
LastRefillTS time.Time
mutex *sync.Mutex
}
func CreateTockenBucket(maxBucketSize, refillRate int) *TockenBucket {
b := &TockenBucket{
MaxBucketSize: maxBucketSize,
RefillRate: refillRate,
CurBucketSize: maxBucketSize,
LastRefillTS: time.Now(),
mutex: new(sync.Mutex),
}
return b
}
func (b *TockenBucket) AllowRequest(tokens int) bool { b.mutex.Lock()
defer b.mutex.Unlock()
b.refillTokens() if b.CurBucketSize >= tokens {
b.CurBucketSize -= tokens
return true
}
return false
}
func (b *TockenBucket) refillTokens() {
curTime := time.Now()
tokensToAdd := int(curTime.Sub(b.LastRefillTS)) * b.RefillRate / int(time.Second)
b.CurBucketSize = min(b.CurBucketSize+tokensToAdd, b.MaxBucketSize)
b.LastRefillTS = curTime
}
func main() {
fmt.Println("Token Bucket algorithm started.")
maxBucketSize, refillRate := 10, 5
tokenBucket := CreateTockenBucket(maxBucketSize, refillRate)
startTime := time.Now()
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
time.Sleep(1 * time.Second)
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
time.Sleep(100 * time.Millisecond)
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
time.Sleep(100 * time.Millisecond)
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
time.Sleep(100 * time.Millisecond)
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
time.Sleep(100 * time.Millisecond)
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
time.Sleep(500 * time.Millisecond)
fmt.Println("Need 4 tokens. Allowed : ", tokenBucket.AllowRequest(4), " Time lapse : ", time.Now().Sub(startTime))
fmt.Println("Token Bucket algorithm ended.")
}
Token Bucket algorithm started.
Need 4 tokens. Allowed : true Time lapse : 0s
Need 4 tokens. Allowed : true Time lapse : 1s
Need 4 tokens. Allowed : true Time lapse : 1.1s
Need 4 tokens. Allowed : false Time lapse : 1.2s
Need 4 tokens. Allowed : false Time lapse : 1.3s
Need 4 tokens. Allowed : false Time lapse : 1.4s
Need 4 tokens. Allowed : true Time lapse : 1.9s
Token Bucket algorithm ended.

--

--

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store