urus/src/serve.rs

//! Listener pool and the `serve` entry point.
//!
//! A small fixed pool of listener actors share the same TCP listen fd (via
//! `dup`) — each blocks in non-blocking `accept4` + `wait_readable` on its
//! own copy. When a connection arrives the listener spawns a connection
//! actor with the `OwnedFd` and immediately returns to `accept`. No
//! coordination needed; the kernel serialises `accept` calls across the fds.
//!
//! Sharing via `dup` rather than the same fd is deliberate — Linux's
//! `accept4` is thread-safe on a single fd, but dup'ing per-listener keeps
//! each actor's epoll registration local to its own RawFd value (so smarm's
//! `waiters: HashMap<RawFd, Pid>` doesn't see collisions between listeners
//! waiting on "the same fd").

use crate::conn_actor::{run_connection, ConnLimits};
use crate::conn_registry::{self, ConnRegistry};
use crate::net::{accept_nonblocking, bind_and_listen, OwnedFd};
use crate::plug::Pipeline;

use smarm::{ChildSpec, OneForOne, Restart, GenServerRef, Strategy};

use std::io::{self, ErrorKind};
use std::net::{SocketAddr, ToSocketAddrs};
use std::os::fd::RawFd;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

#[derive(Clone, Debug)]
pub struct Config {
    pub addr:               SocketAddr,
    pub listener_pool:      usize,
    pub keep_alive_timeout: Duration,
    pub max_header_count:   usize,
    pub read_buf_size:      usize,
    /// Wall-clock budget for reading the request HEAD (from first byte to
    /// full head parse). Kept short — an incomplete head is the classic
    /// slowloris. See `ConnLimits::head_timeout`.
    pub head_timeout:       Duration,
    /// Absolute wall-clock cap on reading the request BODY (from head-parse
    /// to full body). Sized for slow links, so much larger than
    /// `head_timeout`. See `ConnLimits::body_timeout`.
    pub body_timeout:       Duration,
    /// Burst size that resets the body stall clock. A body dribbling fewer
    /// than this per `body_stall_timeout` window is evicted — the slowloris
    /// / slow-legit discriminator. See `ConnLimits::body_burst_bytes`.
    pub body_burst_bytes:   usize,
    /// Max time since the last qualifying body burst before eviction;
    /// backstopped by `body_timeout`. See `ConnLimits::body_stall_timeout`.
    pub body_stall_timeout: Duration,
    /// Per-write budget for response bytes (the fixed head+body write, and
    /// each streamed chunk). See `ConnLimits::write_timeout`.
    pub write_timeout:      Duration,
    pub max_body_bytes:     usize,
    /// How long a graceful shutdown waits for in-flight requests before
    /// force-stopping the remaining connections. Idle keep-alive
    /// connections are closed immediately on shutdown and do not run the
    /// clock out.
    pub drain_timeout:      Duration,
    /// WebSocket: cap on a single frame's payload (header-checked
    /// before buffering; violation closes 1009).
    pub max_frame_payload:  usize,
    /// WebSocket: cap on a complete reassembled message (spans
    /// fragments; violation closes 1009).
    pub max_message_bytes:  usize,
    /// Number of smarm scheduler OS threads. `None` means smarm's default
    /// (one per CPU). Set this to a small fixed number in tests so multiple
    /// concurrent test servers don't oversubscribe the host.
    pub scheduler_threads:  Option<usize>,
    /// Stack reserve (RFC 019 `smarm::SpawnOpts::stack_reserve`) given to
    /// each per-connection actor. Request handlers routinely pull in
    /// application code — DB drivers, (de)compression, templating — whose
    /// stack needs comfortably exceed smarm's bare-actor default of 64 KiB
    /// (the exact shape of bug this exists to head off; see smarm RFC 019).
    /// Default: 256 KiB. The reserve is virtual/demand-paged, so raising it
    /// costs address space, not RSS, until a handler actually uses it.
    pub conn_stack_reserve: usize,
    /// Maximum concurrently-live actors — smarm's fixed slot slab, allocated
    /// once at init. Each connection is one actor, so this is also the hard
    /// cap on concurrent connections. `None` uses smarm's default (16_384).
    /// Slots are ~256 B, so raising this is cheap relative to per-connection
    /// stacks; size it to peak concurrent connections.
    pub max_actors: Option<usize>,
}

/// Default per-connection actor stack reserve (see [`Config::conn_stack_reserve`]).
pub const DEFAULT_CONN_STACK_RESERVE: usize = 256 * 1024;

impl Config {
    pub fn new(addr: SocketAddr) -> Self {
        let pool = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(2)
            .max(2);
        Self {
            addr,
            listener_pool:      pool,
            keep_alive_timeout: Duration::from_secs(60),
            max_header_count:   64,
            read_buf_size:      8 * 1024,
            head_timeout:       Duration::from_secs(30),
            body_timeout:       Duration::from_secs(300),
            body_burst_bytes:   4 * 1024,
            body_stall_timeout: Duration::from_secs(20),
            write_timeout:      Duration::from_secs(30),
            max_body_bytes:     16 * 1024 * 1024,
            drain_timeout:      Duration::from_secs(30),
            max_frame_payload:  1024 * 1024,
            max_message_bytes:  4 * 1024 * 1024,
            scheduler_threads:  None,
            conn_stack_reserve: DEFAULT_CONN_STACK_RESERVE,
            max_actors:         None,
        }
    }

    fn to_conn_limits(&self) -> ConnLimits {
        ConnLimits {
            max_headers:        self.max_header_count,
            initial_read_buf:   self.read_buf_size,
            max_head_bytes:     64 * 1024,
            max_body_bytes:     self.max_body_bytes,
            keep_alive_timeout: self.keep_alive_timeout,
            head_timeout:       self.head_timeout,
            body_timeout:       self.body_timeout,
            body_burst_bytes:   self.body_burst_bytes,
            body_stall_timeout: self.body_stall_timeout,
            write_timeout:      self.write_timeout,
            max_frame_payload:  self.max_frame_payload,
            max_message_bytes:  self.max_message_bytes,
        }
    }
}

// ---------------------------------------------------------------------------
// config-file: TOML overlay for tuning knobs
// ---------------------------------------------------------------------------
//
// urus is a library, so it never presumes a config-file path or reads the
// environment — the embedding binary decides where a file lives and hands
// the text here. This overlays a sparse TOML document onto an existing
// `Config` (built with an addr the binary chose): only the keys present are
// applied, everything else keeps the compiled default. Durations are
// integer seconds. Unknown keys are a hard error so a typo is loud, not a
// silent no-op.
//
// Scope for now: the slowloris-tuning knobs only. Migrating the rest of the
// Config surface into the file is a separate, additive job (the loader
// mechanism is general — it just extends `TomlOverrides`).

/// Error from [`Config::with_toml_str`]: the TOML failed to parse or carried
/// an unknown/mistyped key.
#[cfg(feature = "config-file")]
#[derive(Debug)]
pub enum ConfigError {
    Toml(String),
}

#[cfg(feature = "config-file")]
impl std::fmt::Display for ConfigError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ConfigError::Toml(m) => write!(f, "config TOML error: {m}"),
        }
    }
}

#[cfg(feature = "config-file")]
impl std::error::Error for ConfigError {}

#[cfg(feature = "config-file")]
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct TomlOverrides {
    head_timeout_secs:       Option<u64>,
    body_timeout_secs:       Option<u64>,
    body_burst_bytes:        Option<usize>,
    body_stall_timeout_secs: Option<u64>,
}

#[cfg(feature = "config-file")]
impl Config {
    /// Overlay a TOML document of tuning knobs onto this config (sparse:
    /// only the keys present are applied). Durations are integer seconds.
    ///
    /// Recognized keys: `head_timeout_secs`, `body_timeout_secs`,
    /// `body_burst_bytes`, `body_stall_timeout_secs`. Unknown keys error.
    /// Other `Config` knobs are not yet file-configurable.
    pub fn with_toml_str(mut self, s: &str) -> Result<Self, ConfigError> {
        let o: TomlOverrides =
            toml::from_str(s).map_err(|e| ConfigError::Toml(e.to_string()))?;
        if let Some(v) = o.head_timeout_secs {
            self.head_timeout = Duration::from_secs(v);
        }
        if let Some(v) = o.body_timeout_secs {
            self.body_timeout = Duration::from_secs(v);
        }
        if let Some(v) = o.body_burst_bytes {
            self.body_burst_bytes = v;
        }
        if let Some(v) = o.body_stall_timeout_secs {
            self.body_stall_timeout = Duration::from_secs(v);
        }
        Ok(self)
    }
}

// ---------------------------------------------------------------------------
// dup helper
// ---------------------------------------------------------------------------

fn dup_fd(fd: RawFd) -> io::Result<OwnedFd> {
    let new_fd = unsafe { libc::fcntl(fd, libc::F_DUPFD_CLOEXEC, 0) };
    if new_fd < 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(OwnedFd::from_raw(new_fd))
}

// ---------------------------------------------------------------------------
// listener actor body
// ---------------------------------------------------------------------------

/// Test-only fault injection. When nonzero, the next accept-loop iteration
/// of whichever listener gets there first decrements this and panics —
/// *before* calling `accept`, so a pending connection stays in the kernel
/// backlog and must be picked up by the restarted listener. Cost when idle
/// is one relaxed load per accept-loop iteration (each of which already
/// pays a syscall). Not public API.
#[doc(hidden)]
pub static INJECT_LISTENER_PANICS: AtomicU32 = AtomicU32::new(0);

fn listener_loop(
    listener: Arc<OwnedFd>,
    pipeline: Pipeline,
    limits: ConnLimits,
    conn_stack_reserve: usize,
    registry: GenServerRef<ConnRegistry>,
    shutdown: Arc<AtomicBool>,
) {
    let fd = listener.as_raw();

    loop {
        // Self-termination on shutdown — the ONLY way a listener exits at
        // shutdown, and deliberately a normal return: under
        // `Restart::Transient` a normal exit is terminal, so the
        // supervisor's active-count drains and `sup.run()` returns on its
        // own. No pid is ever `request_stop`ped, which sidesteps both the
        // spawn-to-registration race of self-announced pids and smarm's
        // lossy stop-while-QUEUED window (a flag is wake-free and
        // race-free; a freshly spawned listener observes it on its very
        // first iteration, a parked one within LISTENER_TICK).
        if shutdown.load(Ordering::Relaxed) {
            return;
        }
        if INJECT_LISTENER_PANICS
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1))
            .is_ok()
        {
            panic!("urus: injected listener panic (test hook)");
        }
        match accept_nonblocking(fd) {
            Ok(client) => {
                // Hand the fd off to a new connection actor. spawn() is
                // cheap on smarm — it's a single Vec push under the
                // shared lock.
                let p = pipeline.clone();
                let l = limits;
                let r = registry.clone();
                let opts = smarm::SpawnOpts {
                    stack_reserve: Some(conn_stack_reserve),
                    ..smarm::SpawnOpts::default()
                };
                smarm::spawn_with(opts, move || run_connection(client, p, l, r));
            }
            Err(e) if e.kind() == ErrorKind::WouldBlock => {
                // No pending connection. Park until the listener is
                // readable or the tick elapses; either way we come back
                // around through the shutdown-flag check above.
                match smarm::wait_readable_timeout(fd, LISTENER_TICK) {
                    Ok(_ready) => {} // ready or tick — loop re-checks, retries accept
                    Err(we) => {
                        // epoll registration failed. Under Transient
                        // supervision a normal return is terminal but a
                        // panic restarts us — and a failed wait IS
                        // abnormal, so panic: a transient failure (e.g.
                        // EMFILE on the epoll set) heals by restart
                        // instead of silently shrinking the pool. (smarm
                        // catches actor panics in the trampoline; this is
                        // a Signal::Panic to the supervisor, not process
                        // noise.)
                        panic!("urus: listener wait_readable failed: {we}");
                    }
                }
            }
            Err(e) if e.kind() == ErrorKind::Interrupted => {
                continue;
            }
            Err(e) => {
                // EMFILE / ENFILE / ECONNABORTED etc. Log and continue;
                // the system may recover.
                eprintln!("urus: accept error: {e}");
                // Small backoff via smarm's sleep to avoid spinning if
                // the error is sticky.
                smarm::sleep(Duration::from_millis(10));
            }
        }
    }
    // The Arc clone we were started with drops here (normal exit or
    // unwind), but the ChildSpec factory holds another — the fd outlives
    // any one incarnation of this listener.
}

// ---------------------------------------------------------------------------
// Handle / ShutdownSignal — graceful shutdown plumbing.
// ---------------------------------------------------------------------------

/// How often the root actor polls for a shutdown signal (see
/// `serve_with_shutdown` for why this is a poll); bounds shutdown latency.
const SHUTDOWN_POLL: Duration = Duration::from_millis(100);

/// Listener accept-waits are timed at this tick; the shutdown flag is
/// observed at the top of every accept-loop iteration, so this bounds how
/// long a fully idle listener takes to notice shutdown. It is the SOLE
/// listener-shutdown mechanism (no `request_stop` — see the shutdown
/// sequence notes). Idle cost: one timer wake per listener per tick.
const LISTENER_TICK: Duration = Duration::from_millis(250);

/// A clonable trigger for graceful shutdown. Safe to use from any OS
/// thread (the send only enqueues; the serving side polls), e.g. from a
/// signal-handling thread.
#[derive(Clone)]
pub struct Handle {
    tx: smarm::Sender<()>,
}

impl Handle {
    /// Begin graceful shutdown: stop accepting, close idle keep-alive
    /// connections, drain in-flight requests up to `Config.drain_timeout`,
    /// then force-stop stragglers. `serve_with_shutdown` returns once the
    /// runtime has wound down. Idempotent; extra calls are no-ops.
    pub fn shutdown(&self) {
        let _ = self.tx.send(());
    }
}

/// The receiving half consumed by [`serve_with_shutdown`].
pub struct ShutdownSignal {
    rx: smarm::Receiver<()>,
}

/// Create a connected [`Handle`]/[`ShutdownSignal`] pair.
pub fn shutdown_handle() -> (Handle, ShutdownSignal) {
    let (tx, rx) = smarm::channel();
    (Handle { tx }, ShutdownSignal { rx })
}

// ---------------------------------------------------------------------------
// serve_with_shutdown — main entry. Boots smarm, supervises listeners,
// blocks until shutdown.
// ---------------------------------------------------------------------------
//
// Boots an smarm runtime (one OS thread per CPU by default — see smarm's
// `Config::default()`). The root actor starts the connection registry and
// a one-for-one supervisor over the listener pool, then blocks on the
// shutdown signal. A panicking listener is restarted on the same
// (still-open) fd instead of silently shrinking the accept pool.
// Connection actors stay unsupervised bare spawns — per spec, a connection
// is cheap and its failure is local: a 500 path, not a restart path.
// (`spawn` from a listener parents the conn actor under that listener,
// which never registers a supervisor channel, so conn deaths are invisible
// to the pool supervisor by construction.)
//
// Shutdown sequence (drain-then-stop, the v0.2 chunk 2 decision):
//   1. Set the shared shutdown flag. Every listener observes it at the
//      top of its accept loop (parks are timed at LISTENER_TICK) and
//      returns normally. Under `Restart::Transient` a normal exit is
//      terminal, so no restart happens and the supervisor's active count
//      drains to zero.
//   2. Join the supervisor: `sup.run()` returns on its own once every
//      listener has exited. The join is therefore the barrier "no new
//      connections can ever be accepted" — listener fds are closed (the
//      last Arc clones drop with the supervisor's ChildSpecs), and the
//      kernel refuses new connects.
//
//      Why a flag and not `request_stop`: stopping pids that announce
//      themselves races the spawn-to-registration gap (a fast shutdown
//      CAN beat a fresh listener to the registry), and smarm's
//      `request_stop` is lossy against a QUEUED actor that then parks
//      without passing an observation point — found the hard way; see
//      the chunk 2 commit message. The flag is wake-free and race-free.
//   3. BeginDrain: the registry stops idle conns now and each remaining
//      conn the moment it finishes its in-flight request.
//   4. Poll until no conns remain or the drain deadline passes; past the
//      deadline, ForceStopConns every tick until the set empties (a conn
//      accepted just before listener death may register late).
//   5. Root returns. `rt.run` itself returns only when every actor has
//      exited (force-stopped conns unwind through their fd waits safely —
//      smarm's 06-10 io fix — and close their sockets via OwnedFd::drop).

pub fn serve_with_shutdown(
    config: Config,
    pipeline: Pipeline,
    signal: ShutdownSignal,
) -> io::Result<()> {
    let listener = bind_and_listen(config.addr)?;
    println!("urus: listening on {}", config.addr);

    // One connection-actor-spawning loop per listener pool slot. Each gets
    // its own dup'd fd so epoll registrations don't collide. Each fd is
    // owned by its ChildSpec's factory closure (via `Arc`): a restarted
    // listener re-enters `accept`/`wait_readable` on the same fd — no
    // re-dup, no window where the slot has no fd.
    let mut listener_fds = Vec::with_capacity(config.listener_pool);
    listener_fds.push(Arc::new(listener)); // primary keeps the original

    for _ in 1..config.listener_pool {
        let dup = dup_fd(listener_fds[0].as_raw())?;
        listener_fds.push(Arc::new(dup));
    }

    let limits             = config.to_conn_limits();
    let conn_stack_reserve = config.conn_stack_reserve;
    let drain_timeout      = config.drain_timeout;

    let mut smarm_cfg = match config.scheduler_threads {
        Some(n) => smarm::Config::exact(n),
        None    => smarm::Config::default(),
    };
    if let Some(m) = config.max_actors {
        smarm_cfg = smarm_cfg.max_actors(m);
    }
    let rt = smarm::init(smarm_cfg);
    // Listener self-termination flag — see the shutdown sequence below.
    let shutdown_flag = Arc::new(AtomicBool::new(false));

    rt.run(move || {
        // Registry first: listeners and conns cast into it from birth.
        let registry = conn_registry::start(drain_timeout);

        let mut sup = OneForOne::new().strategy(Strategy::OneForOne);
        for (i, lfd) in listener_fds.into_iter().enumerate() {
            let p = pipeline.clone();
            let r = registry.clone();
            let sf = shutdown_flag.clone();
            sup = sup.child(ChildSpec::new(Restart::Transient, move || {
                println!("urus: listener {} starting", i);
                listener_loop(lfd.clone(), p.clone(), limits, conn_stack_reserve, r.clone(), sf.clone());
            }));
        }
        // Default intensity (3 per 5s) applies; a listener crash-looping
        // faster than that trips the cap and tears the pool down — loud
        // failure over a zombie server.
        let sup_h = smarm::spawn(move || sup.run());
        // The old `urus.server` / `urus.listener.{i}` name registrations
        // are gone with smarm's RFC 014 registry rework: `register` is now
        // `(Name<M>, Sender<M>)`, self-only — a name is a typed messaging
        // endpoint, not a pid tag. urus's bindings were introspection-only
        // with no channel behind them, so they were dropped rather than
        // faked with a unit channel. A real messageable `urus.server`
        // name is in the icebox (ROADMAP.md).

        // Block until told to shut down. We poll `try_recv` + `sleep`
        // rather than parking in `recv`: a smarm `Sender::send` from a
        // foreign OS thread (no runtime in its TLS) enqueues fine but its
        // unpark is a `try_with_runtime` no-op — a parked receiver would
        // never wake. The timer wake comes from inside the runtime, so the
        // poll sees the message within one interval. (A cross-thread-safe
        // unpark is a smarm roadmap candidate; this poll dies with it.)
        // If every Handle was dropped the channel closes and no shutdown
        // can ever arrive: serve forever, exactly v1's semantics.
        loop {
            match signal.rx.try_recv() {
                Ok(Some(())) => break,
                Ok(None)     => smarm::sleep(SHUTDOWN_POLL),
                Err(_)       => smarm::sleep(Duration::from_secs(3600)),
            }
        }

        // ----- Shutdown. -----
        // 1 + 2. Flag the listeners down and join the supervisor; the
        //        join returns once every listener has exited normally
        //        (Transient: normal exit is terminal). After this point
        //        no connection can ever be accepted again.
        shutdown_flag.store(true, Ordering::Relaxed);
        let _ = sup_h.join();

        // 3 + 4. Drain: the registry owns the whole protocol now (idle
        //    stopped immediately, busy until its internal drain_timeout
        //    timer, late registrants stopped on arrival — see
        //    conn_registry docs). `GenServerRef::shutdown()` delivers the
        //    request and blocks on a monitor until the registry has
        //    stopped itself, which it does only once the conn set is
        //    empty: this line IS the "every connection is gone" barrier.
        registry.shutdown();

        // 5. Root returns; the runtime winds down when the last actor
        //    exits.
    });

    Ok(())
}

// ---------------------------------------------------------------------------
// serve_with / serve — convenience entries without a shutdown handle.
// ---------------------------------------------------------------------------

pub fn serve_with(config: Config, pipeline: Pipeline) -> io::Result<()> {
    // The Handle is dropped immediately: shutdown can never be signalled
    // and the server runs until externally killed (v1 semantics).
    let (_handle, signal) = shutdown_handle();
    serve_with_shutdown(config, pipeline, signal)
}

// ---------------------------------------------------------------------------
// serve — convenience over serve_with.
// ---------------------------------------------------------------------------

pub fn serve(addr: impl ToSocketAddrs, pipeline: Pipeline) -> io::Result<()> {
    let addr = addr
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| io::Error::new(ErrorKind::InvalidInput, "no addresses resolved"))?;
    serve_with(Config::new(addr), pipeline)
}

#[cfg(all(test, feature = "config-file"))]
mod config_file_tests {
    use super::*;

    fn base() -> Config {
        Config::new("127.0.0.1:0".parse().unwrap())
    }

    #[test]
    fn toml_empty_keeps_defaults() {
        let d = base();
        let c = base().with_toml_str("").unwrap();
        assert_eq!(c.head_timeout, d.head_timeout);
        assert_eq!(c.body_timeout, d.body_timeout);
        assert_eq!(c.body_burst_bytes, d.body_burst_bytes);
        assert_eq!(c.body_stall_timeout, d.body_stall_timeout);
    }

    #[test]
    fn toml_partial_overrides_only_named() {
        let d = base();
        let c = base().with_toml_str("head_timeout_secs = 5").unwrap();
        assert_eq!(c.head_timeout, Duration::from_secs(5)); // overridden
        assert_eq!(c.body_timeout, d.body_timeout);         // default kept
        assert_eq!(c.body_burst_bytes, d.body_burst_bytes); // default kept
        assert_eq!(c.body_stall_timeout, d.body_stall_timeout);
    }

    #[test]
    fn toml_full_overrides_all() {
        let c = base()
            .with_toml_str(
                "head_timeout_secs = 10\n\
                 body_timeout_secs = 120\n\
                 body_burst_bytes = 8192\n\
                 body_stall_timeout_secs = 15\n",
            )
            .unwrap();
        assert_eq!(c.head_timeout, Duration::from_secs(10));
        assert_eq!(c.body_timeout, Duration::from_secs(120));
        assert_eq!(c.body_burst_bytes, 8192);
        assert_eq!(c.body_stall_timeout, Duration::from_secs(15));
    }

    #[test]
    fn toml_unknown_key_errors() {
        // A mistyped/unknown key is a hard error, not a silent no-op.
        let e = base().with_toml_str("body_timeout_sec = 120"); // typo: missing 's'
        assert!(e.is_err(), "unknown key should error");
    }

    #[test]
    fn toml_malformed_errors() {
        let e = base().with_toml_str("this is not = valid = toml");
        assert!(e.is_err(), "malformed TOML should error");
    }
}