Notes: Token Bucket Rate Limiter in Go

Token Bucket is one of the most used Rate Limiting algorithms. In Token Bucket, as the name suggests, we have a fixed amount of tokens to be exhausted in a timeframe. However, in our implementation timeframe-based filling is not included, rather we have added RefillTokenBucket method which makes it more flexible.

TokenBucket manages a state with availableTokens and last refill time, lastUpdateTime.

type TokenBucket struct {
    availableTokens int `json:"availableTokens"`
    lastUpdateTime  int `json:"lastUpdateTime"`
}

Now, we have four methods implemented in TokenBucket.

• Reset: This resets TokenBucket refilling tokens with MAX_TOKEN_CAPACITY.
• RefillTokenBucket: This refills tokens based on difference between lastUpdateTime and current time.
• ReduceTokensFromBucket: This reduces tokens by the passed integer amount. In case, there are not enough tokens, it will throw an error.
• GetTokenBucketState: This just return current TokenBucket.

And this ends the implementation of our Token Bucket. Now to use this, we can add it as middleware for our webservers and use available methods as per our requirements. We can run a goroutine just to refill the bucket after a certain time by calling RefillTokenBucket. We might want to use a lock as well in case of concurrent reads and writes. We'll implement these in upcoming blogs.

In the meantime, this is the complete code with implementation which you can use to understand Token Bucket:

package main

import (
    "errors"
    "fmt"
    "time"
)

const MAX_TOKEN_CAPACITY = 10
const REFILL_RATE = 2 // 2 tokens per second

type TokenBucket struct {
    availableTokens int `json:"availableTokens"`
    lastUpdateTime  int `json:"lastUpdateTime"` // just whenever we refill
}

type ITokenBucket interface {
    Reset() *TokenBucket
    RefillTokenBucket() *TokenBucket
    ReduceTokensFromBucket(tokenCount int) *TokenBucket
    GetTokenBucketState() *TokenBucket
}

func (b *TokenBucket) Reset() *TokenBucket {
    b.availableTokens = MAX_TOKEN_CAPACITY
    b.lastUpdateTime = currentUnixTime()

    return b
}

func (b *TokenBucket) RefillTokenBucket() *TokenBucket {
    // round to nearest integer
    // not required in this case
    tokenRefillCount := (currentUnixTime() - b.lastUpdateTime) * REFILL_RATE
    netTokenCount := b.availableTokens + tokenRefillCount

    b.availableTokens = min(netTokenCount, MAX_TOKEN_CAPACITY)
    b.lastUpdateTime = min(currentUnixTime(), b.lastUpdateTime+(tokenRefillCount/REFILL_RATE))

    return b
}

func (b *TokenBucket) ReduceTokensFromBucket(tokenCount int) (*TokenBucket, error) {
    if b.availableTokens-tokenCount < 0 {
        return nil, errors.New("not enough tokens")
    }

    b.availableTokens -= tokenCount
    return b, nil
}

func (b *TokenBucket) GetTokenBucketState() *TokenBucket {
    return b
}

func NewTokenBucket() *TokenBucket {
    return &TokenBucket{
        availableTokens: MAX_TOKEN_CAPACITY,
        lastUpdateTime:  currentUnixTime(),
    }
}

// helper functions
func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

func currentUnixTime() int {
    return int(time.Now().Unix())
}

func main() {
    tokenBucket := NewTokenBucket()
    fmt.Printf("bucket current state: %#v\n", tokenBucket.GetTokenBucketState())

    tokenBucket, err := tokenBucket.ReduceTokensFromBucket(7)
    if err != nil {
        fmt.Println(err.Error())
    }

    // adding sleep to fill bucket
    time.Sleep(2 * time.Second)
    fmt.Printf("bucket updated state: %#v\n", tokenBucket.RefillTokenBucket().GetTokenBucketState())

    _, err = tokenBucket.ReduceTokensFromBucket(8)
    if err != nil {
        fmt.Println(err.Error())
    }
}