Skip to main content

Someone wanted to learn this too, so Grasp built them a personal learning path.

You can keep exploring this one or create your own

Token Bucket Rate Limiting in Python and Rust

Hello! Welcome to the fourth module of your course.

In our last lesson, we explored the challenging landscape of race conditions and deadlocks in Python's threading and asyncio models. You learned how to systematically identify these issues using tracebacks and how to resolve them by enforcing atomic operations with locks.

This lesson kicks off our module on Rate Limiting and Traffic Shaping. These concepts are the bedrock of any high-throughput, resilient distributed system. Given your experience designing Rudi, you're already familiar with the critical need to manage load, ensure quality of service, and prevent system overload. Today, we'll dive into one of the most effective and widely-used rate limiting algorithms.

The learning outcome for this lesson is to implement a token bucket rate limiter and compare its Python implementation with Rust patterns used in systems like Rudi. We'll start from the core concept, build a thread-safe Python version, and then directly compare its structure and concurrency model to a tokio-based Rust implementation.


1. The Token Bucket Algorithm: Concept and Rationale

Before we write any code, let's establish a firm conceptual understanding. Rate limiting is about controlling the rate of traffic sent or received by a network interface controller. The token bucket algorithm is an industry standard—used by AWS, Stripe, and others—because it provides a flexible balance between allowing short bursts of traffic and enforcing a long-term average rate.

To get a quick and clear overview of how it works, please watch the following video segment.

Rate Limiter System Design: Token Bucket, Leaky Bucket, Scaling

The video 'Rate Limiter System Design' by ByteByteGo provides an excellent, concise explanation of the token bucket algorithm and its advantages over simpler methods.

Watch the section on the Token Bucket Algorithm (03:06 - 04:20). Focus on these key ideas: The analogy of a bucket being filled with tokens at a steady rate. How each incoming request consumes a token. How this mechanism naturally handles traffic bursts by using accumulated tokens. The two main parameters: bucket capacity (determines the maximum burst size) and refill rate (determines the sustained throughput).

Here is a visual summary of the token bucket model alongside other common strategies.

A Guide to Rate Limiting Strategies
This diagram shows six different rate limiting strategies. The Token Bucket model (bottom-middle) illustrates tokens being added at a steady rate, with requests consuming them. Requests are allowed if tokens are available and rejected or queued otherwise.

The core idea is that the bucket acts as a buffer. During idle periods, the bucket fills up with tokens. When a burst of requests arrives, it can be served quickly by consuming these saved tokens, up to the bucket's capacity. However, the system cannot sustain a rate higher than the refill rate over the long term, effectively "smoothing out" traffic.


2. Implementing a Token Bucket in Python

Now, let's translate this concept into Python code. We'll start with a basic, in-memory implementation. The key is to manage two pieces of state: the current number of tokens and the timestamp of the last refill.

We'll base our implementation on the excellent article from OneUptime.

How to Implement Token Bucket Rate Limiting in Python - OneUptime

The article 'How to Implement Token Bucket Rate Limiting in Python' provides a clean, step-by-step guide to building a token bucket from scratch.

Please read the section Basic In-Memory Implementation. Focus on the TokenBucket class. Pay close attention to: How __post_init__ initializes the bucket to be full at the start. The use of time.monotonic() to get a steadily increasing clock time, which is crucial for preventing issues if the system clock changes. The logic in _refill(): it calculates how much time has passed since the last refill and adds the appropriate number of tokens, capped by the bucket's capacity. The logic in consume(): it refills first, then checks if enough tokens are available. If so, it subtracts them; if not, it calculates the required wait time.

Your Turn: Coding from Scratch

Based on the principles you've just read, try to implement the TokenBucket class yourself in a Python file. This exercise will solidify your understanding of the core mechanics.

Requirements:

  1. An __init__ method that accepts capacity and refill_rate.
  2. It should store tokens and last_refill time.
  3. A consume(tokens_to_consume) method that returns True if the request is allowed and False otherwise.
  4. Ensure you use time.monotonic().

After you've given it a try, you can compare your implementation to the one in the article.


3. Making it Thread-Safe

Our current implementation is not thread-safe. If multiple threads call consume() concurrently, a race condition can occur in the "read-modify-write" sequence on self.tokens, similar to what we saw in the last lesson.

To fix this, we need to protect the critical section with a lock. This ensures that the process of refilling, checking, and consuming tokens is atomic. Given your Rust background, this is directly analogous to wrapping shared state in a std::sync::Mutex.

Let's look at the thread-safe version from the same article.

How to Implement Token Bucket Rate Limiting in Python - OneUptime

The next section of the article extends the basic implementation to be safe for use in multi-threaded applications.

Read the section Thread-Safe Implementation. Note how threading.Lock is introduced. The consume method now uses a with self.lock: block to protect the critical section where _refill_unlocked is called and self.tokens is modified. This pattern guarantees that the lock is released even if an error occurs.

The RateLimiter class in this section also introduces an important pattern for production systems: managing a dictionary of buckets, one for each client_id. This is how you'd implement per-user or per-IP rate limiting.


4. Comparison with Rust (Tokio) Patterns

Now for the core of this lesson: comparing the Python implementation to the patterns you're familiar with in Rust. Your work on Rudi likely involved tokio for handling asynchronous operations and concurrency. The principles are the same, but the ergonomics and safety guarantees differ significantly.

The following resource provides a side-by-side comparison of a web scraper in Python/asyncio and Rust/tokio, with both including a token bucket rate limiter.

convert-python-rust | Skills Marketp... - LobeHub

This resource provides parallel implementations of a RateLimiter, which is perfect for a direct comparison.

First, review the Python asyncio implementation of RateLimiter in section 'Before (Python)'. Then, carefully study the Rust tokio implementation of the RateLimiter struct and its acquire method in section 'After (Rust)'. We will compare them on a few key points.

Let's break down the comparison.

Aspect Python (asyncio from LobeHub) Rust (tokio from LobeHub) Analysis
State Management self.allowance, self.last_check as simple float attributes. allowance: tokio::sync::Mutex<(f64, Instant)>. State is a tuple (allowance, last_check) wrapped in an async mutex. The Rust version is explicit about shared, mutable state. The Mutex enforces that only one task can access the allowance/timestamp pair at a time, preventing race conditions at the language level.
Concurrency Relies on asyncio.sleep() as the only suspension point. No explicit Lock is used. self.allowance.lock().await is called to acquire the lock. The lock guard is explicitly dropped (drop(guard)) before sleeping to prevent holding the lock during a long await. The Rust implementation is more robust. Python's async with lock: (as seen in the OneUptime article) provides similar guard-like behavior automatically. The LobeHub Python example is simpler but could be fragile if more await points were added. Rust's ownership model forces you to be more deliberate about lock lifetimes.
Time Source asyncio.get_event_loop().time() tokio::time::Instant Both correctly use a monotonic clock, which is best practice. Instant is the standard for measuring durations in Rust.
Ergonomics Dynamic and concise. Less boilerplate for a simple case. More verbose but provides compile-time safety guarantees against data races. The type system (f64, Instant) and error handling (Result) are explicit. This is the classic Python vs. Rust trade-off: development speed and flexibility versus performance and compile-time correctness.

This comparison highlights how Rust's design forces you to confront concurrency issues upfront, while in Python, you must be disciplined in applying constructs like asyncio.Lock to prevent subtle race conditions.

5. Performance Trade-offs and Distributed Systems

For an in-process rate limiter, a compiled language like Rust will almost always outperform Python due to the absence of an interpreter and GIL.

Token-Bucket Rate Limiter Benchmarks Across Languages
This benchmark compares various token bucket implementations. Notice the high performance (Sustained GPS, or gets per second) of the Rust `governor` library compared to interpreted languages like Node.js. A similar performance advantage would be expected over Python.

However, in a large-scale distributed system like Bedrock, you wouldn't use an in-memory rate limiter. Each service instance would have its own limiter, leading to incorrect global limits. Instead, you'd use a centralized store like Redis.

Let's briefly examine this superior architecture.

How to Implement Token Bucket Rate Limiting in Python - OneUptime

The OneUptime article also covers a distributed implementation, which is the standard for production-grade systems.

Skim the section Redis-Backed Implementation. The key takeaway is the use of a Lua script. Redis executes Lua scripts atomically. This solves the distributed race condition: the logic to 'get current tokens, refill, check, and consume' all happens as a single, indivisible operation on the Redis server.

In this distributed model, the performance bottleneck shifts from CPU (the rate limit logic) to I/O (the network round-trip to Redis). The performance difference between Python and Rust on the application server becomes negligible because both are spending most of their time waiting for the network. This is a critical trade-off to understand when designing such systems: optimize for the right bottleneck.


Conclusion

In this lesson, we dissected the token bucket algorithm, a cornerstone of traffic management. We've gone from concept to a thread-safe Python implementation and drawn direct parallels to the robust concurrency patterns in Rust/Tokio.

Key Takeaways:

  • The Token Bucket algorithm allows for traffic bursts (up to bucket capacity) while enforcing a sustained average rate (the refill rate).
  • A correct implementation relies on a monotonic clock (time.monotonic or tokio::time::Instant) to be immune to system time changes.
  • In a concurrent environment, the bucket's state must be protected by a lock (threading.Lock, asyncio.Lock, or tokio::sync::Mutex) to prevent race conditions.
  • In distributed systems, the rate limiter state must be centralized (e.g., in Redis), and operations must be atomic (e.g., via a Lua script) to ensure global consistency. In this architecture, network latency often becomes the primary performance bottleneck.

Next Lesson Preview

We've mastered the token bucket. Next, we'll explore another common approach. You will implement a fixed-window counter rate limiter, explaining its memory efficiency and its critical boundary condition issues. This will allow you to start building a mental model of the trade-offs between different rate limiting algorithms.

Can't find a good explanation? Sign up and we'll make it for you

Sign up