Rate Limiting & Load Shedding

[dependencies]
tower_governor = "0.8.0" # depends on tower-http 0.6: a second tower-http compiles alongside 0.7

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use axum::{Router, routing::get};
use tower_governor::{GovernorLayer, governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor};

#[tokio::main]
async fn main() {
    // Bursts of up to 5 requests per client, then one more every 2 seconds.
    let config = Arc::new(
        GovernorConfigBuilder::default()
            .per_second(2)
            .burst_size(5)
            .key_extractor(SmartIpKeyExtractor) // X-Forwarded-For / X-Real-IP / Forwarded, else the peer address
            .use_headers() // x-ratelimit-* headers on every response
            .finish()
            .expect("valid rate-limit configuration"),
    );

    // The limiter keeps one entry per key; purge idle keys periodically.
    let limiter = config.limiter().clone();
    tokio::spawn(async move {
        let mut tick = tokio::time::interval(Duration::from_secs(60));
        loop {
            tick.tick().await;
            limiter.retain_recent();
        }
    });

    let app = Router::new().route("/api/v1/books", get(|| async { "books\n" })).layer(GovernorLayer::new(config));

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    // the peer-IP extractors need the connection's address
    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await.unwrap();
}

$ for i in $(seq 7); do curl -s -o /dev/null -w '%{http_code} ' localhost:3000/api/v1/books; done; echo
200 200 200 200 200 429 429
$ curl -si localhost:3000/api/v1/books | grep -i -E '^HTTP|retry-after|x-ratelimit'
HTTP/1.1 429 Too Many Requests
x-ratelimit-after: 1
retry-after: 1
x-ratelimit-limit: 5
x-ratelimit-remaining: 0

use std::time::Duration;

use axum::{
    BoxError, Router,
    error_handling::HandleErrorLayer,
    http::{StatusCode, header},
    response::IntoResponse,
    routing::get,
};
use tower::ServiceBuilder;

async fn overloaded(err: BoxError) -> impl IntoResponse {
    if err.is::<tower::load_shed::error::Overloaded>() {
        (StatusCode::SERVICE_UNAVAILABLE, [(header::RETRY_AFTER, "1")], "server busy, retry shortly\n")
    } else {
        (StatusCode::INTERNAL_SERVER_ERROR, [(header::RETRY_AFTER, "5")], "internal error\n")
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/api/v1/reports/heavy", get(|| async {
            tokio::time::sleep(Duration::from_millis(500)).await; // simulate expensive work
            "report\n"
        }))
        .layer(
            ServiceBuilder::new()
                .layer(HandleErrorLayer::new(overloaded))
                .load_shed()
                .concurrency_limit(2), // capacity: two concurrent reports per instance
        );
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    // into_make_service() builds the layer stack once, so all connections share one limit (see below)
    axum::serve(listener, app.into_make_service()).await.unwrap();
}

$ for i in 1 2 3 4; do curl -s -o /dev/null -w '%{http_code} ' localhost:3000/api/v1/reports/heavy & done; wait; echo
503 503 200 200

use std::time::{SystemTime, UNIX_EPOCH};

use axum::{
    Router,
    extract::{Request, State},
    http::{HeaderValue, StatusCode, header},
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::get,
};
use redis::aio::ConnectionManager;

const LIMIT: u64 = 3;
const WINDOW_SECS: u64 = 60;

#[derive(Clone)]
struct AppState {
    redis: ConnectionManager,
}

async fn rate_limit(State(state): State<AppState>, req: Request, next: Next) -> Response {
    let client = req
        .headers()
        .get("x-api-key")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("anonymous")
        .to_owned();
    let now = SystemTime::now().duration_since(UNIX_EPOCH).expect("clock after 1970").as_secs();
    let window = now / WINDOW_SECS;
    let key = format!("ratelimit:{client}:{window}");

    let mut redis = state.redis.clone();
    let count: redis::RedisResult<(u64,)> =
        redis::pipe().atomic().incr(&key, 1u64).expire(&key, WINDOW_SECS as i64).ignore().query_async(&mut redis).await;

    match count {
        Ok((n,)) if n > LIMIT => {
            let retry_after = WINDOW_SECS - now % WINDOW_SECS;
            (StatusCode::TOO_MANY_REQUESTS, [(header::RETRY_AFTER, retry_after.to_string())], "rate limit exceeded\n")
                .into_response()
        }
        Ok((n,)) => {
            let mut response = next.run(req).await;
            let remaining = HeaderValue::from(LIMIT - n);
            response.headers_mut().insert("ratelimit-remaining", remaining);
            response
        }
        Err(err) => {
            // fail open: a Redis outage must not take the API down
            tracing::warn!(%err, "rate limiter unavailable");
            next.run(req).await
        }
    }
}

#[tokio::main]
async fn main() {
    let url = std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://127.0.0.1:6379".to_owned());
    let client = redis::Client::open(url).expect("valid Redis URL");
    let state = AppState { redis: ConnectionManager::new(client).await.expect("Redis reachable") };

    let app = Router::new()
        .route("/api/v1/search", get(|| async { "results\n" }))
        .route_layer(middleware::from_fn_with_state(state.clone(), rate_limit))
        .with_state(state);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

$ for i in 1 2 3 4; do curl -s -H 'x-api-key: demo' localhost:3000/api/v1/search; done
results
results
results
rate limit exceeded