How to detect an a dropped SSE (server sent event) client?

[frederikhors]

In the below code I'm handling sse connections within the function sse_handler(). Is there a way to know when that connection is closed so I can remove the client from Bradcaster.inner.clients?

@davidpdrsn, I tried to understand how to use the Drop trait technique you suggested here, with no success.

REPL: https://www.rustexplorer.com/b/xkqyt2

/*
[dependencies]
axum = { version = "0.6.20" }
futures = { version = "0.3.28", default-features = false }
tokio = { version = "1.32.0", default-features = false, features = [
    "macros",
    "process",
    "rt-multi-thread",
] }
tokio-stream = { version = "0.1.14", default-features = false, features = [
    "sync",
] }
*/

use axum::{
    extract::State,
    response::{
        sse::{Event, KeepAlive, Sse},
        Html,
    },
    routing::get,
    Router,
};
use futures::stream::Stream;
use tokio_stream::{wrappers::ReceiverStream, StreamExt};
use std::{
    convert::Infallible,
    sync::{Arc, Mutex},
};
use tokio::sync::mpsc;

struct BroadcasterInner {
    clients: Vec<mpsc::Sender<String>>,
}

pub struct Broadcaster {
    inner: Mutex<BroadcasterInner>,
}

struct AppState {
    broadcaster: Arc<Broadcaster>,
}

#[tokio::main]
async fn main() {
    let broadcaster = Broadcaster::new();

    let app_state = Arc::new(AppState { broadcaster });

    let app = Router::new()
        .route("/send_message", get(send_message))
        .route("/sse", get(sse_handler))
        .with_state(app_state);

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

async fn send_message(State(app_state): State<Arc<AppState>>) -> Html<&'static str> {
    app_state.broadcaster.broadcast("message").await;

    Html("Message sent")
}

async fn sse_handler(
    State(app_state): State<Arc<AppState>>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let rx = app_state.broadcaster.add_client().await;

    let mystream = ReceiverStream::<String>::new(rx).map(|res| Ok(Event::default().data(res)));

    Sse::new(mystream).keep_alive(KeepAlive::default())
}

impl Broadcaster {
    pub fn new() -> Arc<Self> {
        Arc::new(Broadcaster {
            inner: Mutex::new(BroadcasterInner {
                clients: Vec::new(),
            }),
        })
    }

    pub async fn add_client(&self) -> mpsc::Receiver<String> {
        let (tx, rx) = mpsc::channel::<String>(10);

        tx.send("welcome".to_string()).await.unwrap();

        self.inner.lock().unwrap().clients.push(tx);

        rx
    }

    pub async fn broadcast(&self, event: &str) {
        let clients = self.inner.lock().unwrap().clients.clone();

        let send_futures = clients.iter().map(|client| client.send(event.to_string()));

        let _ = futures::future::join_all(send_futures).await;
    }
}

[davidpdrsn]

Using Drop should work. Can you share your code using that?

[frederikhors]

This is the issue: I don't know where to write the Drop.

[zoer]

@frederikhors take a look on this comment #1060 (comment)

[frederikhors]

@zoer, yeah. I know it.

I do not know how to transform this:

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
};

using my code:

pub async fn add_client(&self) -> mpsc::Receiver<String> {
    let (tx, rx) = mpsc::channel::<String>(10);

    tx.send("welcome".to_string()).await.unwrap();

    self.inner.lock().unwrap().clients.push(tx);

    rx
}

pub async fn broadcast(&self, event: &str) {
    let clients = self.inner.lock().unwrap().clients.clone();

    let send_futures = clients.iter().map(|client| client.send(event.to_string()));

    let _ = futures::future::join_all(send_futures).await;
}

[zoer]

I think you chose the wrong channel type for your needs, if I were you I would switch to use tokio::sync::broadcast channel, where you have one sender and multiple receivers, so you would not have problem with removing sender each time when receiver is gone. Eventually that would look something like this:

async fn sse_handler(
    State(app_state): State<Arc<AppState>>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let rx = app_state.broadcaster.add_client().await;

    let mut mystream = ReceiverStream::<String>::new(rx).map(|res| Ok(Event::default().data(res)));
    
    let stream = async_stream::stream! {
        yield Ok(Event::default().data("welcome"));
        loop {
            yield stream.next().await.unwrap()
        }
    };

    Sse::new(mystream).keep_alive(KeepAlive::default())
}

[frederikhors]

Ok. First of all: THANKS!

Then:

1. I would like to use broadcast too and I should use tokio_stream::wrappers::BroadcastStream and not tokio_stream::wrappers::ReceiverStream, right?

2. You said:

so you would not have problem with removing sender each time when receiver is gone

What do you mean? Is the code you posted? Or should I use broadcst's .subscribe() method?

[zoer]

Here's your example adopted to the new approach:

[dependencies]
axum = { version = "0.6.20" }
futures = { version = "0.3.28", default-features = false }
tokio = { version = "1.32.0", default-features = false, features = [
    "macros",
    "process",
    "rt-multi-thread",
] }
tokio-stream = { version = "0.1.14", default-features = false, features = [
    "sync",
] }
async-stream = "0.3.5"

use axum::{
    extract::State,
    response::{
        sse::{Event, KeepAlive, Sse},
        Html,
    },
    routing::get,
    Router,
};
use futures::stream::Stream;
use std::{convert::Infallible, sync::Arc};
use tokio::sync::broadcast;

pub struct Broadcaster {
    fanout: broadcast::Sender<String>,
}

struct AppState {
    broadcaster: Arc<Broadcaster>,
}

#[tokio::main]
async fn main() {
    let broadcaster = Broadcaster::new();

    let app_state = Arc::new(AppState { broadcaster });

    let app = Router::new()
        .route("/send_message", get(send_message))
        .route("/sse", get(sse_handler))
        .with_state(app_state);

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

async fn send_message(State(app_state): State<Arc<AppState>>) -> Html<&'static str> {
    app_state.broadcaster.broadcast("message");

    Html("Message sent")
}

async fn sse_handler(
    State(app_state): State<Arc<AppState>>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let mut rx = app_state.broadcaster.add_client();

    let stream = async_stream::stream! {
        yield Ok(Event::default().data("welcome"));
        loop {
            let msg = rx.recv().await.unwrap();
            yield Ok(Event::default().data(&msg));
        }
    };

    Sse::new(stream).keep_alive(KeepAlive::default())
}

impl Broadcaster {
    pub fn new() -> Arc<Self> {
        let (tx, _) = broadcast::channel(16);
        Arc::new(Broadcaster { fanout: tx })
    }

    pub fn add_client(&self) -> broadcast::Receiver<String> {
        self.fanout.subscribe()
    }

    pub fn broadcast(&self, event: &str) {
        self.fanout
            .send(event.to_string())
            .expect("to have at least one receiver");
    }
}

[frederikhors]

Thanks.

1. yield Ok(Event::default().data("welcome")); send that message to all clients, even those who are already subscribed and to whom it makes no sense to send "welcome" again, right?

2. why the usage of async_stream::stream!? This is another crate to add to Cargo.toml. Can I use tokio_stream::wrappers::BroadcastStream instead like the below?

   let stream = BroadcastStream::<String>::new(rx).map(|msg| Ok(Event::default().data(msg.to_string())));

[zoer]

1. No, each client will only have its own "welcome" message.
2. You can go with BroadcastStream, but in this case to send "welcome" message you will need to make additional channel for each client to be able to send "welcome" message first and then pipe messages from your broadcast receiver.

[jplatte]

An alternative to using the stream! macro would be roughly stream::once(async { Ok(Event::default().data("welcome")) }).chain(broadcast_receiver_stream).

That would also create a stream for each client that first yields the welcome string, then everything from the broadcast subscriber.