Stateful Axum middleware not compiling if wrapped in a Mutex

Question

I'm trying to implement a basic middleware in a Axum project, and I'm facing an issue when my middleware is passed a state wrapped in a Mutex.

Disclaimer

There is a high chance that the shared state used in this sample is not optimized nor recommended to use as-is in Axum. Above all I'm trying to understand the issue and learn how I could spot the underlying error in my code.

The code

cargo.toml

[package]
name = "mutex_sample"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
axum = "0.6.10"
tokio = {version = "1.26.0", features = ["full"]}

main.rs

use std::{
    net::SocketAddr,
    sync::{Arc, Mutex},
};

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

#[tokio::main]
async fn main() {
    let state_with_mutex = Arc::new(Mutex::new(0));
    let state_without_mutex = Arc::new(0);

    let app: Router = Router::new()
        .route("/", get(hello_world))
        // state without mutex compiles correctly
        .layer(middleware::from_fn_with_state(state_without_mutex, test_middleware))
        // state with mutex : trait bound not satisfied
        .layer(middleware::from_fn_with_state(state_with_mutex, test_middleware_mutex));

    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    println!("Listening on {}", addr);

    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn hello_world() -> String {
    String::from("Hello Reachy !")
}

async fn test_middleware_mutex<B>(
    State(state): State<Arc<Mutex<u16>>>,
    request: Request<B>,
    next: Next<B>,
) -> Response {

    // Commenting this let-else expression fixes the compiler error
    let Ok(inner) = state.lock() else {
        return StatusCode::BAD_REQUEST.into_response();
    };

    next.run(request).await
}

async fn test_middleware<B>(
    State(state): State<Arc<u16>>,
    request: Request<B>,
    next: Next<B>,
) -> Response {

    if *state > 1 {
        return StatusCode::BAD_REQUEST.into_response();
    }

    next.run(request).await
}

The error

the trait bound `axum::middleware::FromFn<fn(State<Arc<std::sync::Mutex<u16>>>, Request<_>, Next<_>) -> impl Future<Output = Response<http_body::combinators::box_body::UnsyncBoxBody<axum::body::Bytes, axum::Error>>> {test_middleware_mutex::<_>}, Arc<std::sync::Mutex<{integer}>>, Route<_>, _>: tower_service::Service<Request<_>>` is not satisfied
the following other types implement trait `tower_service::Service<Request>`:
  axum::middleware::FromFn<F, S, I, (T1, T2)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6, T7)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6, T7, T8)>
  axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6, T7, T8, T9)>
and 8 others

Questions

Two behavior I don't understand here (despite the compiler error message) :

• as stated in the code, if I remove the let-else expression from test_middleware_mutex(), then everything compiles correctly. It seems that the multiple possible return type inside the function causes the compiler to complain.
• for the same type of code (let-else inside the async middleware), wrapping the state in a Mutex also causes an issue.

These two behaviors are mutually exclusive, I can fix the error by removing the Mutex or by removing the let-else expression, but I can't keep both.

Does someone know why is this code not working ? Thanks in advance

Answer 1

As stated in my post's comment section, this is a normal behavior for std::sync::Mutex. It is well documented here : https://docs.rs/tokio/latest/tokio/sync/struct.Mutex.html#which-kind-of-mutex-should-you-use

TL;DR

std::sync::Mutex does not do well with async methods since they cannot be held across Futures. If std::sync::Mutex is used, manipulate it in non-async function. Otherwise, use some async-compatible mutex, such as Tokio's https://docs.rs/tokio/latest/tokio/sync/struct.Mutex.html#