如何在Axum中检测断开的SSE客户端并从广播器中移除

嘿，我完全理解你的需求——当SSE客户端断开连接时，你需要及时把对应的mpsc::Sender从广播器的客户端列表中移除，避免无效的客户端占用资源，对吧？这在实时推送场景里是很关键的，我来给你一步步讲解实现方案。

问题分析 #

你当前的代码已经实现了基本的SSE广播功能，但缺少客户端断开后的清理逻辑：当客户端断开，SSE的ReceiverStream会自然结束，但我们没有监听这个事件，所以对应的Sender会一直留在clients列表里，导致后续广播时向无效连接发送消息，还会占用内存。

解决方案步骤 #

1. 给广播器添加移除客户端的方法

我们需要一个方法来从客户端列表中移除指定的Sender，同时顺便清理所有已经断开的客户端，避免内存泄漏：

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 mut inner = self.inner.lock().unwrap();
        // 先清理已断开的客户端，避免向无效连接发送消息
        inner.clients.retain(|client| !client.is_closed());
        let clients = inner.clients.clone();
        drop(inner); // 提前释放锁，避免阻塞其他操作

        let send_futures = clients.iter().map(|client| client.send(event.to_string()));
        let _ = futures::future::join_all(send_futures).await;
    }

    // 新增：移除指定客户端，同时清理无效连接
    pub fn remove_client(&self, tx: &mpsc::Sender<String>) {
        let mut inner = self.inner.lock().unwrap();
        inner.clients.retain(|client| {
            !client.is_closed() && client != tx
        });
    }
}

2. 在SSE处理函数中监听连接断开事件

当客户端建立连接时，我们克隆一份Sender并启动一个异步任务，监听Receiver的关闭事件（也就是客户端断开）。一旦Receiver关闭，就调用清理方法移除对应的客户端：

async fn sse_handler(
    State(app_state): State<Arc<AppState>>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let (tx, rx) = mpsc::channel::<String>(10);
    tx.send("welcome".to_string()).await.unwrap();

    // 克隆广播器和Sender，用于后续清理任务
    let broadcaster_clone = app_state.broadcaster.clone();
    let tx_clone = tx.clone();

    // 将新客户端添加到列表
    app_state.broadcaster.inner.lock().unwrap().clients.push(tx);

    // 启动异步任务：监听Receiver关闭（客户端断开），然后执行清理
    tokio::spawn(async move {
        // 当客户端断开，Receiver会被Drop，closed()会立即完成
        rx.closed().await;
        broadcaster_clone.remove_client(&tx_clone);
    });

    // 构建SSE响应流
    let mystream = ReceiverStream::<String>::new(rx)
        .map(|res| Ok(Event::default().data(res)));

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

为什么这样有效？ #

• 实时清理：当客户端断开连接，Axum会停止处理SSE流，对应的Receiver会被销毁，rx.closed().await会立即触发，清理任务会马上把无效的Sender从列表中移除。
• 双重保障：在broadcast方法中我们也会提前清理已断开的客户端，就算清理任务有延迟，也能避免向无效连接发送消息。
• 性能友好：清理任务是异步执行的，不会阻塞SSE流的处理，也不会影响其他客户端的连接。

完整修改后的代码 #

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 (tx, rx) = mpsc::channel::<String>(10);
    tx.send("welcome".to_string()).await.unwrap();

    let broadcaster_clone = app_state.broadcaster.clone();
    let tx_clone = tx.clone();

    app_state.broadcaster.inner.lock().unwrap().clients.push(tx);

    tokio::spawn(async move {
        rx.closed().await;
        broadcaster_clone.remove_client(&tx_clone);
    });

    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 mut inner = self.inner.lock().unwrap();
        inner.clients.retain(|client| !client.is_closed());
        let clients = inner.clients.clone();
        drop(inner);

        let send_futures = clients.iter().map(|client| client.send(event.to_string()));
        let _ = futures::future::join_all(send_futures).await;
    }

    pub fn remove_client(&self, tx: &mpsc::Sender<String>) {
        let mut inner = self.inner.lock().unwrap();
        inner.clients.retain(|client| {
            !client.is_closed() && client != tx
        });
    }
}

备注：内容来源于stack exchange，提问作者Fred Hors