ratelimiter.go

// Package ratelimiter provides multiple rate limiting algorithms for API protection.
// Includes Token Bucket, Sliding Window Log, and Sliding Window Counter implementations.
package ratelimiter

import (
	"sync"
	"time"
)

// RateLimiter defines the interface for all rate limiters
type RateLimiter interface {
	Allow(key string) (bool, error)
	AllowN(key string, n int) (bool, error)
}

// =============================================================================
// Token Bucket Implementation
// =============================================================================

// TokenBucket implements the token bucket algorithm
type TokenBucket struct {
	rate       int                      // tokens per second
	capacity   int                      // max tokens in bucket
	buckets    map[string]*bucket       // per-key buckets
	mu         sync.RWMutex
	gcInterval time.Duration            // garbage collection interval
	lastGC     time.Time
}

type bucket struct {
	tokens       float64
	lastUpdate   time.Time
}

// NewTokenBucket creates a new token bucket rate limiter
func NewTokenBucket(rate, capacity int) *TokenBucket {
	tb := &TokenBucket{
		rate:       rate,
		capacity:   capacity,
		buckets:    make(map[string]*bucket),
		gcInterval: 5 * time.Minute,
		lastGC:     time.Now(),
	}
	
	// Start garbage collector
	go tb.gc()
	
	return tb
}

// Allow checks if a request is allowed
func (tb *TokenBucket) Allow(key string) (bool, error) {
	return tb.AllowN(key, 1)
}

// AllowN checks if n tokens are available
func (tb *TokenBucket) AllowN(key string, n int) (bool, error) {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	
	now := time.Now()
	b, exists := tb.buckets[key]
	
	if !exists {
		b = &bucket{
			tokens:     float64(tb.capacity),
			lastUpdate: now,
		}
		tb.buckets[key] = b
	}
	
	// Calculate tokens to add since last update
	elapsed := now.Sub(b.lastUpdate).Seconds()
	b.tokens += elapsed * float64(tb.rate)
	
	// Cap at capacity
	if b.tokens > float64(tb.capacity) {
		b.tokens = float64(tb.capacity)
	}
	
	// Update timestamp
	b.lastUpdate = now
	
	// Check if we have enough tokens
	if b.tokens >= float64(n) {
		b.tokens -= float64(n)
		return true, nil
	}
	
	return false, nil
}

// gc performs garbage collection on inactive buckets
func (tb *TokenBucket) gc() {
	ticker := time.NewTicker(tb.gcInterval)
	defer ticker.Stop()
	
	for range ticker.C {
		tb.mu.Lock()
		now := time.Now()
		for key, b := range tb.buckets {
			if now.Sub(b.lastUpdate) > tb.gcInterval {
				delete(tb.buckets, key)
			}
		}
		tb.mu.Unlock()
	}
}

// =============================================================================
// Sliding Window Log Implementation
// =============================================================================

// SlidingWindowLog implements the sliding window log algorithm
type SlidingWindowLog struct {
	limit    int
	window   time.Duration
	requests map[string][]time.Time
	mu       sync.RWMutex
}

// NewSlidingWindowLog creates a new sliding window log rate limiter
func NewSlidingWindowLog(limit int, window time.Duration) *SlidingWindowLog {
	swl := &SlidingWindowLog{
		limit:    limit,
		window:   window,
		requests: make(map[string][]time.Time),
	}
	
	// Start cleaner
	go swl.cleaner()
	
	return swl
}

// Allow checks if a request is allowed
func (swl *SlidingWindowLog) Allow(key string) (bool, error) {
	return swl.AllowN(key, 1)
}

// AllowN checks if n requests are allowed
func (swl *SlidingWindowLog) AllowN(key string, n int) (bool, error) {
	swl.mu.Lock()
	defer swl.mu.Unlock()
	
	now := time.Now()
	windowStart := now.Add(-swl.window)
	
	// Get or create request log
	requests, exists := swl.requests[key]
	if !exists {
		requests = make([]time.Time, 0)
	}
	
	// Remove old requests outside the window
	validRequests := make([]time.Time, 0)
	for _, t := range requests {
		if t.After(windowStart) {
			validRequests = append(validRequests, t)
		}
	}
	
	// Check if adding n requests would exceed limit
	if len(validRequests)+n > swl.limit {
		swl.requests[key] = validRequests
		return false, nil
	}
	
	// Add new requests
	for i := 0; i < n; i++ {
		validRequests = append(validRequests, now)
	}
	swl.requests[key] = validRequests
	
	return true, nil
}

// cleaner periodically removes old entries
func (swl *SlidingWindowLog) cleaner() {
	ticker := time.NewTicker(swl.window)
	defer ticker.Stop()
	
	for range ticker.C {
		swl.mu.Lock()
		now := time.Now()
		windowStart := now.Add(-swl.window)
		
		for key, requests := range swl.requests {
			// Remove entries with all requests outside window
			if len(requests) > 0 && requests[len(requests)-1].Before(windowStart) {
				delete(swl.requests, key)
			}
		}
		swl.mu.Unlock()
	}
}

// =============================================================================
// Sliding Window Counter Implementation
// =============================================================================

// SlidingWindowCounter implements the sliding window counter algorithm
type SlidingWindowCounter struct {
	limit        int
	windowSize   time.Duration
	subWindows   int
	counters     map[string]*windowCounter
	mu           sync.RWMutex
}

type windowCounter struct {
	counts     []int
	timestamps []time.Time
}

// NewSlidingWindowCounter creates a new sliding window counter rate limiter
func NewSlidingWindowCounter(limit int, windowSize time.Duration, subWindows int) *SlidingWindowCounter {
	if subWindows < 2 {
		subWindows = 2 // minimum for sliding window
	}
	
	return &SlidingWindowCounter{
		limit:      limit,
		windowSize: windowSize,
		subWindows: subWindows,
		counters:   make(map[string]*windowCounter),
	}
}

// Allow checks if a request is allowed
func (swc *SlidingWindowCounter) Allow(key string) (bool, error) {
	return swc.AllowN(key, 1)
}

// AllowN checks if n requests are allowed
func (swc *SlidingWindowCounter) AllowN(key string, n int) (bool, error) {
	swc.mu.Lock()
	defer swc.mu.Unlock()
	
	now := time.Now()
	subWindowSize := swc.windowSize / time.Duration(swc.subWindows)
	
	// Get or create counter
	counter, exists := swc.counters[key]
	if !exists {
		counter = &windowCounter{
			counts:     make([]int, swc.subWindows),
			timestamps: make([]time.Time, swc.subWindows),
		}
		swc.counters[key] = counter
	}
	
	// Find current sub-window
	currentSubWindow := int(now.UnixNano() / int64(subWindowSize))
	
	// Calculate total requests in current window
	total := 0
	windowStart := now.Add(-swc.windowSize)
	
	for i := 0; i < swc.subWindows; i++ {
		if counter.timestamps[i].After(windowStart) {
			// Calculate weight for partial windows
			if i == 0 && len(counter.timestamps) > 0 {
				// Weighted count for oldest sub-window
				elapsed := now.Sub(counter.timestamps[i])
				if elapsed < swc.windowSize {
					weight := float64(swc.windowSize-elapsed) / float64(subWindowSize)
					total += int(float64(counter.counts[i]) * weight)
				}
			} else {
				total += counter.counts[i]
			}
		}
	}
	
	// Check if adding n requests would exceed limit
	if total+n > swc.limit {
		return false, nil
	}
	
	// Update current sub-window
	idx := currentSubWindow % swc.subWindows
	if counter.timestamps[idx].IsZero() || 
	   now.Sub(counter.timestamps[idx]) >= subWindowSize {
		// New sub-window
		counter.counts[idx] = n
		counter.timestamps[idx] = now
	} else {
		// Same sub-window
		counter.counts[idx] += n
	}
	
	return true, nil
}

// =============================================================================
// Fixed Window Counter Implementation (Bonus)
// =============================================================================

// FixedWindow implements a simple fixed window counter
type FixedWindow struct {
	limit    int
	window   time.Duration
	counters map[string]*fixedWindowCounter
	mu       sync.RWMutex
}

type fixedWindowCounter struct {
	count      int
	windowStart time.Time
}

// NewFixedWindow creates a new fixed window rate limiter
func NewFixedWindow(limit int, window time.Duration) *FixedWindow {
	return &FixedWindow{
		limit:    limit,
		window:   window,
		counters: make(map[string]*fixedWindowCounter),
	}
}

// Allow checks if a request is allowed
func (fw *FixedWindow) Allow(key string) (bool, error) {
	return fw.AllowN(key, 1)
}

// AllowN checks if n requests are allowed
func (fw *FixedWindow) AllowN(key string, n int) (bool, error) {
	fw.mu.Lock()
	defer fw.mu.Unlock()
	
	now := time.Now()
	counter, exists := fw.counters[key]
	
	if !exists || now.Sub(counter.windowStart) >= fw.window {
		// New window
		fw.counters[key] = &fixedWindowCounter{
			count:       n,
			windowStart: now.Truncate(fw.window),
		}
		return true, nil
	}
	
	// Check if adding n requests would exceed limit
	if counter.count+n > fw.limit {
		return false, nil
	}
	
	counter.count += n
	return true, nil
}