Skip to content
rusthacks.dev
Concurrency

Build a Rate Limiter in Rust With Tokio and a Token Bucket

A token bucket allows bursts up to a capacity and a steady refill rate. The whole algorithm is a timestamp and a float, and the async part is waiting for the next token.

17 min read
Summarize in AI
Flat illustration for the post “Build a Rate Limiter in Rust With Tokio and a Token Bucket”, composed around async.

Rate limiting protects a service from clients that send too much, including by accident: a retry loop with no backoff, a script left running, a crawler with opinions. The token bucket is the most common algorithm because it matches how people think about it, a budget that refills over time. This post builds one with Tokio, runs it, tests it against a simulated clock, finds a rounding problem I did not expect, and then says where to stop writing your own. Everything was compiled and run on Rust 1.99 (edition 2024) with Tokio 1.47.1 on an Apple M1 Max.

TOML
[dependencies]
tokio = { version = "=1.47.1", features = ["full", "test-util"] }

The algorithm

The bucket holds up to capacity tokens. Time adds rate tokens per second, never past capacity. A request takes one token. If none is left, the request is denied, or has to wait. Two numbers describe the policy, and they mean different things. Capacity is the biggest burst you will accept. Rate is the long run average. A bucket of 100 refilling at 10 per second lets a client send 100 requests at once after a quiet minute, and then 10 a second for as long as it keeps going.

No background task is needed to refill. When a request arrives, compute elapsed * rate since the last call, add it, clamp it to capacity and subtract one. State is two numbers and a timestamp. Doing it lazily removes a timer per bucket, which matters the moment there is one bucket per user.

The bucket

rust
use std::time::Duration;
use tokio::time::Instant;

pub struct TokenBucket {
    capacity: f64,
    tokens: f64,
    rate: f64,
    last: Instant,
}

impl TokenBucket {
    pub fn new(capacity: u32, per_second: f64) -> Self {
        Self {
            capacity: f64::from(capacity),
            tokens: f64::from(capacity),
            rate: per_second,
            last: Instant::now(),
        }
    }

    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.rate).min(self.capacity);
        self.last = now;
    }

    pub fn try_acquire(&mut self) -> bool {
        self.refill();
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }

    pub fn wait_time(&mut self) -> Duration {
        self.refill();
        if self.tokens >= 1.0 {
            Duration::ZERO
        } else {
            Duration::from_secs_f64((1.0 - self.tokens) / self.rate)
        }
    }
}

Read it from the constructor. A new bucket starts full, so a fresh client is allowed its burst. refill is the heart: it measures how long it has been since the last call, converts that to tokens, adds them with a ceiling at capacity, and moves last forward. Clamping with min is what stops a bucket that has been idle for an hour from banking an hour of tokens. try_acquire refills, checks that at least one whole token is available, and takes it. wait_time answers the other question a caller has: how long until a token exists. If the bucket has 0.25 tokens at 2 tokens per second, the wait is 0.75 / 2 = 375 ms.

try_acquire is what an HTTP handler wants: allow, or reject with 429 Too Many Requests, and put the number of seconds from wait_time in a Retry-After header. A background job or a client library wants wait_time directly, to sleep rather than fail.

Two choices deserve a comment. The type is tokio::time::Instant rather than std::time::Instant, and that is what makes the tests below possible: Tokio's clock can be paused and moved by hand. And tokens are an f64, because a fractional refill is the normal case. A rate of 5 per second checked every 10 milliseconds adds 0.05 of a token per check. The next section shows what that costs.

Testing without waiting

A rate limiter is a time machine, so test it with a clock you control. With the test-util feature, start_paused = true freezes Tokio's time, and advance moves it forward by an exact amount. The tests run in microseconds and never flake (testing time without sleeping covers the other approach, injecting a clock of your own).

rust
#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test(start_paused = true)]
    async fn refills_over_time() {
        let mut b = TokenBucket::new(2, 1.0);
        assert!(b.try_acquire());
        assert!(b.try_acquire());
        assert!(!b.try_acquire());
        tokio::time::advance(Duration::from_millis(1100)).await;
        assert!(b.try_acquire());
        assert!(!b.try_acquire());
    }

    #[tokio::test(start_paused = true)]
    async fn long_run_average_is_the_rate() {
        let mut b = TokenBucket::new(5, 5.0);
        let mut admitted = 0;
        for _ in 0..1000 {
            if b.try_acquire() {
                admitted += 1;
            }
            tokio::time::advance(Duration::from_millis(10)).await;
        }
        // 10 simulated seconds: the burst of 5 plus 5 per second, less at
        // most one token lost to floating point rounding at a tick boundary.
        assert!((54..=55).contains(&admitted), "admitted {admitted}");
    }

    #[tokio::test(start_paused = true)]
    async fn idle_time_does_not_bank_past_capacity() {
        let mut b = TokenBucket::new(3, 10.0);
        tokio::time::advance(Duration::from_secs(3600)).await;
        let burst = (0..10).filter(|_| b.try_acquire()).count();
        assert_eq!(burst, 3);
    }
}
Text
running 3 tests
test tests::idle_time_does_not_bank_past_capacity ... ok
test tests::refills_over_time ... ok
test tests::long_run_average_is_the_rate ... ok

test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out

Each test pins a different property. refills_over_time drains a bucket of two, sees the third request refused, advances 1.1 seconds and sees exactly one more allowed, not two. idle_time_does_not_bank_past_capacity advances an hour and asserts that only three requests, the capacity, get through. long_run_average_is_the_rate calls the limiter every 10 simulated milliseconds for 10 simulated seconds and counts the admitted calls. The expected total is the burst of 5 plus 5 per second for 10 seconds, which is 55.

That test is not the one I first wrote. My first version asserted exactly 55, and it failed:

Text
assertion `left == right` failed
  left: 54
 right: 55

The bucket admitted 54. The cause is floating point. The refill adds 0.01 * 5.0, which is not exactly 0.05 in binary, and across hundreds of additions the sum can land at 0.9999999 instead of 1.0 at a tick where the exact answer is 1. The check tokens >= 1.0 then fails by a hair and the request is refused one tick early, and the token is spent on the next tick instead. No token is lost for good, but the limiter can be one request behind an ideal, exact counter at any instant. For rate limiting that is irrelevant. For a test, it means asserting a range, 54..=55, and writing down why. If exactness matters, track time in integer nanoseconds and use the approach the governor crate takes, below.

Sharing it and waiting

A limiter is shared by many tasks, which is where the first compile error lives. Here is the natural first attempt: put the bucket in an Arc and call it from a thread.

rust
use std::sync::Arc;

struct TokenBucket {
    tokens: f64,
}

impl TokenBucket {
    fn try_acquire(&mut self) -> bool {
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }
}

fn main() {
    let shared = Arc::new(TokenBucket { tokens: 2.0 });
    let for_task = Arc::clone(&shared);
    let handle = std::thread::spawn(move || for_task.try_acquire());
    println!("{}", handle.join().unwrap());
}
Text
error[E0596]: cannot borrow data in an `Arc` as mutable
  --> src/bin/r_fail.rs:21:45
   |
21 |     let handle = std::thread::spawn(move || for_task.try_acquire());
   |                                             ^^^^^^^^ cannot borrow as mutable
   |
   = help: trait `DerefMut` is required to modify through a dereference, but it is not implemented for `Arc<TokenBucket>`

This is not a problem with the bucket. Arc hands out shared references only, because several owners may exist at once, and try_acquire takes &mut self. Rust refuses to let you mutate through a shared handle without a type that arbitrates, which is the job of a Mutex. The error message is accurate: Arc<TokenBucket> does not implement DerefMut. The fix is Arc<Mutex<TokenBucket>> (sharing state between tasks compares the options), which the next block uses.

rust
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::Instant;

#[derive(Clone)]
pub struct Limiter(Arc<Mutex<TokenBucket>>);

impl Limiter {
    pub fn new(capacity: u32, per_second: f64) -> Self {
        Self(Arc::new(Mutex::new(TokenBucket::new(capacity, per_second))))
    }

    pub async fn acquire(&self) {
        loop {
            let wait = {
                let mut bucket = self.0.lock().await;
                if bucket.try_acquire() {
                    return;
                }
                bucket.wait_time()
            };
            tokio::time::sleep(wait).await;
        }
    }
}

Limiter is a cheap clone around the shared bucket. acquire is the waiting version. It locks, tries to take a token, and if that fails asks for the wait time. The lock is released at the end of the inner block, before the sleep. Holding it across sleep would block every other caller for the whole wait and serialise the very thing the limiter is meant to spread out (mutex guard across await). After sleeping, the loop checks again instead of assuming a token is waiting, because another task may have taken it, or float rounding may have woken it a hair early. Six requests arrive at once at a limiter of capacity 2 and rate 2 per second:

rust
#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
    let limiter = Limiter::new(2, 2.0);
    let start = Instant::now();
    let mut tasks = Vec::new();
    for id in 0..6 {
        let l = limiter.clone();
        tasks.push(tokio::spawn(async move {
            l.acquire().await;
            println!("request {id} admitted at {:>4} ms", start.elapsed().as_millis());
        }));
    }
    for t in tasks {
        t.await.unwrap();
    }
}
Text
request 0 admitted at    0 ms
request 1 admitted at    0 ms
request 2 admitted at  500 ms
request 3 admitted at 1000 ms
request 4 admitted at 1500 ms
request 5 admitted at 2000 ms

The first two get the burst and are admitted at 0 ms. Then one request every 500 ms, which is 2 per second, as configured. With the paused clock the timestamps are exact, which is the other reason to use virtual time in examples. One caveat to read off the output: the order is the order in which tasks happened to wake, not the order they arrived in. This is the fairness limit listed below.

Per key limits

A global limit stops a flood. A per client limit stops one client from using the whole budget, and that needs one bucket per key: IP address, API key, user id. Keep a map from key to bucket, create buckets on demand and, importantly, forget the ones that no longer matter. Without eviction, memory grows with every key ever seen, and an attacker rotating addresses will find that out for you.

rust
struct Keyed {
    capacity: u32,
    rate: f64,
    buckets: Mutex<HashMap<String, (TokenBucket, Instant)>>,
}

impl Keyed {
    fn new(capacity: u32, rate: f64) -> Self {
        Self { capacity, rate, buckets: Mutex::new(HashMap::new()) }
    }

    fn check(&self, key: &str) -> bool {
        let mut map = self.buckets.lock().unwrap();
        let (bucket, seen) = map
            .entry(key.to_owned())
            .or_insert_with(|| (TokenBucket::new(self.capacity, self.rate), Instant::now()));
        *seen = Instant::now();
        bucket.try_acquire()
    }

    // A key idle for longer than a full refill holds a full bucket, which is
    // indistinguishable from a key never seen, so forgetting it changes nothing.
    fn evict_idle(&self) -> usize {
        let full_refill = Duration::from_secs_f64(f64::from(self.capacity) / self.rate);
        let mut map = self.buckets.lock().unwrap();
        let before = map.len();
        map.retain(|_, (_, seen)| seen.elapsed() < full_refill);
        before - map.len()
    }

    fn len(&self) -> usize {
        self.buckets.lock().unwrap().len()
    }
}

The eviction rule is worth a second look. A bucket that has been idle for capacity / rate seconds has refilled completely, so its state is identical to a brand new bucket's, and removing it changes no behaviour. That makes eviction safe at any time. The lock here is a std mutex, held for a map lookup and never across an await, which is the correct choice for a short critical section.

rust
#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
    let limiter = Keyed::new(2, 1.0);
    for key in ["10.0.0.1", "10.0.0.1", "10.0.0.1", "10.0.0.2"] {
        println!("{key}: {}", if limiter.check(key) { "allowed" } else { "429" });
    }
    println!("tracked keys: {}", limiter.len());
    tokio::time::advance(Duration::from_secs(3)).await;
    println!("evicted: {}", limiter.evict_idle());
    println!("tracked keys: {}", limiter.len());
}
Text
10.0.0.1: allowed
10.0.0.1: allowed
10.0.0.1: 429
10.0.0.2: allowed
tracked keys: 2
evicted: 2
tracked keys: 0

With capacity 2 and 1 token per second, the first address got two requests through and a 429 on the third, the second address was unaffected, and the map held two entries. Three simulated seconds later (more than the two seconds a full refill takes), eviction removed both. In production, run it from a periodic task, or sweep a few entries on each insert so there is no cliff.

What the lock costs

A single mutex around every check sounds like a bottleneck, so I measured it. Each of N tasks calls try_acquire one million times on a shared bucket, with a bucket large enough that every call succeeds, in a release build on eight worker threads. I compared std::sync::Mutex against tokio::sync::Mutex.

rust
#[path = "r_bucket.rs"]
#[allow(dead_code)]
mod bucket;
use bucket::TokenBucket;
use std::sync::Arc;
use std::time::Instant;

#[tokio::main(flavor = "multi_thread", worker_threads = 8)]
async fn main() {
    const PER_TASK: u32 = 1_000_000;
    for tasks in [1usize, 8] {
        let std_b = Arc::new(std::sync::Mutex::new(TokenBucket::new(1_000_000, 1e9)));
        let t = Instant::now();
        let hs: Vec<_> = (0..tasks).map(|_| {
            let b = Arc::clone(&std_b);
            tokio::spawn(async move {
                for _ in 0..PER_TASK { let _ = b.lock().unwrap().try_acquire(); }
            })
        }).collect();
        for h in hs { h.await.unwrap(); }
        let std_t = t.elapsed();

        let tok_b = Arc::new(tokio::sync::Mutex::new(TokenBucket::new(1_000_000, 1e9)));
        let t = Instant::now();
        let hs: Vec<_> = (0..tasks).map(|_| {
            let b = Arc::clone(&tok_b);
            tokio::spawn(async move {
                for _ in 0..PER_TASK { let _ = b.lock().await.try_acquire(); }
            })
        }).collect();
        for h in hs { h.await.unwrap(); }
        let tok_t = t.elapsed();
        let n = f64::from(PER_TASK) * tasks as f64;
        println!("{tasks} task(s): std mutex {:>5.0} ns/call, tokio mutex {:>5.0} ns/call",
            std_t.as_nanos() as f64 / n, tok_t.as_nanos() as f64 / n);
    }
}
Text
1 task(s): std mutex    66 ns/call, tokio mutex    60 ns/call
8 task(s): std mutex    97 ns/call, tokio mutex   156 ns/call

One task alone costs 60 to 66 ns per call, so either mutex is cheap; the work is a clock read and a few floating point operations. With eight tasks fighting for it, the std mutex rose to 97 ns per call while the Tokio mutex rose to 156 ns. A std mutex wrapped around a tiny critical section is the better choice, and it is what the Tokio documentation recommends when no await happens inside. At about 100 ns per call, one bucket sustains on the order of ten million checks a second, so the lock is not your limit at any realistic request rate. It becomes visible only when every core hammers one bucket, and then the answer is sharding keys across several buckets or an atomic design, not a faster lock. Note that Limiter::acquire above uses the Tokio mutex for simplicity; swapping in the std one is a safe change because the guard is gone before the sleep.

Choosing capacity and rate

The two numbers come from different questions, and mixing them is the usual error. The rate answers: what is the most sustained load I am willing to accept from one client? Take it from capacity planning, not from what clients usually do. If one instance handles 2,000 requests a second and you want to serve 200 clients fairly, 10 per second per client is where the arithmetic starts. The capacity answers: how bursty are legitimate clients? A web page that fires eight API calls on load needs a capacity of at least eight, or every page view fails its last request. A bucket of 10 at 10 per second allows that page load once per second, sustained.

A useful sanity check is the time to recover: capacity / rate seconds from empty to full. At 10 and 10 per second that is a second, so a client that used its burst gets it back after waiting one second. At 600 and 10 per second the recovery takes a minute, which is a very different promise: a client can send ten minutes' worth of steady traffic in a flash, then must go slow for a minute. The same rate with a larger capacity is a looser policy, and it is the knob to turn when users complain about false positives.

Different routes deserve different buckets. Login attempts want a tiny capacity and a slow refill, since a person types a password once every few seconds and a guesser wants thousands. A search endpoint can take a bigger burst. Key those limiters differently, and apply the strict one before the expensive work, not after.

Telling the client what happened

A rejection should carry enough information to behave well. Return status 429 with a Retry-After header holding wait_time rounded up to whole seconds. Rounding up matters: a client told to retry in 0.4 seconds that retries at 0.4 seconds, with a bucket that has 0.99 of a token, is refused again. Many HTTP clients honour Retry-After automatically, and an API that sends it turns a stampede of retries into an orderly queue.

Many APIs also send the remaining budget on every response, with headers such as RateLimit-Remaining, so well behaved clients throttle themselves before they are refused. The bucket gives you the number for free: it is tokens.floor() after the call. Do this for authenticated clients; for anonymous traffic, the less you reveal about your limits the better.

GCRA, the same thing without the bucket

The Generic Cell Rate Algorithm that governor uses is worth understanding, because it removes the part of this code that caused trouble. Instead of a token count, it stores one timestamp per key: the theoretical arrival time, the moment at which the next request would be exactly on schedule. With a rate of 5 per second the spacing between requests is 200 ms. A request arriving at time t is allowed if t is no earlier than TAT - burst_tolerance, where the tolerance is the spacing times (capacity - 1). When allowed, the new TAT becomes max(TAT, t) + spacing.

Work through capacity 2 at 2 per second: spacing is 500 ms and tolerance is 500 ms. The first request at t = 0 is allowed and moves TAT to 500. The second at t = 0 passes the test because 0 is not earlier than 500 - 500, and moves TAT to 1,000. The third at t = 0 fails, since 0 is earlier than 1,000 - 500. It is allowed at 500, as in the output of the limiter run above. The behaviour matches the bucket exactly, but the state is a single integer of nanoseconds, so there is no floating point drift, no separate "last" field, and an atomic compare and swap can update it without a lock. That is why the crate is lock free, and why a hand written version is a learning tool and a library is the product.

Rate limits and concurrency limits

A rate limit bounds requests per unit time. A concurrency limit bounds requests in flight. They fail differently, and the one you need depends on what is scarce. If each request costs a fixed amount of CPU, a rate limit tracks cost well. If requests vary from 5 ms to 30 seconds, a rate limit lets a client fire its allowance of slow requests all at once and tie up every worker, while a concurrency limit would have stopped it at the cap. Services that talk to a slow downstream usually want both: a token bucket at the edge to shape arrival, and a Semaphore in front of the downstream to bound what is outstanding (limiting concurrency with a semaphore).

The two also compose in the same function. Take a token, then take a permit, and release the permit when the request finishes. Take them in that order: waiting for a permit while holding a token spends budget on requests that are doing nothing yet.

Limits of this design

  • It is per process. With three instances behind a load balancer the effective limit is three times higher, and a client that reconnects to a different instance sees a fresh bucket. A shared limit needs shared state: Redis with an atomic script, or limits enforced at the gateway (distributed rate limiting).
  • It restarts empty. Buckets live in memory, so a deploy resets everyone to a full burst. Usually that is fine; for abuse limits an attacker can time requests to your deploys.
  • Waiters are not served in order. Tasks sleeping in acquire race for the next token, as the output above showed. For strict fairness, use a queue, or tokio::sync::Semaphore, which is fair, with a task that adds permits at the refill rate.
  • The key is only as good as what you key on. An IP address is shared behind a NAT and free to an attacker with a botnet. An authenticated identity is better, applied after authentication (rate limiting for abuse).
  • Burst equals capacity. If a downstream service cannot take a burst, you want a smoothing limiter, which allows no more than one request per interval, or a concurrency limit rather than a rate limit.

The fixed window counter, "100 requests per minute, reset on the minute", is simpler and worse. A client can send 100 requests at the end of one window and 100 at the start of the next, 200 within a second, which is twice the intended rate. A token bucket has no boundary to exploit: the most it can ever send in any interval of length t is capacity + rate * t.

In production

Use the governor crate rather than maintaining this code. It implements the Generic Cell Rate Algorithm, a different way of describing the same behaviour that stores a single timestamp per key and updates it with a compare and swap, so there is no lock and no floating point. It supplies keyed limiters, jitter and async waiting. For HTTP services, Tower layers and the tower_governor crate wire it into Axum (rate limiting in Axum). Writing the bucket yourself is still worth doing once: the tests above are the documentation of what a limiter has to guarantee, and they are how you will check that a library does what you think. A limiter you cannot test deterministically is one you are only guessing about (testing rate limiters goes further with property tests).

Share

Keep reading

Get new posts by email

We send one email when a new post is published. No spam.

Unsubscribe in one click. No spam, ever.