Systems Thinking and Architecture Questions

package ratelimit

import (
	"sync"
	"time"
)

// TokenBucket is a thread-safe token bucket rate limiter.
type TokenBucket struct {
	rate     float64       // tokens per second
	capacity float64       // max tokens (burst)
	tokens   float64       // current tokens
	last     time.Time     // last refill timestamp
	mu       sync.Mutex
}

// NewTokenBucket creates a new TokenBucket.
func NewTokenBucket(tokensPerSecond float64, burst float64) *TokenBucket {
	if tokensPerSecond <= 0 {
		tokensPerSecond = 1
	}
	if burst <= 0 {
		burst = tokensPerSecond
	}
	return &TokenBucket{
		rate:     tokensPerSecond,
		capacity: burst,
		tokens:   burst,
		last:     time.Now(),
	}
}

// AllowRequest attempts to consume one token. Returns true if allowed.
func (tb *TokenBucket) AllowRequest() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.last).Seconds()
	// Refill tokens based on elapsed time
	tb.tokens += elapsed * tb.rate
	if tb.tokens > tb.capacity {
		tb.tokens = tb.capacity
	}
	tb.last = now

	if tb.tokens >= 1.0 {
		tb.tokens -= 1.0
		return true
	}
	return false
}