urus/src/conn_registry.rs

//! 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<Pid, ConnState>,
    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<ConnRegistry> {
    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<ConnRegistry>,
    pid:      Pid,
    make:     fn(Pid) -> Cast,
}

impl DeregisterGuard {
    pub fn new(registry: ServerRef<ConnRegistry>, 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));
    }
}