Day 19: WebSockets with Axum

Phase 2 · Jun 28, 2026

Agenda (2–3 hours)

• Read (45 min): Axum WebSocket example; RFC 6455 §1–4 (WebSocket protocol overview)
• Study (45 min): How does the HTTP upgrade mechanism work? Why is the WebSocket handshake inside an HTTP response?
• Practice (45 min): Build a WebSocket-based chat server: clients join named rooms; messages broadcast to all room members
• Challenge (30 min): Add heartbeat/ping-pong to detect dead connections; clean up the connection registry when connections drop

WebSocket Upgrade in Axum

use axum::extract::ws::{WebSocketUpgrade, WebSocket, Message};

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

async fn handle_socket(mut socket: WebSocket, state: AppState) {
    loop {
        match socket.recv().await {
            Some(Ok(Message::Text(text))) => {
                if socket.send(Message::Text(text)).await.is_err() { break; }
            }
            Some(Ok(Message::Ping(data))) => {
                socket.send(Message::Pong(data)).await.ok();
            }
            _ => break,
        }
    }
}

Broadcast to Room Members

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

struct RoomState {
    tx: broadcast::Sender<String>,
}

struct AppState {
    rooms: Arc<Mutex<HashMap<String, RoomState>>>,
}

async fn handle_socket(mut socket: WebSocket, room: String, state: AppState) {
    let rx = {
        let mut rooms = state.rooms.lock().unwrap();
        let room = rooms.entry(room).or_insert_with(|| {
            RoomState { tx: broadcast::channel(64).0 }
        });
        room.tx.subscribe()
    };
    // Split socket into read/write halves
    let (mut sender, mut receiver) = socket.split();
    // Spawn writer task consuming from broadcast rx
    // Spawn reader task forwarding to broadcast tx
}

Split for Concurrent Read/Write

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

let tx_clone = tx.clone();
let write_task = tokio::spawn(async move {
    while let Ok(msg) = rx.recv().await {
        sender.send(Message::Text(msg)).await.ok();
    }
});

let read_task = tokio::spawn(async move {
    while let Some(Ok(Message::Text(text))) = receiver.next().await {
        tx_clone.send(text).ok();
    }
});

tokio::select! {
    _ = write_task => {}
    _ = read_task => {}
}

Key Takeaways

• WebSocketUpgrade extractor handles the HTTP→WS upgrade handshake automatically
• WebSocket::split() gives separate read and write halves for concurrent use
• broadcast::Sender is the natural primitive for pub/sub within a room
• Clean up connection state (remove from rooms map) when the connection closes

Tomorrow: testing Axum services — how to write unit and integration tests.