429 status code for RateLimitLayer

[kimbauters]

I have a follow-up question from #987 . What is the best way to produce a 429 error when the rate limit is reached?

A hack-y approach is to abuse the HandleErrorLayer for this, and combine it with a TimeoutLayer:

        ServiceBuilder::new()
        .layer(HandleErrorLayer::new(|_err: BoxError| async move {
            (
                StatusCode::TOO_MANY_REQUESTS,
                format!("Too many requests were made today."),
            )
        }))
        .layer(TimeoutLayer::new(Duration::from_secs(5)))
        .layer(BufferLayer::new(1024))
        .layer(RateLimitLayer::new(100, Duration::from_secs(86_400)))

But this doesn't feel quite right. Is there a better way to do this?

[davidpdrsn]

RateLimit will send backpressure up the stack making things wait. To convert that "not ready" state into an error you need LoadShed.

Here is an example:

use std::time::Duration;

use axum::{
    error_handling::HandleErrorLayer,
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::get,
    BoxError, Router,
};
use tower::{
    buffer::BufferLayer,
    limit::RateLimitLayer,
    load_shed::{error::Overloaded, LoadShedLayer},
    ServiceBuilder,
};

#[tokio::main]
async fn main() {
    let app = Router::new().route("/", get(|| async {})).layer(
        ServiceBuilder::new()
            .layer(HandleErrorLayer::new(handle_error))
            .layer(BufferLayer::new(1024))
            .layer(LoadShedLayer::new()) // <- this
            .layer(RateLimitLayer::new(1, Duration::from_secs(5))),
    );

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn handle_error(error: BoxError) -> Response {
    if error.is::<Overloaded>() {
        (StatusCode::TOO_MANY_REQUESTS, "calm down plzz").into_response()
    } else {
        (StatusCode::INTERNAL_SERVER_ERROR, "Something went wrong").into_response()
    }
}

[kimbauters]

Brilliant. That's very clear and works well!