📁 Rust Backend with Axum · Part 29 of 42

Rust Web Server with Axum 0.8: Live Comments over WebSocket with tokio::select!

November 3, 2026

Episode 29 of the Rust Backend series, where we build the backend of a blog step by step. After episode 28 reads come from a Redis cache. Now: when someone comments on a post, everyone reading it should see the comment right away, without pressing reload. Today: WebSocket. All the code below is the final, tested code from the video, ready to copy.

New to Rust? Start with the beginner series.
The full backend series playlist is here.

We continue in the same blog_api project from episode 28. New today: axum's ws feature. The test client is websocat (cargo install websocat).

1. Why WebSocket

HTTP is question and answer: the server can only speak when asked. A WebSocket starts as an HTTP request, is "upgraded", and then stays open, and both sides can send at any time.

One comment, many listeners: Tokio's broadcast channel. Every receiver gets every message (unlike mpsc from episode 27, where each message goes to one receiver).

2. The handler, and a feature flag

The route /posts/{id}/live. The handler gets a WebSocketUpgrade, subscribes to the channel, and answers with on_upgrade. The socket task is a stub for now:

rust

use axum::extract::ws::{WebSocket, WebSocketUpgrade};
use tokio::sync::broadcast;

pub async fn live(
    Extension(live): Extension<broadcast::Sender<Comment>>,
    Path(post_id): Path<i64>,
    ws: WebSocketUpgrade,
) -> Response {
    let rx = live.subscribe();
    ws.on_upgrade(move |socket| live_socket(socket, rx, post_id))
}

async fn live_socket(mut socket: WebSocket, mut rx: broadcast::Receiver<Comment>, post_id: i64) {
    tracing::info!(post_id, "live socket opened");
}

bash

cargo build

text

error[E0432]: unresolved import `axum::extract::ws`

Like multipart in episode 24, WebSocket support is behind a feature:

bash

cargo add axum --features ws

3. The channel

main creates the channel, with room for 100 comments. We don't need the first receiver it returns (every socket subscribes its own), so _:

rust

let (live, _) = broadcast::channel::<Comment>(100);

The app takes it as one more parameter, hands the sender to the handlers, and adds the route:

rust

.route("/posts/{id}/live", get(handlers::live))
// ...
.layer(Extension(live))

add_comment sends each new comment into the channel. send only fails when nobody is listening, which is fine, so we ignore the result:

rust

let _ = live.send(comment.clone());
Ok((StatusCode::CREATED, Json(comment)))

4. tokio::select!: two sources at once

In the socket task we now have two things to wait for: a new comment from the channel, or a message from the client (like close). tokio::select! waits for both at once, and runs whichever branch is ready first:

rust

async fn live_socket(mut socket: WebSocket, mut rx: broadcast::Receiver<Comment>, post_id: i64) {
    loop {
        tokio::select! {
            event = rx.recv() => match event {
                Ok(c) => {
                    if c.post_id != post_id {
                        continue;
                    }
                    let text = serde_json::to_string(&c).unwrap();
                    if socket.send(Message::Text(text.into())).await.is_err() {
                        break;
                    }
                }
                Err(RecvError::Lagged(n)) => {
                    let _ = socket.send(Message::Text(format!("{{\"missed\":{n}}}").into())).await;
                }
                Err(RecvError::Closed) => break,
            },
            msg = socket.recv() => match msg {
                None | Some(Ok(Message::Close(_))) | Some(Err(_)) => break,
                Some(Ok(_)) => {}
            },
        }
    }
    tracing::info!(post_id, "live socket closed");
}

A comment for another post? Skip it. Ours? Send it as JSON text.
When the client closes, or the connection breaks (None or an error), we leave the loop. rx is dropped with the task, so the channel stops keeping messages for it.
Lagged: the channel keeps the last 100 messages. A receiver that falls further behind gets RecvError::Lagged(n). We tell the client how many it missed, and go on.

Add use axum::extract::ws::Message; and use tokio::sync::broadcast::error::RecvError;.

5. Two listeners

Run, log in, and make a second post. Then start two listeners on post 1 in the background, their output marked A: and B::

bash

cargo run
curl -s localhost:8080/posts -H 'content-type: application/json' -H "authorization: Bearer $TOKEN" -d '{"title":"Post two","body":"Hi."}' | jq -c .id
websocat ws://localhost:8080/posts/1/live | sed 's/^/A: /' &
websocat ws://localhost:8080/posts/1/live | sed 's/^/B: /' &

Now a comment on post 1:

bash

curl -s localhost:8080/posts/1/comments -H 'content-type: application/json' -H "authorization: Bearer $TOKEN" -d '{"body":"First!"}' | jq -c .id

Both listeners print the comment as JSON, A: and B:, at the same moment. A comment on post 2:

bash

curl -s localhost:8080/posts/2/comments -H 'content-type: application/json' -H "authorization: Bearer $TOKEN" -d '{"body":"Not for you"}' | jq -c .id

Nobody listening on post 1 sees it. Now we stop listener A (kill %1). The server logs live socket closed post_id=1. One more comment:

bash

curl -s localhost:8080/posts/1/comments -H 'content-type: application/json' -H "authorization: Bearer $TOKEN" -d '{"body":"Only B"}' | jq -c .id

Only B: receives it.

Who may listen? Reading is public on our blog, so anyone. For private data you would check a token (our AuthUser extractor) in live, before on_upgrade.

The final code

src/handlers.rs

rust

use axum::extract::{rejection::QueryRejection, Extension, Multipart, Path, Query, State};
use axum::{http::{header, StatusCode}, response::{IntoResponse, Response}, Json};

use crate::error::ApiError;
use crate::extract::{AdminUser, AuthUser, Role, ValidJson};
use crate::models::{Comment, Hit, Image, ListFilter, ListParams, NewComment, NewPost, NewUser, Page, Post, PostWithTags, SearchParams, User};
use crate::repo::Repo;
use crate::auth;
use crate::config::config;
use serde::Deserialize;
use serde_json::{json, Value};
use axum_extra::extract::cookie::{Cookie, CookieJar, SameSite};
use uuid::Uuid;
use crate::storage::Storage;
use std::sync::Arc;
use std::io;
use crate::jobs::Job;
use tokio::sync::mpsc;
use crate::cache;
use redis::aio::ConnectionManager;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use tokio::sync::broadcast;
use tokio::sync::broadcast::error::RecvError;

#[tracing::instrument(skip_all)]
pub async fn create_post<R: Repo>(
    State(repo): State<R>,
    user: AuthUser,
    ValidJson(new): ValidJson<NewPost>,
) -> Result<(StatusCode, Json<PostWithTags>), ApiError> {
    let post = repo.create(user.id, new).await.map_err(|err| match err.as_database_error() {
        Some(e) if e.is_unique_violation() => ApiError::Conflict("a post with this title already exists".into()),
        _ => ApiError::Database(err),
    })?;
    tracing::debug!(id = post.post.id, "created");
    Ok((StatusCode::CREATED, Json(post)))
}

pub async fn get_post<R: Repo>(
    State(repo): State<R>,
    Extension(redis): Extension<ConnectionManager>,
    Path(id): Path<i64>,
) -> Result<Json<Post>, ApiError> {
    if let Some(post) = cache::get(&redis, id).await {
        tracing::debug!(id, "cache hit");
        return Ok(Json(post));
    }
    tracing::debug!(id, "cache miss");
    let post = repo.get(id).await?.ok_or(ApiError::NotFound(id))?;
    cache::put(&redis, &post).await;
    Ok(Json(post))
}

pub async fn list_posts<R: Repo>(
    State(repo): State<R>,
    params: Result<Query<ListParams>, QueryRejection>,
) -> Result<Json<Page<Post>>, ApiError> {
    let Query(p) = params?;
    let page = p.page.unwrap_or(1).max(1);
    let per_page = p.per_page.unwrap_or(10).clamp(1, 100);
    let (items, total) = repo.list(ListFilter { page, per_page, tag: p.tag, sort: p.sort }).await?;
    Ok(Json(Page { items, page, per_page, total }))
}

pub async fn update_post<R: Repo>(
    State(repo): State<R>,
    Extension(redis): Extension<ConnectionManager>,
    Path(id): Path<i64>,
    user: AuthUser,
    ValidJson(new): ValidJson<NewPost>,
) -> Result<Json<Post>, ApiError> {
    let post = repo.get(id).await?.ok_or(ApiError::NotFound(id))?;
    if !can_edit(&user, &post) {
        return Err(ApiError::Forbidden);
    }
    let post = repo.update(id, new).await?.ok_or(ApiError::NotFound(id))?;
    cache::del(&redis, id).await;
    Ok(Json(post))
}

pub async fn delete_post<R: Repo>(State(repo): State<R>, Extension(redis): Extension<ConnectionManager>, Path(id): Path<i64>, user: AuthUser) -> Result<StatusCode, ApiError> {
    let post = repo.get(id).await?.ok_or(ApiError::NotFound(id))?;
    if !can_edit(&user, &post) {
        return Err(ApiError::Forbidden);
    }
    repo.delete(id).await?;
    cache::del(&redis, id).await;
    Ok(StatusCode::NO_CONTENT)
}

pub async fn list_comments<R: Repo>(
    State(repo): State<R>,
    Path(id): Path<i64>,
) -> Result<Json<Vec<Comment>>, ApiError> {
    repo.get(id).await?.ok_or(ApiError::NotFound(id))?;
    Ok(Json(repo.list_comments(id).await?))
}

pub async fn add_comment<R: Repo>(
    State(repo): State<R>,
    Path(id): Path<i64>,
    Extension(live): Extension<broadcast::Sender<Comment>>,
    _user: AuthUser,
    ValidJson(new): ValidJson<NewComment>,
) -> Result<(StatusCode, Json<Comment>), ApiError> {
    let comment = repo.add_comment(id, new).await.map_err(|err| {
        match err.as_database_error() {
            Some(db_err) if db_err.is_foreign_key_violation() => ApiError::NotFound(id),
            _ => ApiError::Database(err),
        }
    })?;
    let _ = live.send(comment.clone());
    Ok((StatusCode::CREATED, Json(comment)))
}

pub async fn search_posts<R: Repo>(
    State(repo): State<R>,
    Query(p): Query<SearchParams>,
) -> Result<Json<Vec<Hit>>, ApiError> {
    Ok(Json(repo.search(&p.q).await?))
}

pub async fn restore_post<R: Repo>(State(repo): State<R>, Path(id): Path<i64>) -> Result<StatusCode, ApiError> {
    if !repo.restore(id).await? {
        return Err(ApiError::NotFound(id));
    }
    Ok(StatusCode::OK)
}

pub async fn signup<R: Repo>(
    State(repo): State<R>,
    Extension(jobs): Extension<mpsc::Sender<Job>>,
    ValidJson(new): ValidJson<NewUser>,
) -> Result<(StatusCode, Json<User>), ApiError> {
    let hash = crate::auth::hash_password(new.password).await?;
    let user = repo.create_user(&new.email, &hash).await.map_err(|err| match err.as_database_error() {
        Some(e) if e.is_unique_violation() => ApiError::Conflict("email already registered".into()),
        _ => ApiError::Database(err),
    })?;
    if jobs.send(Job::Welcome { email: user.email.clone() }).await.is_err() {
        tracing::error!("job queue closed");
    }
    Ok((StatusCode::CREATED, Json(user)))
}

#[derive(Deserialize)]
pub struct Login {
    pub email: String,
    pub password: String,
}

fn refresh_cookie(token: String) -> Cookie<'static> {
    Cookie::build(("refresh", token))
        .http_only(true)
        .secure(true)
        .same_site(SameSite::Strict)
        .path("/auth")
        .max_age(time::Duration::days(30))
        .build()
}

pub async fn login<R: Repo>(
    State(repo): State<R>,
    jar: CookieJar,
    Json(body): Json<Login>,
) -> Result<(CookieJar, Json<Value>), ApiError> {
    let user = repo.find_user_by_email(&body.email).await?.ok_or(ApiError::Unauthorized)?;
    if !auth::verify_password(body.password, user.password_hash).await? {
        return Err(ApiError::Unauthorized);
    }
    let access = auth::make_token(user.id, &user.role, &config().jwt_secret)?;
    let refresh = auth::new_refresh_token();
    repo.store_refresh(&auth::hash_token(&refresh), user.id, auth::now() + auth::REFRESH_SECS).await?;
    Ok((jar.add(refresh_cookie(refresh)), Json(json!({ "token": access }))))
}

pub async fn me<R: Repo>(State(repo): State<R>, user: AuthUser) -> Result<Json<User>, ApiError> {
    let user = repo.get_user(user.id).await?.ok_or(ApiError::Unauthorized)?;
    Ok(Json(user))
}

pub async fn refresh<R: Repo>(State(repo): State<R>, jar: CookieJar) -> Result<(CookieJar, Json<Value>), ApiError> {
    let old = jar.get("refresh").ok_or(ApiError::Unauthorized)?.value().to_string();
    let user_id = repo.take_refresh(&auth::hash_token(&old), auth::now()).await?.ok_or(ApiError::Unauthorized)?;
    let user = repo.get_user(user_id).await?.ok_or(ApiError::Unauthorized)?;
    let access = auth::make_token(user.id, &user.role, &config().jwt_secret)?;
    let new = auth::new_refresh_token();
    repo.store_refresh(&auth::hash_token(&new), user.id, auth::now() + auth::REFRESH_SECS).await?;
    Ok((jar.add(refresh_cookie(new)), Json(json!({ "token": access }))))
}

pub async fn logout<R: Repo>(State(repo): State<R>, jar: CookieJar) -> Result<(CookieJar, StatusCode), ApiError> {
    if let Some(cookie) = jar.get("refresh") {
        repo.delete_refresh(&auth::hash_token(cookie.value())).await?;
    }
    Ok((jar.remove(Cookie::build("refresh").path("/auth")), StatusCode::NO_CONTENT))
}

pub fn can_edit(user: &AuthUser, post: &Post) -> bool {
    user.role == Role::Admin || post.author_id == Some(user.id)
}

pub async fn list_users<R: Repo>(State(repo): State<R>, _admin: AdminUser) -> Result<Json<Vec<User>>, ApiError> {
    Ok(Json(repo.list_users().await?))
}

pub async fn upload_image<R: Repo, S: Storage>(
    State(repo): State<R>,
    user: AuthUser,
    Path(id): Path<i64>,
    Extension(store): Extension<Arc<S>>,
    mut multipart: Multipart,
) -> Result<(StatusCode, Json<Image>), ApiError> {
    let post = repo.get(id).await?.ok_or(ApiError::NotFound(id))?;
    if !can_edit(&user, &post) {
        return Err(ApiError::Forbidden);
    }
    while let Some(field) = multipart.next_field().await? {
        if field.name() != Some("file") {
            continue;
        }
        let bytes = field.bytes().await?;
        let ext = image_kind(&bytes).ok_or(ApiError::UnsupportedMedia)?;
        let name = format!("{}.{ext}", Uuid::new_v4());
        store.put(&name, &bytes).await?;
        let image = repo.add_image(id, &format!("/files/{name}")).await?;
        return Ok((StatusCode::CREATED, Json(image)));
    }
    Err(ApiError::MissingFile)
}

fn image_kind(bytes: &[u8]) -> Option<&'static str> {
    if bytes.starts_with(&[0x89, b'P', b'N', b'G']) {
        Some("png")
    } else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
        Some("jpg")
    } else {
        None
    }
}

pub async fn get_file<S: Storage>(Extension(store): Extension<Arc<S>>, Path(name): Path<String>) -> Result<Response, ApiError> {
    let bytes = store.get(&name).await.map_err(|err| match err.kind() {
        io::ErrorKind::NotFound | io::ErrorKind::InvalidInput => ApiError::FileNotFound,
        _ => ApiError::Io(err),
    })?;
    let content_type = if name.ends_with(".png") { "image/png" } else { "image/jpeg" };
    Ok(([(header::CONTENT_TYPE, content_type)], bytes).into_response())
}

pub async fn live(
    Extension(live): Extension<broadcast::Sender<Comment>>,
    Path(post_id): Path<i64>,
    ws: WebSocketUpgrade,
) -> Response {
    let rx = live.subscribe();
    ws.on_upgrade(move |socket| live_socket(socket, rx, post_id))
}

async fn live_socket(mut socket: WebSocket, mut rx: broadcast::Receiver<Comment>, post_id: i64) {
    loop {
        tokio::select! {
            event = rx.recv() => match event {
                Ok(c) => {
                    if c.post_id != post_id {
                        continue;
                    }
                    let text = serde_json::to_string(&c).unwrap();
                    if socket.send(Message::Text(text.into())).await.is_err() {
                        break;
                    }
                }
                Err(RecvError::Lagged(n)) => {
                    let _ = socket.send(Message::Text(format!("{{\"missed\":{n}}}").into())).await;
                }
                Err(RecvError::Closed) => break,
            },
            msg = socket.recv() => match msg {
                None | Some(Ok(Message::Close(_))) | Some(Err(_)) => break,
                Some(Ok(_)) => {}
            },
        }
    }
    tracing::info!(post_id, "live socket closed");
}

src/main.rs

rust

mod auth;
mod config;
mod error;
mod extract;
mod handlers;
mod models;
mod repo;
mod storage;
mod mail;
mod jobs;
mod cache;

use axum::{routing::{get, post}, Extension, Router};
use repo::{Repo, SqliteRepo};
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
use std::str::FromStr;
use tower_http::trace::TraceLayer;
use tracing_subscriber::EnvFilter;
use axum::extract::DefaultBodyLimit;
use axum::http::{header, HeaderValue, Method};
use tower_http::cors::CorsLayer;
use tower_http::set_header::SetResponseHeaderLayer;
use std::sync::Arc;
use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
use std::net::SocketAddr;
use storage::{LocalStorage, MemStorage, Storage};
use mail::Mailer;
use jobs::Job;
use tokio::sync::mpsc;
use tokio::time::Duration;
use redis::aio::{ConnectionManager, ConnectionManagerConfig};
use tokio::sync::broadcast;
use models::Comment;

pub fn app<R: Repo, S: Storage>(repo: R, store: S, jobs: mpsc::Sender<Job>, cache: ConnectionManager, live: broadcast::Sender<Comment>) -> Router {
    let cors = CorsLayer::new()
        .allow_origin(HeaderValue::from_static("http://localhost:5500"))
        .allow_methods([Method::GET, Method::POST, Method::PUT, Method::DELETE])
        .allow_headers([header::AUTHORIZATION, header::CONTENT_TYPE])
        .allow_credentials(true);
    let governor = Arc::new(GovernorConfigBuilder::default().per_second(2).burst_size(5).finish().unwrap());
    let auth_routes = Router::new()
        .route("/signup", post(handlers::signup::<R>))
        .route("/login", post(handlers::login::<R>))
        .route("/auth/refresh", post(handlers::refresh::<R>))
        .layer(GovernorLayer::new(governor));
    Router::new()
        .route("/posts", get(handlers::list_posts::<R>).post(handlers::create_post::<R>))
        .route(
            "/posts/{id}",
            get(handlers::get_post::<R>).put(handlers::update_post::<R>).delete(handlers::delete_post::<R>),
        )
        .route("/posts/{id}/comments", get(handlers::list_comments::<R>).post(handlers::add_comment::<R>))
        .route("/search", get(handlers::search_posts::<R>))
        .route("/posts/{id}/restore", post(handlers::restore_post::<R>))
        .route("/me", get(handlers::me::<R>))
        .route("/admin/users", get(handlers::list_users::<R>))
        .route("/auth/logout", post(handlers::logout::<R>))
        .route("/posts/{id}/live", get(handlers::live))
        .route(
            "/posts/{id}/images",
            post(handlers::upload_image::<R, S>).layer(DefaultBodyLimit::max(10 * 1024 * 1024)),
        )
        .route("/files/{name}", get(handlers::get_file::<S>))
        .merge(auth_routes)
        .with_state(repo)
        .layer(Extension(Arc::new(store)))
        .layer(Extension(jobs))
        .layer(Extension(cache))
        .layer(Extension(live))
        .layer(cors)
        .layer(SetResponseHeaderLayer::if_not_present(header::X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff")))
        .layer(SetResponseHeaderLayer::if_not_present(header::X_FRAME_OPTIONS, HeaderValue::from_static("DENY")))
        .layer(SetResponseHeaderLayer::if_not_present(header::REFERRER_POLICY, HeaderValue::from_static("no-referrer")))
        .layer(TraceLayer::new_for_http())
}

#[tokio::main]
async fn main() {
    let cfg = config::config();
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")))
        .init();
    let addr = format!("{}:{}", cfg.host, cfg.port);
    let options = SqliteConnectOptions::from_str(&cfg.database_url)
        .unwrap()
        .create_if_missing(true)
        .journal_mode(SqliteJournalMode::Wal);
    let pool = SqlitePoolOptions::new()
        .max_connections(5)
        .connect_with(options)
        .await
        .unwrap();
    sqlx::migrate!().run(&pool).await.unwrap();
    let repo = SqliteRepo { pool };
    let (tx, rx) = mpsc::channel::<Job>(100);
    let mailer = Mailer::new(cfg);
    tokio::spawn(jobs::worker(rx, mailer));
    tokio::spawn(jobs::cleanup(repo.clone(), Duration::from_secs(cfg.cleanup_secs)));
    let client = redis::Client::open(cfg.redis_url.as_str()).unwrap();
    let redis_config = ConnectionManagerConfig::new()
        .set_number_of_retries(1)
        .set_connection_timeout(Some(Duration::from_millis(300)))
        .set_response_timeout(Some(Duration::from_millis(300)));
    let cache = ConnectionManager::new_with_config(client, redis_config).await.unwrap();
    let (live, _) = broadcast::channel::<Comment>(100);
    match cfg.storage.as_str() {
        "memory" => serve(app(repo, MemStorage::default(), tx, cache, live), &addr).await,
        _ => serve(app(repo, LocalStorage::new("uploads"), tx, cache, live), &addr).await,
    }
}

async fn serve(app: Router, addr: &str) {
    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
    tracing::info!(%addr, "listening");
    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await.unwrap();
}

Recap

A WebSocket starts as HTTP and stays open; both sides can send at any time
WebSocket support needs axum's ws feature; WebSocketUpgrade + on_upgrade hands over the socket
A broadcast channel delivers every comment to every subscriber
tokio::select! waits for the channel and the client at the same time
Handle Lagged (missed messages) and leave the loop on close or error

Next episode: the simpler cousin, server-sent events. One way only, and plain HTTP. Watch the series on YouTube.