Rate Limiting with Token Buckets

Rate Limiting with Token Buckets

In this lesson, you will learn how to throttle asynchronous requests to an external service by implementing a Token Bucket algorithm that manages your request budget over time.

The Token Bucket Concept

Imagine you have a bucket that holds a maximum number of tokens. Each token represents the permission to perform one API request. The bucket refills at a steady, fixed rate. If the bucket is full, new tokens overflow and are discarded. If you want to make a request, you must first "spend" a token from the bucket. If the bucket is empty, you must wait until a new token is added.

This mechanism is ideal for asynchronous services because it allows for "bursts" (when the bucket is full) while strictly enforcing a long-term average request rate.

Implementation Logic

To implement this, we need to track two pieces of state: the current number of tokens and the last time the bucket was refilled. When an async task attempts to fire, it checks if enough tokens exist. If not, it sleeps until the refill timer provides the necessary capacity.

Here is a simplified look at how you might track this state in an async context:

struct TokenBucket {
    tokens: f64,
    max_tokens: f64,
    refill_rate: f64, // tokens per second
    last_refill: std::time::Instant,
}

impl TokenBucket {
    fn refill(&mut self) {
        let now = std::time::Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens);
        self.last_refill = now;
    }

    async fn acquire(&mut self) {
        loop {
            self.refill();
            if self.tokens >= 1.0 {
                self.tokens -= 1.0;
                return;
            }
            // If empty, wait for the next potential refill
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        }
    }
}

By wrapping this TokenBucket in an Async Mutex, you can safely share the bucket across multiple concurrent tasks. When a task calls acquire(), it asynchronously suspends execution if the limit is reached, ensuring your system remains responsive without overwhelming the downstream service. This approach effectively converts a high-frequency spike of tasks into a steady, manageable stream of outgoing requests.

Key Takeaway: The Token Bucket pattern uses a shared state of refillable credits to transform uncontrollable traffic bursts into a predictable, service-friendly flow of asynchronous events.