bucket.go

package rate

import (
	"fmt"
	"hash/maphash"
	"math"
	"time"

	"github.com/webriots/rate/time56"
)

type Limiter interface {
	CheckToken([]byte) bool
	CheckTokens([]byte, uint8) bool
	TakeToken([]byte) bool
	TakeTokens([]byte, uint8) bool
}

// TokenBucketLimiter implements the token bucket algorithm for rate
// limiting. It maintains multiple buckets to distribute load and
// reduce contention. Each bucket has a fixed capacity and refills at
// a specified rate.
type TokenBucketLimiter struct {
	buckets       atomicSliceUint64 // Array of token buckets
	burstCapacity uint8             // Maximum tokens per bucket
	nanosPerToken int64             // Nanoseconds per token refill
	seed          maphash.Seed      // Used for index bucket hash gen
}

// Compile-time assertion that TokenBucketLimiter implements Limiter
var _ Limiter = (*TokenBucketLimiter)(nil)

// NewTokenBucketLimiter creates a new token bucket rate limiter with
// the specified parameters:
//
//   - numBuckets: number of token buckets (automatically rounded up
//     to the nearest power of two if not already a power of two, for
//     efficient hashing)
//   - burstCapacity: maximum number of tokens that can be consumed at
//     once
//   - refillRate: rate at which tokens are refilled (must be positive
//     and finite)
//   - refillRateUnit: time unit for refill rate calculations (e.g.,
//     time.Second, must be a positive duration)
//
// Returns a new TokenBucketLimiter instance and any error that
// occurred during creation. The numBuckets parameter is automatically
// rounded up to the nearest power of two if not already a power of
// two, for efficient hashing.
//
// Input validation:
//
//   - If refillRate is not a positive, finite number (e.g., negative,
//     zero, NaN, or infinity), returns an error with message
//     "refillRate must be a positive, finite number"
//   - If refillRateUnit is not a positive duration, returns an error
//     with message "refillRateUnit must represent a positive
//     duration"
//   - If the product of refillRate and refillRateUnit (in
//     nanoseconds) exceeds maximum representable value, returns an
//     error with message "refillRate per duration is too large"
func NewTokenBucketLimiter(
	numBuckets uint,
	burstCapacity uint8,
	refillRate float64,
	refillRateUnit time.Duration,
) (*TokenBucketLimiter, error) {
	if math.IsNaN(refillRate) || math.IsInf(refillRate, 0) || refillRate <= 0 {
		return nil, fmt.Errorf("refillRate must be a positive, finite number")
	}

	if rate := float64(refillRateUnit.Nanoseconds()); rate <= 0 {
		return nil, fmt.Errorf("refillRateUnit must represent a positive duration")
	} else if rate > math.MaxFloat64/refillRate {
		return nil, fmt.Errorf("refillRate per duration is too large")
	}

	n := ceilPow2(uint64(numBuckets))
	bucket := newTokenBucket(burstCapacity, time56.Unix(nowfn()))
	packed := bucket.packed()

	buckets := newAtomicSliceUint64(int(n))
	for i := range buckets.Len() {
		buckets.Set(i, packed)
	}

	return &TokenBucketLimiter{
		buckets:       buckets,
		burstCapacity: burstCapacity,
		nanosPerToken: nanoRate(refillRateUnit, refillRate),
		seed:          maphash.MakeSeed(),
	}, nil
}

// CheckToken returns whether a token would be available for the given
// ID without actually taking it. This is useful for preemptively
// checking if an operation would be rate limited before attempting
// it. Returns true if a token would be available, false otherwise.
func (t *TokenBucketLimiter) CheckToken(id []byte) bool {
	return t.CheckTokens(id, 1)
}

// CheckTokens returns whether n tokens would be available for the
// given ID without actually taking them. This is useful for
// preemptively checking if an operation would be rate limited before
// attempting it. Returns true if all n tokens would be available,
// false otherwise.
func (t *TokenBucketLimiter) CheckTokens(id []byte, n uint8) bool {
	return t.checkTokensWithNow(id, n, nowfn())
}

func (t *TokenBucketLimiter) checkTokensWithNow(id []byte, n uint8, now int64) bool {
	return t.checkInner(t.index(id), t.nanosPerToken, now, n)
}

// TakeToken attempts to take a token for the given ID. It returns
// true if a token was successfully taken, false if the operation
// should be rate limited. This method is thread-safe and can be
// called concurrently from multiple goroutines.
func (t *TokenBucketLimiter) TakeToken(id []byte) bool {
	return t.TakeTokens(id, 1)
}

// TakeTokens attempts to take n tokens for the given ID. It returns
// true if all n tokens were successfully taken, false if the
// operation should be rate limited. This method is thread-safe and
// can be called concurrently from multiple goroutines. The operation
// is atomic: either all n tokens are taken, or none are taken.
func (t *TokenBucketLimiter) TakeTokens(id []byte, n uint8) bool {
	return t.takeTokensWithNow(id, n, nowfn())
}

func (t *TokenBucketLimiter) takeTokensWithNow(id []byte, n uint8, now int64) bool {
	return t.takeTokenInner(t.index(id), t.nanosPerToken, now, n)
}

// checkInner is an internal method that checks if n tokens are
// available in the bucket at the specified index using the given
// refill rate. This is used by CheckToken/CheckTokens and is also
// used by other limiters that wrap this one.
func (t *TokenBucketLimiter) checkInner(index int, rate int64, now int64, n uint8) bool {
	existing := t.buckets.Get(index)
	unpacked := unpack(existing)
	refilled := unpacked.refill(now, rate, t.burstCapacity)
	return refilled.level >= n
}

// takeTokenInner is an internal method that attempts to take n tokens
// from the bucket at the specified index using the given refill rate.
// This is used by TakeToken/TakeTokens and is also used by other
// limiters that wrap this one. It uses atomic operations to ensure
// thread safety.
func (t *TokenBucketLimiter) takeTokenInner(index int, rate int64, now int64, n uint8) bool {
	for {
		existing := t.buckets.Get(index)
		unpacked := unpack(existing)
		refilled := unpacked.refill(now, rate, t.burstCapacity)
		consumed, ok := refilled.take(n)

		if consumed != unpacked && !t.buckets.CompareAndSwap(
			index,
			existing,
			consumed.packed(),
		) {
			continue
		}

		return ok
	}
}

// index calculates the bucket index for the given ID using maphash.
// The result is masked to ensure it falls within the range of valid
// buckets.
func (t *TokenBucketLimiter) index(id []byte) int {
	return int(maphash.Bytes(t.seed, id) & uint64(t.buckets.Len()-1))
}

// tokenBucket represents a single token bucket with a certain level
// (number of tokens) and a timestamp of when it was last refilled.
type tokenBucket struct {
	level uint8       // Current number of tokens in the bucket
	stamp time56.Time // Last time the bucket was refilled
}

// newTokenBucket creates a new token bucket with the specified level
// and timestamp.
func newTokenBucket(level uint8, stamp time56.Time) tokenBucket {
	return tokenBucket{level: level, stamp: stamp}
}

// refill updates the token bucket based on elapsed time since the
// last refill. It calculates how many tokens should be added based on
// the elapsed time and refill rate, and updates the bucket's level
// and timestamp accordingly. The bucket level will not exceed
// maxLevel.
func (b tokenBucket) refill(nowNS, rate int64, maxLevel uint8) tokenBucket {
	now := time56.Unix(nowNS)

	elapsed := now.Since(b.stamp)
	if elapsed <= 0 {
		return b
	}

	tokens := elapsed / rate
	if tokens <= 0 {
		return b
	}

	if avail := maxLevel - b.level; tokens < int64(avail) {
		b.level += uint8(tokens)
	} else {
		b.level = maxLevel
	}

	if remainder := elapsed % rate; remainder > 0 {
		b.stamp = now.Sub(remainder)
	} else {
		b.stamp = now
	}

	return b
}

// take attempts to take n tokens from the bucket. Returns the updated
// bucket and a boolean indicating whether all n tokens were taken. If
// insufficient tokens are available, the bucket remains unchanged and
// false is returned. The operation is atomic: either all n tokens are
// taken, or none are taken.
func (b tokenBucket) take(n uint8) (tokenBucket, bool) {
	if b.level >= n {
		b.level -= n
		return b, true
	} else {
		return b, false
	}
}

// packed converts the token bucket to a packed uint64 representation
// where the level is stored in the high 8 bits and the timestamp in
// the low 56 bits.
func (b tokenBucket) packed() uint64 {
	return b.stamp.Pack(b.level)
}

// unpack extracts a token bucket from its packed uint64
// representation. This is the inverse operation of packed().
func unpack(packed uint64) tokenBucket {
	return newTokenBucket(time56.Unpack(packed))
}

// nanoRate converts a refill rate from tokens per unit to nanoseconds
// per token. This is used to calculate how frequently tokens should
// be added to buckets.
//
// Given a rate like "5 tokens per second", the returned value would
// be the number of nanoseconds per token (200,000,000 ns/token or 0.2
// seconds per token). The formula used is: (duration in nanoseconds)
// / (tokens per unit).
//
// For example:
//   - 1 token per second → 1,000,000,000 ns per token
//   - 2 tokens per second → 500,000,000 ns per token
//   - 0.5 tokens per second → 2,000,000,000 ns per token
func nanoRate(refillRateUnit time.Duration, refillRate float64) int64 {
	return int64(float64(refillRateUnit.Nanoseconds()) / refillRate)
}

// unitRate converts a rate in nanoseconds per token back to tokens
// per unit. This is the inverse operation of nanoRate and is used to
// provide human-readable rate values for APIs returning rate
// information.
//
// Given a rate in nanoseconds per token, it returns the number of
// tokens per time unit. The formula used is: (unit duration in
// nanoseconds) / (nanoseconds per token).
//
// For example:
//   - 1,000,000,000 ns per token → 1.0 tokens per second
//   - 500,000,000 ns per token → 2.0 tokens per second
//   - 2,000,000,000 ns per token → 0.5 tokens per second
func unitRate(refillRateUnit time.Duration, refillRateNanos int64) float64 {
	return float64(refillRateUnit.Nanoseconds()) / float64(refillRateNanos)
}

// maxPow2 defines the maximum power of two value that ceilPow2 will
// return. This prevents potential overflow issues when rounding up
// values close to the maximum uint64 value. Using 2^62 allows safe
// bit manipulation while still providing an extremely large maximum
// bucket count.
const maxPow2 = 1 << 62

// ceilPow2 rounds up the given number to the nearest power of two. If
// the input is already a power of two, it returns the input
// unchanged. This implementation uses a bit manipulation algorithm
// for efficiency.
//
// Special cases:
//   - If input is 0, returns 1 (2^0)
//   - If input exceeds maxPow2 (2^62), returns maxPow2 to avoid
//     overflow
//
// Examples:
//   - ceilPow2(0) → 1
//   - ceilPow2(1) → 1
//   - ceilPow2(2) → 2
//   - ceilPow2(3) → 4
//   - ceilPow2(4) → 4
//   - ceilPow2(5) → 8
func ceilPow2(x uint64) uint64 {
	if x == 0 {
		return 1
	}

	if x >= maxPow2 {
		return maxPow2
	}

	x = x - 1
	x |= x >> 1
	x |= x >> 2
	x |= x >> 4
	x |= x >> 8
	x |= x >> 16
	x |= x >> 32
	return x + 1
}