Rate Limiter : Token Bucket Algorithm

Implement Rate Limiter using Token Bucket Algorithms

Source Code in Golang

package main
import (
   "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.")
}

Output

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.