TOKEN BUCKET ALGORITHM

The token bucket rate limiting algorithm is a popular method for rate limiting, used by companies like Amazon and Stripe.

A token bucket is a container that has a pre-defined capacity. Tokens are put in the bucket at preset rates periodically. Once the bucket is full, no more tokens are added. Each request consumes one token. When a request arrives, we first check if there are available tokens in the bucket. If there are no tokens available, the request is denied.

The token bucket algorithm takes two parameters. Bucket size (the number of tokens the bucket can store) and refill rate (number of tokens put into the bucket every second).

Example:

package main

import (
	"fmt"
	"sync"
	"time"
)

// TokenBucket represents a token bucket rate limiter
type TokenBucket struct {
	tokens           int
	capacity         int
	tokenRate        time.Duration
	lastRefill       time.Time
	mu               sync.Mutex
}

// NewTokenBucket creates a new token bucket
func NewTokenBucket(capacity int, refillRate time.Duration) *TokenBucket {
	return &TokenBucket{
		tokens:           capacity,
		capacity:         capacity,
		tokenRate:        refillRate,
		lastRefill:       time.Now(),
	}
}

// refill refills tokens in the bucket based on the elapsed time since the last refill
func (tb *TokenBucket) refill() {
	now := time.Now()
	elapsed := now.Sub(tb.lastRefill)
	tokensToAdd := int(elapsed / tb.tokenRate)
	if tokensToAdd > 0 {
		tb.tokens = min(tb.capacity, tb.tokens+tokensToAdd)
		tb.lastRefill = now
	}
}

// Consume consumes a token from the bucket if available
func (tb *TokenBucket) Consume() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	tb.refill()
	if tb.tokens > 0 {
		tb.tokens--
		return true
	}
	return false
}

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

func main() {
	tb := NewTokenBucket(10, time.Second)  // Capacity of 10 tokens, and refills 1 token every second

	// Simulating rapid requests
	for i := 0; i < 15; i++ {
		if tb.Consume() {
			fmt.Println("Request", i, "allowed")
		} else {
			fmt.Println("Request", i, "denied")
		}
		time.Sleep(500 * time.Millisecond)
	}
}