Rate Limiting with Token Bucket Algorithm

package main

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"
)

// TokenBucket implements a thread-safe token bucket rate limiter
type TokenBucket struct {
	mu           sync.Mutex
	capacity     int           // Maximum number of tokens
	tokens       int           // Current number of tokens
	refillRate   int           // Tokens to add per refill interval
	refillPeriod time.Duration // Interval for refilling tokens
	lastRefill   time.Time     // Timestamp of last refill
}

// NewTokenBucket creates a new TokenBucket with specified capacity and refill rate
func NewTokenBucket(capacity, refillRate int, refillPeriod time.Duration) *TokenBucket {
	return &TokenBucket{
		capacity:     capacity,
		tokens:       capacity, // Start with full bucket
		refillRate:   refillRate,
		refillPeriod: refillPeriod,
		lastRefill:   time.Now(),
	}
}

// refill adds tokens based on time elapsed since last refill
func (tb *TokenBucket) refill() {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.lastRefill)
	if elapsed < tb.refillPeriod {
		return
	}

	// Calculate number of refills to perform
	refills := int(elapsed / tb.refillPeriod)
	newTokens := refills * tb.refillRate
	tb.tokens = min(tb.tokens+newTokens, tb.capacity)
	tb.lastRefill = now.Add(-(elapsed % tb.refillPeriod)) // Preserve leftover time
}

// Allow checks if a token is available (blocks until token is available or context expires)
func (tb *TokenBucket) Allow(ctx context.Context) bool {
	for {
		tb.refill()

		tb.mu.Lock()
		if tb.tokens > 0 {
			tb.tokens--
			tb.mu.Unlock()
			return true
		}
		tb.mu.Unlock()

		// Check if context is done before waiting
		select {
		case <-ctx.Done():
			return false
		case <-time.After(100 * time.Millisecond): // Short sleep to prevent busy waiting
			continue
		}
	}
}

// Example usage
func main() {
	// Create rate limiter: 5 tokens capacity, 1 token per second refill rate
	limiter := NewTokenBucket(5, 1, 1*time.Second)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var wg sync.WaitGroup
	// Simulate 10 concurrent requests
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(requestID int) {
			defer wg.Done()
			
			start := time.Now()
			if limiter.Allow(ctx) {
				log.Printf("Request %d: allowed (waited %v)", requestID, time.Since(start))
			} else {
				log.Printf("Request %d: denied (context expired)", requestID)
			}
		}(i)
	}

	wg.Wait()
	fmt.Println("All requests processed")
}