# 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())
	}
}
```