Axum WebSockets With Broadcast Channels as a Weekend Project

A chat room is the smallest program that needs everything WebSockets bring: long lived connections, messages in both directions at once, fan-out to many peers and a failure mode for the peer that cannot keep up. This guide builds one in Axum 0.8 on Tokio and tests it with real clients. The layout and conventions follow the production Axum service; here the interest is the socket.

Everything here compiles and passes on Rust 1.99, edition 2024, axum 0.8.9 with the ws feature, tokio 1.54, futures-util 0.3 and, for the tests only, tokio-tungstenite 0.30.

The shape of the problem

HTTP handlers return once. A WebSocket handler does not: after the upgrade it owns a connection that stays open for minutes or hours, and it must do two jobs at the same time. It reads what the client sends, and it writes what other clients sent. If you write that as a single loop with socket.recv().await you can only wait for one thing at once, so a client that stops talking never receives anything, and a client that cannot be written to blocks everything behind it.

The fix is structural. Split the socket into its two halves, give each half to its own task, and connect the tasks through a channel. For fan-out the right channel is tokio::sync::broadcast: one sender, any number of receivers, every receiver sees every message. That gives the room its three parts:

1. A Hub that holds the broadcast sender.
2. A writer task per client that turns hub messages into frames.
3. A reader task per client that turns frames into hub messages.

Dependencies

[package]
name = "chat"
version = "0.1.0"
edition = "2024"

[dependencies]
axum = { version = "0.8.9", features = ["ws"] }
futures-util = "0.3.34"
tokio = { version = "1.54.1", features = ["full"] }

[dev-dependencies]
tokio-tungstenite = "0.30.0"

The ws feature of axum is not on by default. Without it WebSocketUpgrade does not exist, and the compiler says so with an unresolved import, not a hint about the feature. futures-util provides split, SinkExt::send and StreamExt::next; tokio-tungstenite is a dev-dependency because only the tests need a client.

The complete server

The library holds the hub and the handler. It is a library so that tests can start the same router that main does.

use std::sync::Arc;

use axum::{
    Router,
    extract::{
        State,
        ws::{Message, WebSocket, WebSocketUpgrade},
    },
    response::IntoResponse,
    routing::get,
};
use futures_util::{SinkExt, StreamExt};
use tokio::sync::broadcast::{self, error::RecvError};

#[derive(Clone)]
pub struct Hub {
    tx: broadcast::Sender<Arc<str>>,
}

impl Hub {
    pub fn new(capacity: usize) -> Self {
        let (tx, _) = broadcast::channel(capacity);
        Hub { tx }
    }
}

pub fn app(hub: Hub) -> Router {
    Router::new()
        .route("/ws", get(upgrade))
        .route("/online", get(online))
        .with_state(hub)
}

async fn online(State(hub): State<Hub>) -> String {
    // Every connected client holds exactly one receiver.
    hub.tx.receiver_count().to_string()
}

async fn upgrade(ws: WebSocketUpgrade, State(hub): State<Hub>) -> impl IntoResponse {
    // Chat lines are small; refuse anything bigger before it is buffered.
    ws.max_message_size(4 * 1024)
        .on_upgrade(move |socket| client(socket, hub))
}

async fn client(socket: WebSocket, hub: Hub) {
    let (mut sink, mut stream) = socket.split();
    let mut rx = hub.tx.subscribe();

    // Hub -> this client. Runs on its own task so a slow socket never blocks the reader.
    let mut send_task = tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Ok(line) => {
                    if sink.send(Message::text(line.to_string())).await.is_err() {
                        break;
                    }
                }
                Err(RecvError::Lagged(skipped)) => {
                    let note = format!("[{skipped} messages dropped]");
                    if sink.send(Message::text(note)).await.is_err() {
                        break;
                    }
                }
                Err(RecvError::Closed) => break,
            }
        }
    });

    // This client -> hub.
    let tx = hub.tx.clone();
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = stream.next().await {
            match msg {
                Message::Text(text) => {
                    // No subscribers is not an error worth reporting.
                    let _ = tx.send(Arc::from(text.as_str()));
                }
                Message::Close(_) => break,
                _ => {}
            }
        }
    });

    // When either half ends, stop the other so the task and the socket are released.
    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    }
}

Go through it in the order a connection experiences it.

app(hub) registers /ws and a small /online route, and gives both the hub as state. Hub wraps a broadcast::Sender<Arc<str>>. The payload type is Arc<str> rather than String on purpose. A broadcast channel clones the message for every receiver, so with 500 clients a String would be allocated and copied 500 times per line, while an Arc<str> clone is a reference count bump. For a chat line this hardly matters; for a market data feed it is the difference between fine and a profile dominated by memcpy.

upgrade is the HTTP half. WebSocketUpgrade is an extractor: it checks the Upgrade and Sec-WebSocket-Key headers and gives you a builder. max_message_size(4 * 1024) tells the protocol layer to reject any single message larger than 4 KiB. The default is far bigger (64 MiB for a message at the time of writing), which means an anonymous client can make your server buffer 64 MiB per connection. Always set this explicitly. on_upgrade takes a closure that receives the WebSocket once the handshake has finished and returns a future; Axum spawns that future for you.

client is the connection half. socket.split() consumes the WebSocket and returns a SplitSink and a SplitStream. hub.tx.subscribe() is called before either task starts, so there is no window in which this client misses messages sent after it connected.

The writer task loops on rx.recv(). A received line becomes Message::text(...) and goes down the sink. If the send fails the client is gone, so the loop ends. The interesting arm is RecvError::Lagged, covered below. RecvError::Closed only happens when every sender has been dropped, which means the hub itself is shutting down.

The reader task pulls from the stream and forwards text frames into the hub. tx.send returns an error when there are no subscribers; the result is ignored on purpose, because "nobody is listening" is not a failure of the sender. Pings and pongs fall into the _ arm and need no handling: the underlying library answers pings itself. A Close frame ends the loop.

The last statement is the part people forget. Two tasks mean two lifetimes. tokio::select! waits for the first to finish and aborts the other. Without it, a client that closes its socket ends the reader but leaves the writer parked forever on rx.recv(), holding a receiver slot in the channel and an open half of the connection. After a day of mobile clients dropping off networks you have thousands of those. Pass &mut send_task rather than the handle by value so the other branch can still call abort on the loser.

A detail that changed in Axum 0.8: Message::Text now carries Utf8Bytes, not a String, so the text was validated once when the frame arrived and text.as_str() is free. Older tutorials call text.to_string() or match on Message::Text(String), and the compiler error you get from copying one is a type mismatch on the variant, which looks worse than it is. Binary frames land in Message::Binary, which this room ignores; a client that sends only binary receives nothing and never learns why, so a production server should close such a connection with an explicit reason instead of dropping its frames.

Why the sink can only live in one task

You will be tempted to answer each message with an acknowledgement from the reader task. Try it and the compiler objects:

error[E0382]: use of moved value: `sink`
  --> src/lib.rs:63:38
   |
38 |     let (mut sink, mut stream) = socket.split();
   |          -------- move occurs because `sink` has type `SplitSink<WebSocket, axum::extract::ws::Message>`, which does not implement the `Copy` trait
...
42 |     let mut send_task = tokio::spawn(async move {
   |                                      ---------- value moved here
43 |         loop {
   |         ---- inside of this loop
...
46 |                     if sink.send(Message::text(line.to_string())).await.is_err() {
   |                        ---- variable moved due to use in coroutine
...
63 |     let mut recv_task = tokio::spawn(async move {
   |                                      ^^^^^^^^^^ value used here after move
...

sink was moved into the writer's async move block, so the reader cannot use it. The compiler is right to refuse: a sink is a single writer, and two tasks writing to it would interleave frames. The cure is not Arc<Mutex<Sink>>, which works but serialises the room behind a lock. Send the acknowledgement through the same path as everything else, either by pushing into the hub or by giving the writer task a small mpsc channel of its own. One task owns the write half; everything else sends it a message.

Run it

use chat::{Hub, app};

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await?;
    println!("ws://{}/ws", listener.local_addr()?);
    axum::serve(listener, app(Hub::new(64))).await
}

cargo run prints the address, and any WebSocket client can connect. In a browser console:

const ws = new WebSocket("ws://127.0.0.1:3000/ws");
ws.onmessage = (e) => console.log("received:", e.data);
ws.onopen = () => ws.send("hello room");

Open the console in two tabs and a line typed in one appears in both, including the sender's own, because the hub does not know who sent what. If you want the sender excluded, prefix each line with a client id inside the hub and compare it in the writer task before sending; the cheap way is to make the channel carry (ClientId, Arc<str>).

What a slow client costs

A broadcast channel is bounded. The capacity passed to channel is the number of messages kept in the ring buffer. If a receiver falls more than that far behind, the oldest messages are overwritten and its next recv returns Lagged(n). This is the property that protects the server: one phone on a bad connection cannot make memory grow without limit, because there is no per-client queue to grow.

Here is the behaviour in isolation, with a channel of capacity 2 and five messages sent before the receiver reads anything:

use tokio::sync::broadcast::{self, error::RecvError};

#[tokio::main]
async fn main() {
    // Capacity 2: the channel keeps only the two newest messages.
    let (tx, mut rx) = broadcast::channel::<u32>(2);

    for n in 1..=5 {
        tx.send(n).unwrap();
    }

    loop {
        match rx.recv().await {
            Ok(n) => println!("got {n}"),
            Err(RecvError::Lagged(skipped)) => println!("lagged: skipped {skipped}"),
            Err(RecvError::Closed) => break,
        }
        if rx.is_empty() {
            break;
        }
    }
}

lagged: skipped 3
got 4
got 5

Three messages were skipped and the receiver got the two newest. The channel did not block the sender and did not fail. In the chat server the writer task turns that error into a [3 messages dropped] line, which tells the user the truth instead of leaving a silent gap. Whether dropping is right depends on the application. For chat, dropping the oldest is fine. For a ledger of money movements it is not; there you want a per-client mpsc queue with a policy for what to do when it fills (usually, disconnect the client and let it resync).

Pick the capacity from the burst you want to absorb: the message rate at the busiest moment multiplied by the longest stall you will tolerate. A room that peaks at 200 messages a second and tolerates a one second stall wants a capacity of at least 200. The server above uses 64, which is a toy value.

Test with real clients

A WebSocket server is easy to break in ways a unit test never sees, so the integration test uses actual clients. Binding to port 0 lets the operating system pick a free port, which keeps parallel tests from colliding:

use std::time::Duration;

use chat::{Hub, app};
use futures_util::{SinkExt, StreamExt};
use tokio::net::TcpListener;
use tokio_tungstenite::{connect_async, tungstenite::Message};

async fn start(capacity: usize) -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app(Hub::new(capacity)))
            .await
            .unwrap();
    });
    format!("ws://{addr}/ws")
}

async fn next_text(
    ws: &mut (impl StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>> + Unpin),
) -> String {
    let msg = tokio::time::timeout(Duration::from_secs(2), ws.next())
        .await
        .expect("timed out waiting for a message")
        .expect("socket closed")
        .expect("socket error");
    msg.into_text().unwrap().to_string()
}

#[tokio::test]
async fn every_client_sees_every_line() {
    let url = start(16).await;
    let (mut alice, _) = connect_async(&url).await.unwrap();
    let (mut bob, _) = connect_async(&url).await.unwrap();
    // Subscription happens inside the upgraded task; give both a moment to register.
    tokio::time::sleep(Duration::from_millis(100)).await;

    alice.send(Message::text("hello from alice")).await.unwrap();

    assert_eq!(next_text(&mut alice).await, "hello from alice");
    assert_eq!(next_text(&mut bob).await, "hello from alice");
}

#[tokio::test]
async fn oversized_message_closes_the_socket() {
    let url = start(16).await;
    let (mut ws, _) = connect_async(&url).await.unwrap();
    ws.send(Message::text("x".repeat(10_000))).await.unwrap();

    let end = tokio::time::timeout(Duration::from_secs(2), ws.next())
        .await
        .unwrap();
    assert!(matches!(
        end,
        None | Some(Err(_)) | Some(Ok(Message::Close(_)))
    ));
}

// The same server answers plain HTTP on /online, so a bare TCP request is enough.
async fn online(url: &str) -> usize {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let addr = url.trim_start_matches("ws://").trim_end_matches("/ws");
    let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
    let request = b"GET /online HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n";
    stream.write_all(request).await.unwrap();

    let mut raw = String::new();
    stream.read_to_string(&mut raw).await.unwrap();
    raw.rsplit("\r\n\r\n")
        .next()
        .unwrap()
        .trim()
        .parse()
        .unwrap()
}

#[tokio::test]
async fn a_closed_client_releases_its_receiver() {
    let url = start(16).await;
    let (alice, _) = connect_async(&url).await.unwrap();
    let (bob, _) = connect_async(&url).await.unwrap();
    tokio::time::sleep(Duration::from_millis(100)).await;
    assert_eq!(online(&url).await, 2);

    drop(bob);
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(online(&url).await, 1);
    drop(alice);
}

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
test oversized_message_closes_the_socket ... ok
test every_client_sees_every_line ... ok
test a_closed_client_releases_its_receiver ... ok
test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.32s
test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

every_client_sees_every_line connects two clients, sends from one and expects the text on both. The 100 ms sleep is there because subscription happens inside the upgraded task, a moment after the handshake returns; a sleep is crude, and the sturdier version has the server send a hello frame that the test waits for. next_text wraps every read in a two second timeout, so a regression shows up as a clear "timed out waiting for a message" instead of a test run that hangs until CI kills it.

oversized_message_closes_the_socket sends 10,000 bytes against a 4 KiB limit and expects the connection to end. This pins the max_message_size setting: deleting it in a refactor turns a green test red.

The third test is the one that earns its keep. a_closed_client_releases_its_receiver asks the server how many clients it thinks it has, through the /online route, which reports receiver_count() from the hub. Two clients connect and the count is 2. One drops, and 200 ms later the count must be 1. This test exists to prove the select! at the end of client. Replace its two abort calls with empty blocks, so the surviving task is simply left running, and the test fails:

thread 'a_closed_client_releases_its_receiver' panicked at tests/broadcast.rs:88:5:
assertion `left == right` failed
  left: 2
 right: 1

The server still counts a client that left. The reader task saw the closed socket and ended; the writer task, nobody's problem any more, stayed parked on rx.recv() with its receiver registered. It would only have noticed when the next chat line arrived and the send failed, and in a quiet room that can be never. This is a leak you cannot see in a test that only checks messages, which is why the test asks for the count directly.

Heartbeats and idle connections

TCP does not tell you when a mobile client walks out of Wi-Fi range; the connection stays "open" on the server until a timeout the operating system chooses, often two hours. For a chat room that means ghosts. The protocol answer is a ping on a timer: send Message::Ping every 30 seconds from the writer task, and treat a missing pong, or any failed send, as a dead peer. Add an interval with tokio::time::interval to the writer's select! and the half-dead connections clean themselves up. Load balancers have their own idle timeout (60 seconds is a common default), so the heartbeat has to be more frequent than the lowest idle limit between the client and your process.

Authentication happens before the upgrade

The browser WebSocket API cannot set an Authorization header. Authenticate in the HTTP request that precedes the upgrade: a session cookie, which the browser sends automatically, or a short lived token in the query string that the handler validates before calling on_upgrade. Reject with 401 at that point, as the extractor for any other route would. Once the connection is upgraded, there is no clean way to send a status code, only to close the socket. If you check the Origin header too, you stop another site from opening sockets with your visitors' cookies.

Choosing the channel: broadcast, mpsc or watch

The hub uses broadcast because every client wants every message. That is one of three reasonable choices, and picking the wrong one is the usual cause of memory growth in socket servers.

A rule that has held up: if the message is an event, you want broadcast or mpsc; if it is a state, you want watch. A scoreboard pushed over a socket does not need every intermediate score, only the last one, and watch makes the server do the minimum amount of work for a slow client. A chat line is an event and losing it matters a little, so broadcast with an honest "messages dropped" note is a fair trade. Anything financial is an event that must not be lost, which means a per-client mpsc and a policy for a full queue.

Close sockets politely

A server that stops by being killed leaves every client to discover the loss through a timeout. When you control the shutdown, send each client a Message::Close first, with a code and a reason, so a well written client reconnects at once rather than after its own 30 second guess. The simplest implementation is a second broadcast, or a CancellationToken, that every writer task selects on: when it fires, the writer sends the close frame, then returns. Axum graceful shutdown explains why with_graceful_shutdown alone does not do this: it waits for connections to end, and a WebSocket connection never ends by itself, so a deploy would hang until the platform killed it.

Before you put it in front of users

1. max_message_size is set to the largest legitimate frame, not the default.
2. Authentication and an Origin check happen before on_upgrade.
3. The channel capacity covers the busiest burst you will tolerate, and the lag behaviour is a decision rather than a default.
4. Both tasks end together; a test like a_closed_client_releases_its_receiver keeps it that way.
5. A heartbeat is faster than every idle timeout between you and the client.
6. There is a limit on connections per user or per address, because each socket costs two tasks and a file descriptor.

Where this design stops

The hub lives in one process. Two instances of this server behind a load balancer are two separate rooms: a message sent to instance A never reaches clients on instance B. To scale out you put a shared bus in the middle (Redis pub/sub, NATS, or a database LISTEN/NOTIFY) and have each instance forward bus messages into its local hub. The local broadcast channel stays; it just stops being the only source of messages.

If your messages are server to client only (notifications, progress, a live feed), skip the socket and use server sent events, which are plain HTTP, pass through proxies without special configuration and reconnect on their own. WebSockets earn their complexity when the client sends frequently too.

Next steps

Add a ClientId and a name to each message, serialise them with serde_json, and send JSON frames instead of bare text. Put the heartbeat in. Then follow the rest of the Axum series: timeouts and limits apply to the upgrade request, and graceful shutdown needs a plan for sockets that never finish on their own, such as sending a close frame to every client when the server stops.