use dashmap::DashMap;
use std::time::Instant;
pub struct TokenBucket {
tokens: f64,
last_refill: Instant,
}
impl TokenBucket {
fn new(capacity: f64) -> Self {
TokenBucket { tokens: capacity, last_refill: Instant::now() }
}
fn try_consume(&mut self, capacity: f64, refill_per_sec: f64) -> bool {
let now = Instant::now();
let elapsed = now.duration_since(self.last_refill).as_secs_f64();
self.tokens = (self.tokens + elapsed * refill_per_sec).min(capacity);
self.last_refill = now;
if self.tokens >= 1.0 {
self.tokens -= 1.0;
true
} else {
false
}
}
}
#[derive(Clone)]
pub struct RateLimiter {
buckets: std::sync::Arc<DashMap<String, TokenBucket>>,
capacity: f64,
refill_per_sec: f64,
}
impl RateLimiter {
pub fn new(capacity: f64, refill_per_sec: f64) -> Self {
RateLimiter {
buckets: std::sync::Arc::new(DashMap::new()),
capacity,
refill_per_sec,
}
}
pub fn check(&self, key: &str) -> bool {
let mut entry = self
.buckets
.entry(key.to_owned())
.or_insert_with(|| TokenBucket::new(self.capacity));
entry.try_consume(self.capacity, self.refill_per_sec)
}
}
use axum::{
extract::{ConnectInfo, State},
http::{header, Request, StatusCode},
middleware::Next,
response::{IntoResponse, Response},
};
use std::net::SocketAddr;
use crate::token_bucket::RateLimiter;
pub async fn rate_limit<B>(
State(limiter): State<RateLimiter>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
req: Request<B>,
next: Next<B>,
) -> Result<Response, Response> {
let key = addr.ip().to_string();
if limiter.check(&key) {
Ok(next.run(req).await)
} else {
let body = "rate limit exceeded";
let resp = (
StatusCode::TOO_MANY_REQUESTS,
[(header::RETRY_AFTER, "1")],
body,
)
.into_response();
Err(resp)
}
}
mod middleware;
mod token_bucket;
use axum::{routing::get, Router};
use std::net::SocketAddr;
use crate::token_bucket::RateLimiter;
async fn hello() -> &'static str {
"ok"
}
#[tokio::main]
async fn main() {
// 10 tokens capacity, refilling 5 per second per client IP.
let limiter = RateLimiter::new(10.0, 5.0);
let app = Router::new()
.route("/", get(hello))
.layer(axum::middleware::from_fn_with_state(
limiter.clone(),
middleware::rate_limit,
))
.with_state(limiter);
let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.await
.unwrap();
}
This snippet shows a per-client rate limiter for an axum service built on the token-bucket algorithm. The bucket refills continuously at a fixed rate and each request consumes one token; when the bucket is empty the request is rejected with 429 Too Many Requests. This model is preferable to a fixed-window counter because it smooths bursts naturally and never suffers the boundary spike where two windows meet.
In token_bucket.rs, TokenBucket stores tokens as an f64 plus the last_refill instant. try_consume computes elapsed time lazily on each call, adds elapsed * refill_per_sec tokens capped at capacity, and only then checks whether one token is available. Computing refill lazily avoids a background timer task per client — the state is cheap and only advances when the client is actually seen. RateLimiter wraps a DashMap<String, TokenBucket> so buckets for different keys can be read and mutated concurrently without a global lock; DashMap shards internally, and entry().or_insert_with() makes first-touch creation atomic per key.
In middleware.rs, rate_limit is an axum::middleware::from_fn_with_state handler. It derives a client key from the connection's SocketAddr (a real deployment behind a proxy would prefer a trusted X-Forwarded-For or an API key), calls limiter.check(&key), and either forwards the request via next.run(req).await or short-circuits with a 429. The rejection response sets a Retry-After header so well-behaved clients back off. Returning Result<Response, StatusCode> keeps the handler ergonomic while still producing a proper HTTP response.
In main.rs, the limiter is constructed once, wrapped in shared state, and attached with .layer(...) so every route inherits the limit. Because the middleware runs before the handler, rejected requests never touch business logic. into_make_service_with_connect_info is required so the peer SocketAddr is available for keying.
The main trade-off is that this state is per-process: horizontal scaling needs a shared store like Redis instead of DashMap. It also grows memory with distinct keys, so a production version would evict idle buckets. For a single instance, however, this approach is fast, lock-light, and simple to reason about.
Related snips
export interface RetryOptions {
retries: number;
baseMs: number;
maxMs: number;
signal?: AbortSignal;
onRetry?: (attempt: number, delay: number, err: unknown) => void;Exponential backoff with jitter for retries
struct Config<'a> {
name: &'a str,
value: &'a str,
}
fn parse_config(line: &str) -> Config {Lifetime annotations for flexible borrowing in structs
use crossbeam::channel::unbounded;
use std::thread;
fn main() {
let (tx, rx) = unbounded();Crossbeam for advanced concurrent data structures
// Creating a Promise
const myPromise = new Promise((resolve, reject) => {
const success = true;
setTimeout(() => {
if (success) {Promises and async/await patterns for asynchronous JavaScript
use tracing::{info, instrument};
#[instrument]
fn process_request(user_id: u64) {
info!(user_id, "Processing request");
// Work happens heretracing for structured logging and distributed tracing
use std::sync::mpsc;
use std::thread;
fn main() {
let (tx, rx) = mpsc::channel();Channels (mpsc) for message passing between threads
Share this code
Here's the card — post it anywhere.
