diff --git a/src/lib.rs b/src/lib.rs index cb910ab..8004bac 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,12 +28,14 @@ pub mod conn_registry; pub mod serve; pub mod sse; pub mod ws; +pub mod pubsub; // Re-exports — what most users want at the crate root. pub use conn::{Assigns, Body, Conn, HeaderMap, HttpVersion, Method, Params, RespBody, StreamBody}; pub use plug::{Next, Pipeline, Plug}; pub use router::Router; pub use sse::{EventSender, SseClosed}; +pub use pubsub::{PubSub, PubSubDown}; pub use ws::{Message, WsClosed, WsHandler, WsSender}; pub use serve::{ serve, serve_with, serve_with_shutdown, shutdown_handle, Config, Handle, ShutdownSignal, diff --git a/src/pubsub.rs b/src/pubsub.rs new file mode 100644 index 0000000..4991b5b --- /dev/null +++ b/src/pubsub.rs @@ -0,0 +1,495 @@ +//! `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`): 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>`**: 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`, 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. +//! +//! # Why there is no `register(name)` helper (yet) +//! +//! smarm's registry maps `name → Pid`, but a `Pid` cannot be turned back +//! into a `ServerRef` (the ref *is* the inbox sender). A useful named +//! lookup therefore needs either smarm support (registry-held senders) +//! or a process-global type-erased map here — both against the grain of +//! the ratified design. Deferred; pass the handle. + +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::sync::Arc; + +use smarm::gen_server::{self, GenServer, ServerCtx}; +use smarm::{channel, Down, Pid, Receiver, Sender, ServerRef, 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 {} + +/// A clonable handle to one pub/sub instance (one topic table actor). +/// +/// Payloads are broadcast as `Arc`: one allocation per broadcast, not +/// per subscriber. +pub struct PubSub { + server: ServerRef>, +} + +impl Clone for PubSub { + fn clone(&self) -> Self { + PubSub { server: self.server.clone() } + } +} + +impl Default for PubSub { + fn default() -> Self { + Self::new() + } +} + +impl PubSub { + /// Start a fresh topic table. Must run inside the smarm runtime (it + /// spawns the table's gen_server actor). The table lives until the + /// last handle is dropped. + pub fn new() -> Self { + PubSub { server: gen_server::start(Table::new()) } + } + + /// Subscribe the **calling actor** to `topic`. Returns the receiving + /// end; messages arrive as `Arc`. 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) -> Result>, 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, + ) -> Result>, 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) -> 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) -> 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, 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, + 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) -> Result { + 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 — for introspection / registry use. + pub fn pid(&self) -> Pid { + self.server.pid() + } + + fn cast_broadcast(&self, topic: String, msg: M, skip: Option) -> Result<(), PubSubDown> { + self.server + .cast(Cast::Broadcast { topic, msg: Arc::new(msg), skip }) + .map_err(|_| PubSubDown) + } +} + +// --------------------------------------------------------------------------- +// The table gen_server +// --------------------------------------------------------------------------- + +enum Call { + Subscribe { topic: String, pid: Pid, tx: Sender> }, + Count { topic: String }, +} + +enum Reply { + Subscribed, + Count(usize), +} + +enum Cast { + Unsubscribe { topic: String, pid: Pid }, + Broadcast { topic: String, msg: Arc, skip: Option }, +} + +struct Table { + topics: HashMap>>>, + /// 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, + watcher: Option, +} + +impl Table { + fn new() -> Self { + Table { topics: HashMap::new(), monitored: HashSet::new(), watcher: None } + } +} + +impl GenServer for Table { + type Call = Call; + type Reply = Reply; + type Cast = Cast; + type Info = (); + + fn init(&mut self, ctx: &ServerCtx) { + self.watcher = Some(ctx.watcher()); + } + + fn handle_call(&mut self, request: Call) -> 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) { + 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 + } + + #[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::::new(); + let mut handles = Vec::new(); + for i in 0..2u32 { + let ps = ps.clone(); + 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(); + } + }); + 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::::new())); + let ptrs2 = ptrs.clone(); + smarm::run(move || { + let ps = PubSub::>::new(); + let mut handles = Vec::new(); + for _ in 0..2 { + let ps = ps.clone(); + 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(); + } + }); + 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::::new())); + let got2 = got.clone(); + smarm::run(move || { + let ps = PubSub::::new(); + let ps_loud = ps.clone(); + 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(); + }); + 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::::new(); + 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; + }); + 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::::new(); + 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); + }); + 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::::new(); + let ps2 = ps.clone(); + 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); + }); + 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::::new(); + 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); + }); + assert_eq!(*counts.lock().unwrap(), (1, 0)); + } + + #[test] + fn topics_are_independent() { + let got = Arc::new(Mutex::new(Vec::::new())); + let got2 = got.clone(); + smarm::run(move || { + let ps = PubSub::::new(); + 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()); + }); + assert_eq!(*got.lock().unwrap(), vec![1]); + } +}