Token Bucket Rate Limiter Concurrency Safe

package main

import (
	"math"
	"sync"
	"time"
)

type TokenBucket struct {
	mu       sync.Mutex
	tokens   float64
	capacity float64
	refill   float64 // tokens per second
	last     time.Time
}

func NewTokenBucket(capacity float64, refillPerSec float64) *TokenBucket {
	return &TokenBucket{
		tokens:   capacity,
		capacity: capacity,
		refill:   refillPerSec,
		last:     time.Now(),
	}
}

func (b *TokenBucket) take(now time.Time, n float64) bool {
	elapsed := now.Sub(b.last).Seconds()
	b.tokens = math.Min(b.capacity, b.tokens+elapsed*b.refill)
	b.last = now

	if b.tokens >= n {
		b.tokens -= n
		return true
	}
	return false
}

func (b *TokenBucket) TryTake(n float64) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.take(time.Now(), n)
}

type KeyedLimiter struct {
	mu       sync.Mutex
	buckets  map[string]*TokenBucket
	capacity float64
	refill   float64
}

func (k *KeyedLimiter) TryTake(key string, n float64) bool {
	k.mu.Lock()
	b, ok := k.buckets[key]
	if !ok {
		b = NewTokenBucket(k.capacity, k.refill)
		k.buckets[key] = b
	}
	k.mu.Unlock()
	return b.TryTake(n)
}