urus/src/pubsub.rs

//! `urus::pubsub` — local-node topic pub/sub, phoenix_pubsub-shaped.
//!
//! Deliberately independent of HTTP: nothing here imports the rest of
//! urus, only smarm. Usable from any smarm app; candidate for crate
//! extraction later.
//!
//! # Shape (v0.5 design, user-ratified)
//!
//! - **Generic per instance** (`PubSub<M>`): zero-cost dispatch, no
//!   downcasts. One instance per message domain; heterogeneous events on
//!   one bus are an app-side `enum M`.
//! - **Subscribe returns a `Receiver<Arc<M>>`**: the subscriber owns its
//!   receive loop and composes with the v0.4 ws outbound pattern (a relay
//!   actor pumping the receiver into a `WsSender` clone). One fresh
//!   receiver per subscription; fan-in (`subscribe_with(topic, sender)`)
//!   is a compatible later addition if wanted.
//! - **Unbounded mailboxes**: `broadcast` never blocks the topic table —
//!   a bounded-block policy would head-of-line-block *every* topic behind
//!   one slow subscriber, and bounded-drop silently loses messages.
//!   Memory risk is a live-but-not-receiving subscriber, bounded by two
//!   cleanup paths: monitor `Down` on subscriber death, and pruning on
//!   send failure when a receiver was dropped. Bench before sharding.
//! - **User-owned handle**: construct a `PubSub<M>`, clone it around
//!   (handler closures already thread state this way). No global named
//!   instance — see the module note on `register` below.
//! - **One subscription per `(pid, topic)`**: subscribe is idempotent;
//!   re-subscribing replaces the sender (the old receiver's channel
//!   closes). Kills the silent double-delivery bug class.
//!
//! # Cleanup mechanics
//!
//! The table monitors each subscriber pid **once** (first subscribe) via
//! the gen_server [`Watcher`]; the `Down` removes the pid from every
//! topic — a full table scan, deliberately: deaths are rare and topic
//! counts small at this stage, a pid→topics reverse index is the
//! documented optimisation if that ever measures hot (same spirit as the
//! roadmap's "don't pre-shard"). A pid that unsubscribes from everything
//! but stays alive keeps its (one-shot, inert) monitor until death;
//! that's a bounded bookkeeping entry, not a leak.
//!
//! # The handle is an address, not an owner (v0.7)
//!
//! [`PubSub::new`] takes a **name** and spawns nothing. It is `const`,
//! costs a `&'static str`, and may be built anywhere — out of the
//! runtime, in a `static`, in a route closure, held by a long-lived
//! relay. Every operation resolves the name through smarm's registry, so
//! a table restarted by its supervisor is reached transparently.
//!
//! The table actor is started by [`PubSub::child`], a `ChildSpec` you put
//! in your supervision tree (or hand to `serve_with*`, which puts it in
//! the root ahead of the endpoint). Its lifetime is the supervisor's:
//! it stops when the supervisor shuts it down, in reverse start order,
//! after the endpoint has drained.
//!
//! ```ignore
//! const BUS: PubSub<Event> = PubSub::new("events");
//! serve_with(cfg, rt_cfg, pipeline, vec![BUS.child()])?;
//! ```
//!
//! This replaces the v0.5–v0.6 idiom where `PubSub::new()` spawned the
//! table and the handle owned its life — a non-static
//! `Arc<OnceLock<PubSub<M>>>` lazily initialised from the first handler,
//! with matching rules that the cell must not be `static` and that
//! relays must never hold a clone. All of that existed to hand-manage a
//! refcount. smarm 0.7 made a server's lifetime its own (refs are
//! addresses; the inbox no longer closes when the last ref drops), which
//! both removed the mechanism those rules relied on and made the named
//! lookup below possible.
//!
//! Cost: one registry resolution per operation, including per broadcast.
//! Caching a `GenServerRef` in the handle would save it and go stale
//! across exactly the restart the supervisor exists to perform. Measure
//! before optimising — see the bench item in `ROADMAP.md`.

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::sync::Arc;

use smarm::gen_server::{self, GenServer, GenServerBuilder, GenServerCtx, GenServerName};
use smarm::{channel, ChildSpec, Down, GenServerRef, Pid, Receiver, Restart, Sender, Watcher};

// ---------------------------------------------------------------------------
// Public handle
// ---------------------------------------------------------------------------

/// The topic table is gone — its gen_server actor exited (panic or runtime
/// shutdown). Every operation on the handle fails with this from then on.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PubSubDown;

impl fmt::Display for PubSubDown {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "pubsub topic table is down")
    }
}

impl std::error::Error for PubSubDown {}

/// The address of one pub/sub instance: a name, resolved per operation.
///
/// Cheap (`Copy`, a `&'static str`), constructible anywhere including
/// `const` context, and owns nothing — the table actor it addresses is
/// started by [`child`](Self::child) under a supervisor. Every operation
/// returns [`PubSubDown`] if no live table currently holds the name.
///
/// Payloads are broadcast as `Arc<M>`: one allocation per broadcast, not
/// per subscriber.
pub struct PubSub<M: Send + Sync + 'static> {
    name: GenServerName<Table<M>>,
}

impl<M: Send + Sync + 'static> Clone for PubSub<M> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<M: Send + Sync + 'static> Copy for PubSub<M> {}

impl<M: Send + Sync + 'static> PubSub<M> {
    /// Address the topic table registered under `name`. Spawns nothing
    /// and never fails: the name is resolved at each use.
    pub const fn new(name: &'static str) -> Self {
        PubSub { name: GenServerName::new(name) }
    }

    /// The `ChildSpec` that runs this instance's table actor. Put it in
    /// your supervision tree ahead of anything that broadcasts —
    /// `serve_with*` takes a `Vec<ChildSpec>` for exactly this.
    ///
    /// `Permanent`: a table that dies is a bug, and its subscribers'
    /// receivers died with it, so the restart is only half a repair —
    /// pair it with a `RestForOne` parent (as `serve_with*` does) so the
    /// endpoint restarts behind it and connections re-subscribe.
    pub fn child(&self) -> ChildSpec {
        let name = self.name;
        ChildSpec::new(Restart::Permanent, move || {
            if let Err(e) = GenServerBuilder::new(Table::<M>::new()).named(name).run() {
                panic!("urus pubsub '{}': name already taken: {e:?}", name.as_str());
            }
        })
    }

    /// The registry key this handle resolves.
    pub const fn name(&self) -> &'static str {
        self.name.as_str()
    }

    fn server(&self) -> Result<GenServerRef<Table<M>>, PubSubDown> {
        gen_server::whereis_server(self.name).ok_or(PubSubDown)
    }

    /// Subscribe the **calling actor** to `topic`. Returns the receiving
    /// end; messages arrive as `Arc<M>`. Idempotent per `(pid, topic)`:
    /// subscribing again replaces the sender, closing the previously
    /// returned receiver.
    ///
    /// The subscription is cleaned up when the calling actor dies. If the
    /// receive loop runs in a *different* actor than the session that
    /// should scope the subscription, use [`subscribe_as`](Self::subscribe_as)
    /// to pin cleanup to the right pid (e.g. a ws connection actor
    /// subscribing, with a spawned relay doing the receiving).
    pub fn subscribe(&self, topic: impl Into<String>) -> Result<Receiver<Arc<M>>, PubSubDown> {
        self.subscribe_as(smarm::self_pid(), topic)
    }

    /// [`subscribe`](Self::subscribe), but the subscription's lifetime is
    /// tied to `pid` instead of the calling actor.
    pub fn subscribe_as(
        &self,
        pid: Pid,
        topic: impl Into<String>,
    ) -> Result<Receiver<Arc<M>>, PubSubDown> {
        let (tx, rx) = channel();
        // A call, not a cast: when this returns the table is updated, so
        // a broadcast issued right after by the same caller is seen.
        match self.server()?.call(Call::Subscribe { topic: topic.into(), pid, tx }) {
            Ok(_) => Ok(rx),
            Err(_) => Err(PubSubDown),
        }
    }

    /// Drop the calling actor's subscription to `topic` (its receiver's
    /// channel closes). No-op if not subscribed.
    pub fn unsubscribe(&self, topic: impl Into<String>) -> Result<(), PubSubDown> {
        self.unsubscribe_as(smarm::self_pid(), topic)
    }

    /// [`unsubscribe`](Self::unsubscribe) for an explicit pid.
    pub fn unsubscribe_as(&self, pid: Pid, topic: impl Into<String>) -> Result<(), PubSubDown> {
        self.server()?
            .cast(Cast::Unsubscribe { topic: topic.into(), pid })
            .map_err(|_| PubSubDown)
    }

    /// Broadcast `msg` to every subscriber of `topic`. Never blocks on
    /// subscribers (unbounded mailboxes); returns once the table has the
    /// request queued.
    pub fn broadcast(&self, topic: impl Into<String>, msg: M) -> Result<(), PubSubDown> {
        self.cast_broadcast(topic.into(), msg, None)
    }

    /// [`broadcast`](Self::broadcast), skipping delivery to `from` —
    /// pass `smarm::self_pid()` for the phoenix `broadcast_from` shape
    /// ("everyone in the room but me").
    pub fn broadcast_from(
        &self,
        from: Pid,
        topic: impl Into<String>,
        msg: M,
    ) -> Result<(), PubSubDown> {
        self.cast_broadcast(topic.into(), msg, Some(from))
    }

    /// Number of live subscriptions on `topic`. Counts entries in the
    /// table — a subscriber whose receiver was dropped but hasn't been
    /// pruned yet (no broadcast since, still alive) is still counted.
    pub fn subscriber_count(&self, topic: impl Into<String>) -> Result<usize, PubSubDown> {
        match self.server()?.call(Call::Count { topic: topic.into() }) {
            Ok(Reply::Count(n)) => Ok(n),
            Ok(Reply::Subscribed) => unreachable!("Count call answered with Subscribed"),
            Err(_) => Err(PubSubDown),
        }
    }

    /// The topic table actor's pid, or `None` if no live table holds the
    /// name (not started yet, or between a crash and its restart).
    pub fn pid(&self) -> Option<Pid> {
        self.server().ok().map(|s| s.pid())
    }

    fn cast_broadcast(&self, topic: String, msg: M, skip: Option<Pid>) -> Result<(), PubSubDown> {
        self.server()?
            .cast(Cast::Broadcast { topic, msg: Arc::new(msg), skip })
            .map_err(|_| PubSubDown)
    }
}

// ---------------------------------------------------------------------------
// The table gen_server
// ---------------------------------------------------------------------------

enum Call<M> {
    Subscribe { topic: String, pid: Pid, tx: Sender<Arc<M>> },
    Count { topic: String },
}

enum Reply {
    Subscribed,
    Count(usize),
}

enum Cast<M> {
    Unsubscribe { topic: String, pid: Pid },
    Broadcast { topic: String, msg: Arc<M>, skip: Option<Pid> },
}

struct Table<M: Send + Sync + 'static> {
    topics: HashMap<String, HashMap<Pid, Sender<Arc<M>>>>,
    /// Pids we already hold a monitor for. Monitors are one-shot and
    /// created at most once per live pid (first subscribe); `handle_down`
    /// retires the entry so a reused-slot pid (fresh generation) gets a
    /// fresh monitor.
    monitored: HashSet<Pid>,
    watcher: Option<Watcher<Table<M>>>,
}

impl<M: Send + Sync + 'static> Table<M> {
    fn new() -> Self {
        Table { topics: HashMap::new(), monitored: HashSet::new(), watcher: None }
    }
}

impl<M: Send + Sync + 'static> GenServer for Table<M> {
    type Call = Call<M>;
    type Reply = Reply;
    type Cast = Cast<M>;
    type Info = ();
    type Timer = ();

    fn init(&mut self, ctx: &GenServerCtx<Self>) {
        self.watcher = Some(ctx.watcher());
    }

    fn handle_call(&mut self, request: Call<M>) -> Reply {
        match request {
            Call::Subscribe { topic, pid, tx } => {
                if self.monitored.insert(pid) {
                    let m = smarm::monitor(pid);
                    self.watcher
                        .as_ref()
                        .expect("watcher set in init")
                        .watch(m);
                }
                // Insert replaces: idempotent per (pid, topic); the old
                // sender drops and the stale receiver's channel closes.
                self.topics.entry(topic).or_default().insert(pid, tx);
                Reply::Subscribed
            }
            Call::Count { topic } => {
                Reply::Count(self.topics.get(&topic).map_or(0, HashMap::len))
            }
        }
    }

    fn handle_cast(&mut self, request: Cast<M>) {
        match request {
            Cast::Unsubscribe { topic, pid } => {
                if let Some(subs) = self.topics.get_mut(&topic) {
                    subs.remove(&pid);
                    if subs.is_empty() {
                        self.topics.remove(&topic);
                    }
                }
            }
            Cast::Broadcast { topic, msg, skip } => {
                let Some(subs) = self.topics.get_mut(&topic) else { return };
                // Prune-on-send-failure: a dropped receiver (subscriber
                // alive but done listening, e.g. ws conn reaped at
                // write_timeout) is removed lazily here; subscriber
                // *death* is handled eagerly by the monitor.
                subs.retain(|pid, tx| {
                    if skip == Some(*pid) {
                        return true;
                    }
                    tx.send(Arc::clone(&msg)).is_ok()
                });
                if subs.is_empty() {
                    self.topics.remove(&topic);
                }
            }
        }
    }

    fn handle_down(&mut self, down: Down) {
        self.monitored.remove(&down.pid);
        // Full scan, by design — see the module docs.
        self.topics.retain(|_, subs| {
            subs.remove(&down.pid);
            !subs.is_empty()
        });
    }
}

// ---------------------------------------------------------------------------
// Tests (runtime-backed: results collected outside, asserted after run)
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;
    use std::time::Duration;

    /// Poll `f` until it returns true or ~2s elapse. For the
    /// asynchronous cleanup paths (monitor Down delivery).
    fn eventually(mut f: impl FnMut() -> bool) -> bool {
        for _ in 0..200 {
            if f() {
                return true;
            }
            smarm::sleep(Duration::from_millis(10));
        }
        false
    }

    /// Start `ps`'s table under a supervisor and hand back the sup's pid.
    ///
    /// The readiness poll is smarm's "start order is not start readiness"
    /// gap (see smarm ROADMAP): `start_child` spawns and moves on, so the
    /// name may not be bound when this returns. Real apps don't hit it —
    /// a handler only runs once a connection has been accepted, long
    /// after the tree is up — but a test that broadcasts immediately does.
    fn start_table<M: Send + Sync + 'static>(ps: PubSub<M>) -> Pid {
        let sup = smarm::spawn(move || smarm::OneForOne::new().child(ps.child()).run());
        assert!(eventually(|| ps.pid().is_some()), "table never bound its name");
        sup.pid()
    }

    /// Ordered shutdown of the tree `start_table` built.
    fn stop_table(sup: Pid) {
        smarm::request_shutdown(sup);
    }

    #[test]
    fn broadcast_reaches_all_subscribers_once() {
        let got = Arc::new(Mutex::new(Vec::<(u32, String)>::new()));
        let got2 = got.clone();
        smarm::run(move || {
            let ps = PubSub::<String>::new("test-bus");
            let sup = start_table(ps);
            let mut handles = Vec::new();
            for i in 0..2u32 {
                let got = got2.clone();
                handles.push(smarm::spawn(move || {
                    let rx = ps.subscribe("room:a").unwrap();
                    let msg = rx.recv().unwrap();
                    got.lock().unwrap().push((i, (*msg).clone()));
                }));
            }
            // Subscribes are calls from the spawned actors; wait until
            // both are in the table before broadcasting.
            assert!(eventually(|| ps.subscriber_count("room:a").unwrap() == 2));
            ps.broadcast("room:a", "hello".to_string()).unwrap();
            for h in handles {
                h.join().unwrap();
            }
            stop_table(sup);
        });
        let mut v = got.lock().unwrap().clone();
        v.sort();
        assert_eq!(v, vec![(0, "hello".into()), (1, "hello".into())]);
    }

    #[test]
    fn arc_payload_is_shared_not_cloned() {
        let ptrs = Arc::new(Mutex::new(Vec::<usize>::new()));
        let ptrs2 = ptrs.clone();
        smarm::run(move || {
            let ps = PubSub::<Vec<u8>>::new("test-bus");
            let sup = start_table(ps);
            let mut handles = Vec::new();
            for _ in 0..2 {
                let ptrs = ptrs2.clone();
                handles.push(smarm::spawn(move || {
                    let rx = ps.subscribe("t").unwrap();
                    let msg = rx.recv().unwrap();
                    ptrs.lock().unwrap().push(Arc::as_ptr(&msg) as usize);
                }));
            }
            assert!(eventually(|| ps.subscriber_count("t").unwrap() == 2));
            ps.broadcast("t", vec![1, 2, 3]).unwrap();
            for h in handles {
                h.join().unwrap();
            }
            stop_table(sup);
        });
        let v = ptrs.lock().unwrap();
        assert_eq!(v.len(), 2);
        assert_eq!(v[0], v[1], "both subscribers should see the same allocation");
    }

    #[test]
    fn broadcast_from_skips_the_sender() {
        let got = Arc::new(Mutex::new(Vec::<String>::new()));
        let got2 = got.clone();
        smarm::run(move || {
            let ps = PubSub::<String>::new("test-bus");
            let sup = start_table(ps);
            let ps_loud = ps;
            let got = got2.clone();
            let loud = smarm::spawn(move || {
                let rx = ps_loud.subscribe("room").unwrap();
                ps_loud
                    .broadcast_from(smarm::self_pid(), "room", "from loud".into())
                    .unwrap();
                // Must receive the OTHER broadcast only.
                let msg = rx.recv().unwrap();
                got.lock().unwrap().push(format!("loud got {}", *msg));
            });
            let rx = ps.subscribe("room").unwrap();
            assert!(eventually(|| ps.subscriber_count("room").unwrap() == 2));
            ps.broadcast("room", "for everyone".into()).unwrap();
            let first = rx.recv().unwrap();
            let second = rx.recv().unwrap();
            got2.lock()
                .unwrap()
                .push(format!("root got {} then {}", *first, *second));
            loud.join().unwrap();
            stop_table(sup);
        });
        let v = got.lock().unwrap();
        assert!(v.contains(&"loud got for everyone".to_string()), "{v:?}");
        // Root (not the broadcast_from sender) receives both, in order.
        assert!(
            v.contains(&"root got from loud then for everyone".to_string())
                || v.contains(&"root got for everyone then from loud".to_string()),
            "{v:?}"
        );
    }

    #[test]
    fn unsubscribe_closes_the_receiver_and_stops_delivery() {
        let ok = Arc::new(Mutex::new(false));
        let ok2 = ok.clone();
        smarm::run(move || {
            let ps = PubSub::<u32>::new("test-bus");
            let sup = start_table(ps);
            let rx = ps.subscribe("t").unwrap();
            ps.unsubscribe("t").unwrap();
            // Unsubscribe is a cast; the table holds the only sender, so
            // once processed the channel closes and recv errs.
            assert!(eventually(|| ps.subscriber_count("t").unwrap() == 0));
            assert!(rx.recv().is_err());
            ps.broadcast("t", 7).unwrap(); // no subscribers: no-op, no panic
            *ok2.lock().unwrap() = true;
            stop_table(sup);
        });
        assert!(*ok.lock().unwrap());
    }

    #[test]
    fn resubscribe_replaces_the_old_subscription() {
        let got = Arc::new(Mutex::new((0u32, false)));
        let got2 = got.clone();
        smarm::run(move || {
            let ps = PubSub::<u32>::new("test-bus");
            let sup = start_table(ps);
            let rx_old = ps.subscribe("t").unwrap();
            let rx_new = ps.subscribe("t").unwrap();
            assert_eq!(ps.subscriber_count("t").unwrap(), 1, "idempotent per (pid, topic)");
            ps.broadcast("t", 42).unwrap();
            let v = *rx_new.recv().unwrap();
            let old_closed = rx_old.recv().is_err();
            *got2.lock().unwrap() = (v, old_closed);
            stop_table(sup);
        });
        assert_eq!(*got.lock().unwrap(), (42, true));
    }

    #[test]
    fn dead_subscriber_is_pruned_by_the_monitor() {
        let ok = Arc::new(Mutex::new(false));
        let ok2 = ok.clone();
        smarm::run(move || {
            let ps = PubSub::<u32>::new("test-bus");
            let sup = start_table(ps);
            let ps2 = ps;
            let h = smarm::spawn(move || {
                let _rx = ps2.subscribe("t").unwrap();
                // Exit without ever receiving: rx drops with the stack.
            });
            h.join().unwrap();
            // No broadcast issued — this MUST be the monitor path, not
            // prune-on-send-failure.
            *ok2.lock().unwrap() = eventually(|| ps.subscriber_count("t").unwrap() == 0);
            stop_table(sup);
        });
        assert!(*ok.lock().unwrap(), "monitor Down never pruned the dead subscriber");
    }

    #[test]
    fn dropped_receiver_is_pruned_on_next_broadcast() {
        let counts = Arc::new(Mutex::new((0usize, 0usize)));
        let counts2 = counts.clone();
        smarm::run(move || {
            let ps = PubSub::<u32>::new("test-bus");
            let sup = start_table(ps);
            let rx = ps.subscribe("t").unwrap();
            drop(rx);
            let before = ps.subscriber_count("t").unwrap();
            ps.broadcast("t", 1).unwrap();
            let mut after = ps.subscriber_count("t").unwrap();
            // Casts are ordered but count is a call that can overtake
            // nothing here — same inbox, so after the broadcast cast is
            // handled. Still poll defensively.
            if after != 0 {
                after = if eventually(|| ps.subscriber_count("t").unwrap() == 0) { 0 } else { after };
            }
            *counts2.lock().unwrap() = (before, after);
            stop_table(sup);
        });
        assert_eq!(*counts.lock().unwrap(), (1, 0));
    }

    #[test]
    fn topics_are_independent() {
        let got = Arc::new(Mutex::new(Vec::<u32>::new()));
        let got2 = got.clone();
        smarm::run(move || {
            let ps = PubSub::<u32>::new("test-bus");
            let sup = start_table(ps);
            let rx_a = ps.subscribe("a").unwrap();
            ps.broadcast("b", 99).unwrap(); // nobody on b; must not reach a
            ps.broadcast("a", 1).unwrap();
            got2.lock().unwrap().push(*rx_a.recv().unwrap());
            stop_table(sup);
        });
        assert_eq!(*got.lock().unwrap(), vec![1]);
    }
}