如何用Rocket-WS编写消息转发至所有客户端的Demo

可行实现方案 #

核心思路是利用Tokio无界通道管理所有客户端的消息发送器，避免直接存储WebSocket流（流是独占资源，无法被多线程安全访问）。每个客户端连接时创建一个通道发送器，存入全局集合；收到消息时遍历所有发送器转发，客户端断开时自动清理无效发送器。

完整代码示例 #

Cargo.toml 依赖配置

[package]
name = "rocket-ws-broadcast"
version = "0.1.0"
edition = "2021"

[dependencies]
rocket = { version = "0.5.0-rc.3", features = ["json"] }
rocket_ws = "0.5.0-rc.3"
tokio = { version = "1.35.0", features = ["full"] }
futures = "0.3.29"
serde = { version = "1.0.193", features = ["derive"] }
serde_json = "1.0.108"

主程序代码

use rocket::{get, routes, State};
use rocket_ws::{ws, Message};
use tokio::sync::{mpsc, Mutex};
use std::sync::Arc;

// 定义全局共享的客户端发送器集合类型
type ClientSender = mpsc::UnboundedSender<Message>;
type SharedClients = Arc<Mutex<Vec<ClientSender>>>;

#[get("/ws")]
async fn ws_broadcast(ws: ws::WebSocket, clients: &State<SharedClients>) -> ws::Channel<'static> {
    ws.channel(move |mut stream| {
        // 为当前客户端创建消息通道
        let (client_tx, mut client_rx) = mpsc::unbounded_channel();
        // 将发送器加入全局集合
        clients.lock().await.push(client_tx.clone());

        Box::pin(async move {
            // 任务1：读取客户端发来的消息，转发给所有在线客户端
            let read_task = async {
                while let Some(msg_result) = stream.next().await {
                    match msg_result {
                        Ok(msg) => {
                            let mut clients_lock = clients.lock().await;
                            // 遍历所有发送器转发消息，同时清理已断开的客户端
                            clients_lock.retain(|sender| sender.send(msg.clone()).is_ok());
                        }
                        Err(e) => eprintln!("客户端消息读取失败: {}", e),
                    }
                }
            };

            // 任务2：监听通道消息，发送给当前客户端
            let write_task = async {
                while let Some(msg) = client_rx.recv().await {
                    if stream.send(msg).await.is_err() {
                        break;
                    }
                }
            };

            // 同时运行两个任务，任意一个结束则终止
            tokio::select! {
                _ = read_task => {},
                _ = write_task => {},
            }

            // 客户端断开，从全局集合中移除当前发送器
            clients.lock().await.retain(|sender| sender != &client_tx);

            Ok(())
        })
    })
}

#[rocket::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 初始化全局客户端集合
    let shared_clients = Arc::new(Mutex::new(Vec::new()));

    rocket::build()
        .manage(shared_clients)
        .mount("/", routes![ws_broadcast])
        .launch()
        .await?;

    Ok(())
}

关键实现说明 #

1. 通道替代直接存储流：使用mpsc::UnboundedSender作为客户端的消息发送代理，避免直接操作独占的WebSocket流，解决多线程安全访问问题。
2. 双任务处理：每个客户端连接后启动两个异步任务：
   • 读取客户端输入消息，遍历全局发送器转发给所有在线客户端
   • 监听通道消息，将其他客户端的消息发送给当前客户端
3. 自动清理无效连接：转发消息时通过retain过滤掉发送失败的发送器（对应已断开的客户端），客户端断开时主动从集合中移除自己的发送器。
4. 错误处理：避免使用unwrap，改用match处理消息读取/发送错误，防止程序崩溃。

针对你原有代码的问题分析 #

• 直接存储DuplexStream错误：WebSocket流是独占资源，无法被多线程同时访问，且所有权转移后无法再次获取。
• 缺少连接清理：客户端断开后未从HashMap中移除无效条目，导致后续操作出错。
• 错误处理不严谨：serde_json::from_str(&concrete).unwrap()可能因格式错误崩溃，建议改用match或?处理。

内容的提问来源于stack exchange，提问作者Looouiiis