urus/examples/ws_chat.rs

//! WebSocket chat rooms (v0.5): `urus::pubsub` wired to the v0.4 duplex.
//!
//!   cargo run --example ws_chat
//!   websocat ws://127.0.0.1:8080/chat/lobby     (run two of these)
//!
//! The shapes this demonstrates:
//!
//! - **One `PubSub<String>` for the whole app**, topics are rooms
//!   (`room:{name}`). Built lazily on the first connection via a
//!   NON-static `Arc<OnceLock<PubSub>>` captured by the route closure:
//!   `PubSub::new()` spawns the table actor, which smarm only allows
//!   in-runtime — and keeping the cell non-static means the table's
//!   last handle drops in-runtime when the drained pipeline drops, so
//!   `serve_with_shutdown` actually returns. A `static` cell would pin
//!   the table forever and block smarm's all-done. (Same constraint
//!   crud's store solves with its OnceLock; the non-static refinement
//!   is what makes graceful shutdown compose.)
//!
//! - **`on_open` subscribes and spawns the relay** — a listen-only
//!   client receives the room without ever sending. The subscription is
//!   pinned to the CONNECTION actor (`on_open` runs inside it), so the
//!   monitor cleans up exactly when the connection dies.
//!
//! - **The relay holds the `Receiver` and a `WsSender` clone — and
//!   deliberately NOT a `PubSub` handle.** A relay holding the handle
//!   keeps the table's inbox open while the table keeps the relay's
//!   receiver open: neither ever exits, and shutdown hangs. Receiver
//!   only: conn dies → monitor prunes → sender drops → relay's recv
//!   errs → relay exits. Every link in that chain is in-runtime.

use std::sync::{Arc, OnceLock};

use urus::{
    serve_with_shutdown, shutdown_handle, Config, Conn, Message, Next, Pipeline, PubSub, Router,
    WsHandler, WsSender,
};

type Bus = PubSub<String>;

struct ChatHandler {
    bus: Arc<OnceLock<Bus>>,
    room: String,
}

impl ChatHandler {
    fn topic(&self) -> String {
        format!("room:{}", self.room)
    }

    fn bus(&self) -> &Bus {
        // First connection anywhere spawns the table; we are inside the
        // connection actor here, so the spawn is legal.
        self.bus.get_or_init(PubSub::new)
    }
}

impl WsHandler for ChatHandler {
    fn on_open(&mut self, sender: &WsSender) {
        let topic = self.topic();
        let rx = match self.bus().subscribe(&topic) {
            Ok(rx) => rx,
            Err(_) => {
                let _ = sender.close(1011, "chat bus down");
                return;
            }
        };
        let _ = self
            .bus()
            .broadcast_from(smarm::self_pid(), &topic, format!("* someone joined {topic}"));

        // The relay: room messages -> this socket. Exits when the
        // subscription is pruned (conn death / unsubscribe) or the
        // socket is gone (WsClosed). See the module docs for why it
        // must not capture a Bus handle.
        let out = sender.clone();
        smarm::spawn(move || {
            while let Ok(msg) = rx.recv() {
                if out.send(Message::Text((*msg).clone())).is_err() {
                    break; // client gone or server draining
                }
            }
        });
    }

    fn on_message(&mut self, msg: Message, sender: &WsSender) {
        let Message::Text(text) = msg else {
            let _ = sender.close(1003, "text only");
            return;
        };
        // broadcast_from: the sender's own relay is skipped — no echo.
        // self_pid() here is the connection actor, the pid on_open
        // subscribed as.
        let _ = self
            .bus()
            .broadcast_from(smarm::self_pid(), self.topic(), text);
    }

    fn on_close(&mut self, _code: Option<u16>, _reason: &str) {
        let topic = self.topic();
        let _ = self
            .bus()
            .broadcast_from(smarm::self_pid(), &topic, format!("* someone left {topic}"));
        // No explicit unsubscribe: the connection actor is about to
        // exit and the monitor prunes the subscription (which is also
        // what stops the relay).
    }
}

fn main() {
    let bus: Arc<OnceLock<Bus>> = Arc::new(OnceLock::new());

    let pipeline = Pipeline::new().plug(Router::new().get("/chat/:room", move |c: Conn, _n: Next| {
        let room = c.params.get("room").unwrap_or("lobby").to_string();
        let bus = bus.clone();
        c.upgrade(ChatHandler { bus, room })
    }));

    let (handle, signal) = shutdown_handle();
    std::thread::spawn(move || {
        println!("ws_chat on ws://127.0.0.1:8080/chat/:room — press Enter to shut down");
        let mut line = String::new();
        let _ = std::io::stdin().read_line(&mut line);
        handle.shutdown();
    });

    serve_with_shutdown(Config::new("0.0.0.0:8080".parse().unwrap()), pipeline, signal).unwrap();
    println!("ws_chat: drained, bye");
}