Rust 用 Axum + Tokio 写 WebSocket 服务：广播、心跳与连接清理

需求：一个能广播的实时通道

聊天、行情、协作编辑都要服务端主动推。Axum 的 ws 提取器把握手交给框架，我们只需在 on_upgrade 之后拿到 socket 自己处理。下面是一个带广播和心跳的最小可用服务。

全局状态放一个广播通道

use axum::{
  extract::ws::{Message, WebSocket, WebSocketUpgrade},
  response::Response,
  routing::get,
  Router, State,
};
use tokio::sync::broadcast;
use std::sync::Arc;

#[derive(Clone)]
struct AppState {
  tx: broadcast::Sender<String>,
}

async fn ws_handler(
  ws: WebSocketUpgrade,
  State(state): State<AppState>,
) -> Response {
  ws.on_upgrade(move |socket| handle_socket(socket, state))
}

每个连接：读客户端、广播、心跳

use futures::{SinkExt, StreamExt};

async fn handle_socket(socket: WebSocket, state: AppState) {
  let (mut sender, mut receiver) = socket.split();
  let mut rx = state.tx.subscribe();

  // 收客户端消息，广播给所有人
  let mut recv_task = tokio::spawn(async move {
    while let Some(Ok(msg)) = receiver.next().await {
      if let Message::Text(text) = msg {
        let _ = state.tx.send(text.to_string());
      }
    }
  });

  // 把广播消息发回这个客户端
  let mut send_task = tokio::spawn(async move {
    while let Ok(text) = rx.recv().await {
      if sender.send(Message::Text(text)).await.is_err() {
        break;
      }
    }
  });

  // 任一边结束就取消另一边
  tokio::select! {
    _ = &mut recv_task => send_task.abort(),
    _ = &mut send_task => recv_task.abort(),
  }
}

心跳与房间

浏览器和代理会掐掉空闲连接，定期发 Message::Ping 保活。要按房间隔离，把全局 broadcast 换成“房间 id -> Sender”的 HashMap，客户端带上房间号再订阅对应通道即可。

// 每 30 秒发一次 Ping
let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
tokio::spawn(async move {
  loop {
    interval.tick().await;
    if sender.send(Message::Ping(vec![])).await.is_err() {
      break;
    }
  }
});

几个坑

• split 后 sender/receiver 要 move 到各自的 task，不能再借用原 socket。
• 客户端断开要两边都 abort，否则 goroutine 泄漏。
• 广播用 broadcast 是扇出模型；要单收单发的私信，得给每个连接单独建 channel 并维护路由表。