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, Call, Cast, ConnRegistry, Reply};
use crate::net::{accept_nonblocking, bind_and_listen, OwnedFd};
use crate::plug::Pipeline;

use smarm::{ChildSpec, OneForOne, Restart, ServerRef, 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,
    pub request_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,
    /// 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>,
}

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,
            request_timeout:    Duration::from_secs(30),
            max_body_bytes:     16 * 1024 * 1024,
            drain_timeout:      Duration::from_secs(30),
            scheduler_threads:  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,
        }
    }
}

// ---------------------------------------------------------------------------
// 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,
    registry: ServerRef<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();
                smarm::spawn(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 drain_timeout = config.drain_timeout;

    let smarm_cfg = match config.scheduler_threads {
        Some(n) => smarm::Config::exact(n),
        None    => smarm::Config::default(),
    };
    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();

        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, 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());

        // 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. Same sweep discipline as listeners on the force-
        // stop path: a conn accepted just before its listener died may
        // register after the deadline, so keep force-stopping until the
        // set is empty (each pass kills everything registered; new
        // registrants are a strictly shrinking population once listeners
        // are gone).
        let _ = registry.cast(Cast::BeginDrain);
        let deadline = std::time::Instant::now() + drain_timeout;
        let mut force = false;
        loop {
            match registry.call(Call::ConnCount) {
                Ok(Reply::ConnCount(0)) => break,
                Ok(_) => {}
                Err(_) => break, // registry gone; nothing left to track
            }
            let now = std::time::Instant::now();
            if force || now >= deadline {
                force = true;
                let _ = registry.cast(Cast::ForceStopConns);
                smarm::sleep(Duration::from_millis(10));
            } else {
                smarm::sleep(Duration::from_millis(50).min(deadline - now));
            }
        }

        // 5. Our ServerRef drops here. The registry's inbox closes once
        //    the last conn's clone drops with it, and 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)
}