//! Connection registry — the shutdown coordinator (v0.2 chunk 2). //! //! A `gen_server` tracking live connection pids and their busy/idle //! state. Connection actors self-register as their first action and //! self-deregister via a drop guard (so panic unwinds deregister too); //! both casts come from the same sender, so Started always precedes //! Ended in the inbox. (The roadmap sketched the *listener* casting //! `{Started, pid}`, but then a short-lived conn's Ended could overtake //! its Started and leak a dead pid into the set forever. Self- //! registration makes the order a per-sender FIFO guarantee instead of //! a race.) //! //! Listeners are deliberately NOT tracked here: they shut down via a //! shared flag + timed accept-waits (see `serve`), never via //! `request_stop` — stopping pids that announce themselves is racy (the //! spawn-to-registration gap), and smarm's `request_stop` is lossy //! against an actor that is QUEUED and then parks without passing an //! observation point (see the v0.2 shutdown notes in the commit //! message). Connections don't suffer this in practice: every stop the //! registry issues targets a pid that just sent us a cast (so it is //! running or parked, both covered), and the force-stop path re-sweeps //! until the set empties. //! //! Drain protocol: `BeginDrain` stops every idle connection immediately //! and flips the registry into draining mode, in which any connection //! that *becomes* idle (finishes its in-flight request) is stopped on the //! spot. Busy connections are left to finish; `ForceStopConns` (sent by //! `serve` at the drain deadline) stops whatever remains. `request_stop` //! unwinds a conn actor parked in `wait_readable` safely (smarm's 06-10 //! io fix) and `OwnedFd::drop` closes its socket on the way out. //! //! This server is also the planned introspection point for ws/channels //! (roadmap v0.4+), which is why it exists as its own module rather than //! being inlined into `serve`. use smarm::{GenServer, Pid, ServerBuilder, ServerRef}; use std::collections::HashMap; // --------------------------------------------------------------------------- // Messages // --------------------------------------------------------------------------- pub enum Cast { /// A connection actor started (self-registered, initially idle: it has /// not parsed a request head yet). ConnStarted(Pid), /// Parsed a request head; a response is now owed. ConnBusy(Pid), /// Response written; parked (or about to park) waiting for the next /// keep-alive request. ConnIdle(Pid), ConnEnded(Pid), /// Stop idle conns now and stop each remaining conn as it goes idle. BeginDrain, /// Drain deadline passed: stop every remaining conn. ForceStopConns, } pub enum Call { ConnCount, } pub enum Reply { ConnCount(usize), } // --------------------------------------------------------------------------- // Server // --------------------------------------------------------------------------- #[derive(Clone, Copy, PartialEq, Eq)] enum ConnState { Busy, Idle, } #[derive(Default)] pub struct ConnRegistry { conns: HashMap, draining: bool, } impl GenServer for ConnRegistry { type Call = Call; type Reply = Reply; type Cast = Cast; type Info = (); fn handle_call(&mut self, request: Call) -> Reply { match request { Call::ConnCount => Reply::ConnCount(self.conns.len()), } } fn handle_cast(&mut self, request: Cast) { match request { Cast::ConnStarted(pid) => { self.conns.insert(pid, ConnState::Idle); if self.draining { // Listener-stop race: this conn was accepted just // before its listener died. Drain means no new work. smarm::request_stop(pid); } } Cast::ConnBusy(pid) => { if let Some(s) = self.conns.get_mut(&pid) { *s = ConnState::Busy; } } Cast::ConnIdle(pid) => { if let Some(s) = self.conns.get_mut(&pid) { *s = ConnState::Idle; if self.draining { // Finished its in-flight request; nothing more is // owed. The pid leaves the map via its drop // guard's ConnEnded once the unwind completes. smarm::request_stop(pid); } } } Cast::ConnEnded(pid) => { self.conns.remove(&pid); } Cast::BeginDrain => { self.draining = true; for (pid, state) in &self.conns { if *state == ConnState::Idle { smarm::request_stop(*pid); } } } Cast::ForceStopConns => { for pid in self.conns.keys() { smarm::request_stop(*pid); } } } } } pub fn start() -> ServerRef { ServerBuilder::new(ConnRegistry::default()).start() } // --------------------------------------------------------------------------- // Drop guard — self-deregistration on any exit path. // --------------------------------------------------------------------------- /// Casts `make(pid)` on drop. Runs on normal return, `request_stop` /// unwind, and panic unwind alike; the cast is infallible from the /// guard's perspective (a dead registry just returns an ignored Err). pub struct DeregisterGuard { registry: ServerRef, pid: Pid, make: fn(Pid) -> Cast, } impl DeregisterGuard { pub fn new(registry: ServerRef, pid: Pid, make: fn(Pid) -> Cast) -> Self { Self { registry, pid, make } } } impl Drop for DeregisterGuard { fn drop(&mut self) { let _ = self.registry.cast((self.make)(self.pid)); } }