webriots/rate

The numBuckets parameter determines how many internal buckets are used to distribute rate limits across different IDs. The library automatically rounds this value up to the nearest power of two for efficient hashing.

Key Concepts

• Bucket Selection: The library hashes the ID (using hash/maphash) and uses hash & (numBuckets - 1) to select a bucket.
• Collisions: When two IDs map to the same bucket, they share the same rate limit. While the probability of at least one collision is high, the percentage of IDs affected is what matters for system fairness.
• Load Factor: To keep the impact of collisions below 5% of your total IDs, aim to maintain a load factor (IDs / buckets) below 0.1.

Guidelines by System Size

• Small (≤100 IDs): Use 4,096 or 16,384 buckets.
• Medium (101-1,000 IDs): Use 16,384 or 65,536 buckets.
• Large (1,001-10,000 IDs): Use 65,536 or 262,144 buckets.
• Very Large (>10,000 IDs): Use 1,048,576 buckets.

Memory Usage

Memory scales linearly with numBuckets. Each bucket consumes 8 bytes (uint64).

• 65,536 buckets $\approx$ 512 KB
• 1,048,576 buckets $\approx$ 8 MB

// Example for a system with ~5,000 distinct IDs
limiter, err := rate.NewTokenBucketLimiter(
    65536,    // numBuckets (will remain 2^16 since it's already a power of 2) to maintain <10% collision probability
    10,       // burstCapacity
    100,      // refillRate
    time.Second,
)

All rate limiters in this package implement the Limiter interface. This allows you to write generic code that can work with any implementation (Token Bucket, AIMD, or Rotating Token Bucket) by accepting the interface instead of a concrete type.

Interface Methods:

• CheckToken([]byte) bool: Check if a single token is available without consuming it.
• CheckTokens([]byte, uint8) bool: Check if n tokens are available without consuming them.
• TakeToken([]byte) bool: Attempt to consume a single token; returns true if successful.
• TakeTokens([]byte, n uint8) bool: Attempt to consume n tokens atomically; returns true only if all n tokens were successfully taken.

func processRequest(limiter rate.Limiter, clientID []byte, tokensNeeded uint8) error {
	// Check if we have enough tokens without consuming them
	if !limiter.CheckTokens(clientID, tokensNeeded) {
		return fmt.Errorf("rate limited: need %d tokens", tokensNeeded)
	}

	// Actually take the tokens
	if !limiter.TakeTokens(clientID, tokensNeeded) {
		return fmt.Errorf("rate limited: tokens no longer available")
	}

	// Process the request...
	return nil
}

The rate package is designed for high-concurrency environments with the following guarantees:

Atomic Operations

• Lock-free: All token bucket operations use lock-free atomic operations (e.g., Compare-And-Swap/CAS).
• No Mutexes: The library avoids traditional mutexes to eliminate lock contention.
• Thread-safe Hashing: Bucket selection uses thread-safe hash calculations.

Concurrent Access Patterns

• Read/Write Safety: Multiple goroutines can safely check token availability (TakeToken/Check) and update rates concurrently.
• Automatic Refill: Token refill happens automatically during token operations without requiring explicit locks.

Limitations

• Fairness: The library does not guarantee perfect fairness across all IDs.
• Collisions: When multiple IDs hash to the same bucket, they share the same rate limit.
• Contention: Very high contention on a single ID might experience CAS retry loops.

The TokenBucketLimiter implements a classic token bucket algorithm. It uses multiple buckets to distribute load across different request IDs, making it highly scalable for concurrent systems.

Initialization

Use rate.NewTokenBucketLimiter(numBuckets, burstCapacity, refillRate, refillRateUnit) to create a new limiter.

• numBuckets: The number of buckets (automatically rounded to the nearest power of 2).
• burstCapacity: Maximum number of tokens available for a burst.
• refillRate: The amount of tokens added per unit of time.
• refillRateUnit: The time duration for the refill rate (e.g., time.Second).

Core Methods

• TakeToken(id []byte) bool: Attempts to consume a single token for the given ID. Returns true if successful.
• TakeTokens(id []byte, n int) bool: Attempts to consume n tokens for the given ID.
• CheckToken(id []byte) bool: Checks if a token is available for the ID without consuming it.
• CheckTokens(id []byte, n int) bool: Checks if n tokens are available for the ID without consuming them.

package main

import (
	"fmt"
	"time"

	"github.com/webriots/rate"
)

func main() {
	// Create a new token bucket limiter:
	// - 1024 buckets (automatically rounded to nearest power of 2 if
	//   not already a power of 2)
	// - 10 tokens burst capacity
	// - 100 tokens per second refill rate
	limiter, err := rate.NewTokenBucketLimiter(1024, 10, 100, time.Second)
	if err != nil {
		panic(err)
	}

	// Try to take a token for a specific ID
	id := []byte("user-123")

	if limiter.TakeToken(id) {
		fmt.Println("Token acquired, proceeding with request")
	} else {
		fmt.Println("Rate limited, try again later")
	}

	// Check without consuming
	if limiter.CheckToken(id) {
		fmt.Println("Token would be available")
	}

	// Take multiple tokens at once
	if limiter.TakeTokens(id, 5) {
		fmt.Println("Successfully took 5 tokens")
	}

	// Check if multiple tokens would be available
	if limiter.CheckTokens(id, 3) {
		fmt.Println("3 tokens would be available")
	}
}

The AIMDTokenBucketLimiter uses an Additive-Increase/Multiplicative-Decrease algorithm inspired by TCP congestion control. This is ideal for adjusting request rates dynamically based on the success or failure of downstream services.

Initialization

Use rate.NewAIMDTokenBucketLimiter(...) with the following parameters:

• numBuckets: Number of buckets (rounded to power of 2).
• burstCapacity: Maximum burst capacity.
• rateMin: Minimum allowed rate.
• rateMax: Maximum allowed rate.
• rateInit: Initial starting rate.
• rateAdditiveIncrease: The amount to increase the rate by on success.
• rateMultiplicativeDecrease: The factor to decrease the rate by on failure.
• refillRateUnit: The time duration for the rate units.

Dynamic Rate Adjustment

Unlike the standard limiter, you must manually signal the outcome of your request to adjust the rate:

• IncreaseRate(id []byte): Call this when a request is successful to increase the allowed rate.
• DecreaseRate(id []byte): Call this when a request fails (e.g., due to downstream overload) to back off the rate.

package main

import (
	"fmt"
	"time"

	"github.com/webriots/rate"
)

func main() {
	// Create an AIMD token bucket limiter:
	// - 1024 buckets
	// - 10 tokens burst capacity
	// - Min rate: 1 token/s, Max rate: 100 tokens/s, Initial rate: 10 tokens/s
	// - Increase by 1 token/s on success, decrease by factor of 2 on failure
	limiter, err := rate.NewAIMDTokenBucketLimiter(
		1024,  // numBuckets
		10,    // burstCapacity
		1.0,   // rateMin
		100.0, // rateMax
		10.0,  // rateInit
		1.0,   // rateAdditiveIncrease
		2.0,   // rateMultiplicativeDecrease
		time.Second,
	)
	if err != nil {
		panic(err)
	}

	id := []byte("api-client-123")

	if limiter.TakeToken(id) {
		success := processRequest()

		if success {
			limiter.IncreaseRate(id)
		} else {
			limiter.DecreaseRate(id)
		}
	} else {
		fmt.Println("Rate limited, try again later")
	}
}

func processRequest() bool {
	fmt.Println("Processing request")
	return true
}

The RotatingTokenBucketLimiter is a collision-resistant rate limiter that uses periodic hash seed rotation. This is useful for high-throughput systems where minimizing collisions across many request IDs is critical.

Initialization

Use rate.NewRotatingTokenBucketLimiter(numBuckets, burstCapacity, refillRate, refillRateUnit).

• numBuckets: Number of buckets (rounded to power of 2).
• burstCapacity: Maximum burst capacity.
• refillRate: Amount of tokens added per unit of time.
• refillRateUnit: The time duration for the refill rate.

Usage

It supports the same interface as the standard TokenBucketLimiter, including TakeToken, TakeTokens, CheckToken, and CheckTokens.

package main

import (
	"fmt"
	"time"

	"github.com/webriots/rate"
)

func main() {
	// Create a rotating token bucket limiter:
	// - 1024 buckets
	// - 10 tokens burst capacity
	// - 100 tokens per second refill rate
	limiter, err := rate.NewRotatingTokenBucketLimiter(
		1024,        // numBuckets
		10,          // burstCapacity
		100,         // refillRate
		time.Second, // refillRateUnit
	)
	if err != nil {
		panic(err)
	}

	id := []byte("user-456")

	if limiter.TakeToken(id) {
		fmt.Println("Token acquired, proceeding with request")
	} else {
		fmt.Println("Rate limited, try again later")
	}

	if limiter.CheckToken(id) {
		fmt.Println("Token would be available")
	}

	if limiter.TakeTokens(id, 3) {
		fmt.Println("Successfully took 3 tokens")
	}
}

A standard TokenBucketLimiter using hashing can suffer from accuracy issues if different IDs hash to the same bucket (a collision). RotatingTokenBucketLimiter mitigates this through a dual-limiter rotation strategy:

1. Dual Limiters: It maintains a rotatingPair consisting of a checked limiter (used for decisions) and an ignored limiter (used to maintain state during rotation).
2. Rotation: Periodically, the roles swap. The current ignored limiter becomes the new checked limiter, and a new ignored limiter is created with a fresh hash seed.
3. Synchronization: Every TakeTokens or CheckTokens call is applied to both limiters. This ensures that when a rotation occurs, the new checked limiter is already in a consistent state.
4. Impact: Because of this dual-consumption, the effective token consumption rate is 2x. However, the benefit is that any hash collisions only persist for the duration of one RotationInterval.

The AIMDTokenBucketLimiter (Additive Increase, Multiplicative Decrease) is designed for dynamic backend protection. It allows you to adjust the rate based on backend health.

• Increase Rate: Call IncreaseRate(id) when the backend call succeeds.
• Decrease Rate: Call DecreaseRate(id) when the backend returns an overload error.

// Create AIMD limiter for dynamic backend protection
limiter, _ := rate.NewAIMDTokenBucketLimiter(
    4096,   // numBuckets
    20,     // burstCapacity
    1.0,    // rateMin
    1000.0, // rateMax
    100.0,  // rateInit
    10.0,   // rateAdditiveIncrease
    2.0,    // rateMultiplicativeDecrease
    time.Second,
)

func handleRequest(backendID []byte) {
    if limiter.TakeToken(backendID) {
        err := callBackendService()
        if err == nil {
            // Backend is healthy, increase rate
            limiter.IncreaseRate(backendID)
        } else if isOverloadError(err) {
            // Backend is overloaded, decrease rate
            limiter.DecreaseRate(backendID)
        }
    } else {
        // Circuit breaking, return error immediately
    }
}

The RotatingTokenBucketLimiter addresses hash collision unfairness by maintaining two TokenBucketLimiters with different hash seeds and rotating between them. This ensures that if two IDs collide in one limiter, they will likely have independent buckets in the next rotation.

Constructor: NewRotatingTokenBucketLimiter(numBuckets, burstCapacity, refillRate, refillRateUnit)

Parameters:

• numBuckets: Number of buckets per limiter (rounded to nearest power of 2).
• burstCapacity: Maximum tokens per bucket.
• refillRate: Rate of refill.
• refillRateUnit: Time unit for refill rate.

Key Features:

• Automatic Rotation: The rotation interval is calculated as (burstCapacity / refillRate * refillRateUnit) * 5.0 to ensure statistical convergence.
• Collision Resistance: Uses a dual-limiter architecture where tokens are consumed from both, but only the "checked" limiter determines the result.
• Method RotationInterval() time.Duration: Returns the calculated interval at which the limiters swap roles.

Trade-offs:

• Higher Resource Usage: Consumes ~2x tokens and ~2x memory compared to a standard TokenBucketLimiter.
• Performance: Approximately ~2x slower due to dual-limiter management.

When to use: Use when hash collision fairness is critical or when you have many distinct IDs with similar access patterns and can afford the resource overhead.

limiter, err := rate.NewRotatingTokenBucketLimiter(
    numBuckets,     // Number of buckets
    burstCapacity,  // Max tokens per bucket
    refillRate,     // Rate of refill
    refillRateUnit, // Time unit
)