urus/examples/ws_chat.rs

//! WebSocket chat rooms (v0.7): `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}`). `BUS` is a `const`: the handle is an address, not
//!   the table. The table actor is `BUS.child()`, handed to
//!   `serve_with_shutdown` as an app child, so it starts before the
//!   endpoint and — shutdown being ordered in reverse — stops after the
//!   endpoint has drained.
//!
//!   The v0.5 idiom this replaces was a non-static `Arc<OnceLock<PubSub>>`
//!   lazily initialised from the first connection, because `PubSub::new()`
//!   used to spawn the table and the handle used to own its life. Neither
//!   is true any more.
//!
//! - **`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.** It may
//!   also hold `BUS` — a handle pins nothing now — but it has no use for
//!   one. Conn dies → monitor prunes → sender drops → relay's recv errs
//!   → relay exits.

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

const BUS: PubSub<String> = PubSub::new("chat");

struct ChatHandler {
    room: String,
}

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

impl WsHandler for ChatHandler {
    fn on_open(&mut self, sender: &WsSender) {
        let topic = self.topic();
        let rx = match BUS.subscribe(&topic) {
            Ok(rx) => rx,
            Err(_) => {
                let _ = sender.close(1011, "chat bus down");
                return;
            }
        };
        let _ = 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).
        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 _ = BUS.broadcast_from(smarm::self_pid(), self.topic(), text);
    }

    fn on_close(&mut self, _code: Option<u16>, _reason: &str) {
        let topic = self.topic();
        let _ = 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 pipeline = Pipeline::new().plug(Router::new().get("/chat/:room", |c: Conn, _n: Next| {
        let room = c.params.get("room").unwrap_or("lobby").to_string();
        c.upgrade(ChatHandler { 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()),
        smarm::Config::default(),
        pipeline,
        vec![BUS.child()],
        signal,
    )
    .unwrap();
    println!("ws_chat: drained, bye");
}