第 13 課：WebSocket 和即時

Rust：從基礎到高級

第 4 部分：進階後端

亞洲開發網

1.Axum WebSocket

use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::IntoResponse,
    routing::get,
    Router,
};

async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    while let Some(msg) = socket.recv().await {
        let msg = match msg {
            Ok(msg) => msg,
            Err(_) => return,
        };

        match msg {
            Message::Text(text) => {
                let response = format!("Echo: {}", text);
                if socket.send(Message::Text(response)).await.is_err() {
                    return;
                }
            }
            Message::Close(_) => return,
            _ => {}
        }
    }
}

let app = Router::new().route("/ws", get(ws_handler));

2. 有廣播的聊天室

use tokio::sync::broadcast;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Clone)]
struct ChatState {
    rooms: Arc<RwLock<HashMap<String, broadcast::Sender<ChatMessage>>>>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct ChatMessage {
    room: String,
    user: String,
    content: String,
    timestamp: i64,
}

impl ChatState {
    fn new() -> Self {
        Self { rooms: Arc::new(RwLock::new(HashMap::new())) }
    }

    async fn join_room(&self, room: &str) -> broadcast::Receiver<ChatMessage> {
        let mut rooms = self.rooms.write().await;
        let tx = rooms.entry(room.to_string())
            .or_insert_with(|| broadcast::channel(100).0);
        tx.subscribe()
    }

    async fn send_message(&self, msg: ChatMessage) {
        let rooms = self.rooms.read().await;
        if let Some(tx) = rooms.get(&msg.room) {
            let _ = tx.send(msg);
        }
    }
}

async fn ws_chat(
    ws: WebSocketUpgrade,
    State(state): State<ChatState>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_chat(socket, state))
}

async fn handle_chat(mut socket: WebSocket, state: ChatState) {
    let mut rx = state.join_room("general").await;

    let (mut sender, mut receiver) = socket.split();

    // Task nhận broadcast
    let mut send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            let text = serde_json::to_string(&msg).unwrap();
            if sender.send(Message::Text(text)).await.is_err() { break; }
        }
    });

    // Task nhận từ client
    let state_clone = state.clone();
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(Message::Text(text))) = receiver.next().await {
            if let Ok(msg) = serde_json::from_str::<ChatMessage>(&text) {
                state_clone.send_message(msg).await;
            }
        }
    });

    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    }
}

3. 伺服器發送的事件

use axum::response::sse::{Event, Sse};
use tokio_stream::StreamExt;

async fn sse_handler() -> Sse<impl tokio_stream::Stream<Item = Result<Event, std::convert::Infallible>>> {
    let stream = tokio_stream::wrappers::IntervalStream::new(tokio::time::interval(Duration::from_secs(1)))
        .map(|_| {
            let data = serde_json::json!({ "time": Utc::now().to_rfc3339() });
            Ok(Event::default().data(data.to_string()))
        });
    Sse::new(stream)
}

下一篇： gRPC 與 Tonic。