I want to get the client interrupt signal in the handler, or let the handler execute normally and get the interrupt signal in the response

[Axb12]

Summary

Hi everyone, I am a rust and axum newbie. I recently encountered a problem. I want the handler to be able to complete normally when the client closes the connection during execution, instead of being interrupted directly like a program crash. I hope that when the client is interrupted, I can do some corresponding processing here. Is there any good way to solve this problem?

I saw two related discussions, but they don't quite cover my needs.(#2196) (#2817)

axum version

0.8.4

[mladedav]

If you want to react to the disconnection, the second discussion you linked provides the answer, i.e. you'll want to use a drop guard. If you just want the handler to execute till end even if the response is discarded, you'll want to use tokio::spawn to ensure the part you want to finish always runs even if the spawning future is not polled anymore.

[Axb12]

Thank you very much for your example. I can use it successfully in block responses. But when I use streaming responses, there is a problem. I found that because of the streaming, Axum directly let the request link go through and end normally? It triggered drop, which made it impossible for me to distinguish in stream_handler whether it was a client disconnection or normal streaming.
What should I do?

use std::{
    convert::Infallible,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
};

use axum::{
    body::Body,
    extract::Request,
    http,
    middleware::{self, Next},
    response::Response,
    routing::post,
};
use futures::stream;
use tokio::{
    sync::mpsc,
    time,
};

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt().init();

    let app = axum::Router::new()
        .route("/stream", post(stream_handler))
        .layer(middleware::from_fn(log_middleware));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

struct DropGuard(Arc<AtomicBool>);
impl Drop for DropGuard {
    fn drop(&mut self) {
        self.0.store(true, Ordering::Relaxed);
    }
}

async fn stream_handler(req: Request) -> Response {
    let finished = req.extensions().get::<Arc<AtomicBool>>().cloned().unwrap();
    let (tx, rx) = mpsc::channel(1);
    tokio::spawn(async move {
        // Do stuff
        for i in 0..10 {
            tracing::info!("loop {}", i);
            if let Err(e) = tx.send(format!("{i}")).await {
                tracing::error!("send error: {}", &e);
            };
            if finished.load(Ordering::Relaxed) {
                // React to disconnection
                tracing::warn!("connection is closed");
                break;
            }
            time::sleep(time::Duration::from_millis(500)).await;
        }
        // Do other stuff
    });

    let sse_stream = stream::unfold(rx, |mut catch_stream: mpsc::Receiver<String>| async move {
        match catch_stream.recv().await {
            Some(data) => {
                let event = format!("data: {}\n\n", data);
                Some((Ok::<_, Infallible>(event), catch_stream))
            }
            None => {
                // let event = format!("data: [DONE]\n\n");
                None
            }
        }
    });

    Response::builder()
        .status(http::StatusCode::OK)
        .header(http::header::CONTENT_TYPE, "text/event-stream")
        .body(Body::from_stream(sse_stream))
        .unwrap()
}

pub async fn log_middleware(mut req: Request, next: Next) -> Response {
    let (_tx, mut rx) = mpsc::channel(1);
    let finished = Arc::new(AtomicBool::new(false));
    let _guard = DropGuard(finished.clone());
    let f = finished.clone();
    req.extensions_mut().insert(finished);
    if let Err(e) = tokio::spawn(async move {
        let path = req.uri().path().to_owned();
        let method = req.method().clone();
        let t = time::Instant::now();
        let resp = next.run(req).await;
        let duration = t.elapsed();
        tracing::info!(
            "Request: {} {} in {}ms - Status {}",
            method,
            path,
            duration.as_millis(),
            resp.status()
        );
        if let Err(e) = _tx.send(resp).await {
            if !f.load(Ordering::Relaxed) {
                tracing::error!("log middleware send error: {:?}", &e);
            }
        };
    })
    .await
    {
        tracing::info!("log middleware spawn error: {}", &e);
        return Response::new(Body::empty());
    };

    match rx.recv().await {
        Some(res) => res,
        None => {
            tracing::info!("log middleware recv None");
            return Response::new(Body::empty());
        }
    }
}

[Axb12]

Oh, I found that I can use tx in stream_handler to determine whether the client closes the connection, because if the client closes the connection, tx will be closed, resulting in failure to send.
But in order to make the log in log_middleware more accurate, I had to nest a tokio::spawn, this looks a bit ugly.
The good news is that it looks like I can achieve the effect I originally wanted.

use std::{
    convert::Infallible,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
};

use axum::{
    body::Body,
    extract::Request,
    http,
    middleware::{self, Next},
    response::Response,
    routing::post,
};
use futures::stream;
use tokio::{sync::mpsc, time};

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt().init();

    let app = axum::Router::new()
        .route("/stream", post(stream_handler))
        .layer(middleware::from_fn(log_middleware));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

struct DropGuard(Arc<AtomicBool>);
impl Drop for DropGuard {
    fn drop(&mut self) {
        tracing::info!("DropGuard");
        self.0.store(true, Ordering::Relaxed);
    }
}

async fn stream_handler(req: Request) -> Response {
    let ttx = req.extensions().get::<mpsc::Sender<()>>().cloned().unwrap();
    let (tx, rx) = mpsc::channel(1);
    tokio::spawn(async move {
        // Do stuff
        for i in 0..10 {
            tracing::info!("loop {}", i);
            if let Err(_) = tx.send(format!("{i}")).await {
                tracing::warn!("client is closed");
                break;
            };
            time::sleep(time::Duration::from_millis(500)).await;
        }
        ttx.send(()).await.unwrap();
        // Do other stuff
    });

    let sse_stream = stream::unfold(rx, |mut catch_stream: mpsc::Receiver<String>| async move {
        match catch_stream.recv().await {
            Some(data) => {
                let event = format!("data: {}\n\n", data);
                Some((Ok::<_, Infallible>(event), catch_stream))
            }
            None => {
                // let event = format!("data: [DONE]\n\n");
                None
            }
        }
    });

    Response::builder()
        .status(http::StatusCode::OK)
        .header(http::header::CONTENT_TYPE, "text/event-stream")
        .body(Body::from_stream(sse_stream))
        .unwrap()
}

pub async fn log_middleware(mut req: Request, next: Next) -> Response {
    let (_tx, mut rx) = mpsc::channel(1);
    let (ttx, mut rrx) = mpsc::channel::<()>(1);
    let finished = Arc::new(AtomicBool::new(false));
    let _guard = DropGuard(finished.clone());
    let f = finished.clone();
    req.extensions_mut().insert(finished);
    req.extensions_mut().insert(ttx);
    if let Err(e) = tokio::spawn(async move {
        let path = req.uri().path().to_owned();
        let method = req.method().clone();
        let t = time::Instant::now();
        let resp = next.run(req).await;
        let status_code = resp.status();
        tokio::spawn(async move {
            rrx.recv().await;
            let duration = t.elapsed();
            tracing::info!(
                "Request: {} {} in {}ms - Status {}",
                method,
                path,
                duration.as_millis(),
                status_code,
            );
        });
        if let Err(e) = _tx.send(resp).await {
            if !f.load(Ordering::Relaxed) {
                tracing::error!("spawn error: {:?}", &e);
            }
        };
    })
    .await
    {
        tracing::info!("spawn error: {}", &e);
        return Response::new(Body::empty());
    };

    match rx.recv().await {
        Some(res) => res,
        None => {
            tracing::info!("log middleware recv None");
            return Response::new(Body::empty());
        }
    }
}

[mladedav]

The guard will be dropped when you finish your handler, which happens when you return the response even if it is streamed. You'll need to create a custom body and put the guard in there so that it is dropped when the response body is dropped.

[Axb12]

Thank you very much for your help. But I'm sorry, I'm a rust newbie. I need to learn more about your answer.
Can you help me check the other answer I just posted? It seems to do what I want.

[Axb12]

I found that if I use channels, I don’t need to use

struct DropGuard(Arc<AtomicBool>);
impl Drop for DropGuard {
    fn drop(&mut self) {
        self.0.store(true, Ordering::Relaxed);
    }
}

I only use the channel inside the handler to determine whether the client closes the connection.