A token-bucket rate limiter in 60 lines of Rust

Sign-in endpoints without a rate limit are one of the most common findings in the reports I write. So here’s the limiter I reach for: small enough to read in one sitting, and correct under concurrency.

Why a token bucket#

A fixed window (“100 requests per minute”) has a known weakness. A client can send 100 requests at 12:00:59 and another 100 at 12:01:00, which is 200 requests in two seconds. A token bucket smooths that out. Tokens refill continuously, and each request spends one.

The code#

use std::time::Instant;

pub struct Bucket {
    capacity: f64,
    tokens: f64,
    refill_per_sec: f64,
    last: Instant,
}

impl Bucket {
    pub fn new(capacity: u32, refill_per_sec: f64) -> Self {
        Self { capacity: capacity as f64, tokens: capacity as f64, refill_per_sec, last: Instant::now() }
    }

    pub fn try_take(&mut self) -> bool {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.refill_per_sec).min(self.capacity);
        self.last = now;
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }
}

The highlighted lines are the whole algorithm. Everything else is bookkeeping.

For many keys (one bucket per IP or per account), wrap it in a sharded map so threads don’t fight over a single lock:

use dashmap::DashMap;
use std::net::IpAddr;

pub struct Limiter {
    buckets: DashMap<IpAddr, Bucket>,
    capacity: u32,
    refill: f64,
}

impl Limiter {
    pub fn check(&self, ip: IpAddr) -> bool {
        self.buckets
            .entry(ip)
            .or_insert_with(|| Bucket::new(self.capacity, self.refill))
            .try_take()
    }
}

The Go version, for comparison#

func (b *Bucket) TryTake() bool {
	b.mu.Lock()
	defer b.mu.Unlock()

	now := time.Now()
	b.tokens = min(b.capacity, b.tokens+now.Sub(b.last).Seconds()*b.refill)
	b.last = now
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

Same algorithm, different ergonomics. Go makes the lock explicit. Rust’s &mut self forces you to decide who owns the bucket before it compiles.

Configuration#

limits:
  sign_in:
    capacity: 5 # burst
    refill_per_sec: 0.1 # one attempt every 10s after that
  password_reset:
    capacity: 3
    refill_per_sec: 0.0167

Numbers#

Measured on a laptop, single process, 8 threads hammering 10,000 distinct keys: