How to know when Sse connection is closed?

[tmpfs]

Hi, thanks a lot for axum, i much prefer it to the other rust web frameworks i have used!

I want to use SSE behind authentication which necessitates storing some state for the channel/stream for each authenticated client.

How would i know when the SSE connection is closed to i can remove stale state? Looking at the examples and docs didn't yield anything, hopefully i am missing something obvious, any pointers much appreciated 🙏

[davidpdrsn]

You can make a type that implements Drop and hold that until the stream closes. Doing that with async-stream is quite easy:

async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    struct Guard {
        // whatever state you need here
    }

    impl Drop for Guard {
        fn drop(&mut self) {
            tracing::info!("stream closed");
        }
    }

    let stream = async_stream::stream! {
        let _guard = Guard {};
        let mut interval = tokio::time::interval(Duration::from_secs(1));
        loop {
            interval.tick().await;
            yield Ok(Event::default().data("hi"));
        }
        // `_guard` is dropped
    };

    Sse::new(stream)
}

[tmpfs]

Thanks @davidpdrsn, that looks perfect 👍

[tmpfs]

Unfortunately @davidpdrsn, this does not appear to work.

I created an MVP to reproduce the issue here: https://github.com/tmpfs/axum-sse-close-mvp.

The client correctly closes the event source but the server does not execute the code after the loop.

Am I doing something wrong or any other ideas on how to detect the stream close event?

Scrap that I wasn't putting the drop in place, sorry!

[davidpdrsn]

You're missing the type that implements Drop. Putting stuff after the loop is not enough on its own.

[brentmjohnson]

Thanks @davidpdrsn! Exactly what I needed for cleanup tasks on stream closure.

[JoelSatkas]

In case anyone might find it useful, the suggested solution didnt apply to me as i was wrapping a receiver in a stream instead of wrapping a loop or tokio task. So i found that you can create a new tokio task to listen and wait for the transmitter to be closed.

pub async fn sse_connection_creation() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    tracing::debug!("Request for new SSE connection");

    let (tx, rx) = mpsc::unbounded_channel();

    // Store your transmitter to be used later.

    handle_client_disconnected(tx);

    let stream = UnboundedReceiverStream::<YourResponseType>::new(rx).map(|res| {
        /* Your logic to create your SSE packets from the data type*/
    });

    Sse::new(stream).keep_alive(
        axum::response::sse::KeepAlive::new()
            .interval(Duration::from_secs(15))
            .text("keep-alive-text"),
    )
}

pub fn handle_client_disconnected(client_transmitter) {
        tokio::spawn(async move {
            // Wait for client to close.
            client_transmitter.closed().await;

            // Your clean up code.
        });
    }

Seems to be working for me, not sure if im not seeing something in how it could not work.