arcature 2026.2.1

//! Live connection registry for realtime drain (PROGRAM.md §AP2.1-8).
//!
//! Tracks live WebSocket and SSE connections so the application can drain
//! them on shutdown. The registry is **app-owned**, not a process-global
//! singleton: it is built once and stored in `AppState` (or resolved via a
//! typed service), and every connection registers itself against the one
//! value the application owns (AGENTS.md §20 — no hidden mutable global).
//!
//! The registry counts connections by `Arc<AtomicUsize>` and exposes a
//! [`Registry::drain`] future that signals every live connection to close
//! and waits for the count to reach zero (or a bound). It is intentionally
//! self-contained: the realtime module does not edit
//! `crates/arcature/src/application/shutdown.rs` this wave (the Production
//! lane owns the global drain state machine). The application calls
//! `registry.drain(...)` from its own shutdown hook; the master wires it
//! into the global drain after the Production lane lands (see the report's
//! integration-seam note).
//!
//! # Why not a global?
//!
//! A global registry would be reachable from any task without explicit
//! ownership, recreating the "hidden request-scoped global" anti-pattern
//! (§20) at connection scope. Requiring the application to hold the
//! `Registry` in `AppState` makes the ownership explicit and the drain
//! deterministic: only connections that registered against *this* registry
//! are drained, and the application knows exactly where it lives.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use tokio::sync::Notify;

use crate::realtime::error::RealtimeError;

#[derive(Debug)]
struct Shared {
    live: AtomicUsize,
    // Signaled by `drain` so every live connection task observes the
    // shutdown request. Connections select against this `Notify` (or the
    // `drain` returns a guard that wraps it) alongside their receive loop.
    drain: Notify,
}

/// A live-connection registry, shared across all connections the
/// application accepts.
///
/// Construct with [`Registry::new`]; clone cheaply into every connection
/// task. [`Registry::acquire`] returns a guard that increments the count
/// and decrements it on drop — so the count reflects only live connections
/// (no orphans). [`Registry::drain`] signals shutdown and awaits the count
/// reaching zero (or the bound).
#[derive(Clone)]
pub struct Registry {
    shared: Arc<Shared>,
}

/// A guard that increments the live-connection count on creation and
/// decrements it on drop. Dropping it also notifies the registry so a
/// `drain` waiting on the count does not miss the decrement.
#[derive(Debug)]
pub struct ConnectionGuard {
    shared: Arc<Shared>,
}

impl Registry {
    /// Create a new registry with zero live connections.
    #[must_use]
    pub fn new() -> Self {
        Self {
            shared: Arc::new(Shared {
                live: AtomicUsize::new(0),
                drain: Notify::new(),
            }),
        }
    }

    /// The number of currently-live connections.
    #[must_use]
    pub fn live_count(&self) -> usize {
        self.shared.live.load(Ordering::Relaxed)
    }

    /// Acquire a connection guard. The application calls this at the start
    /// of a connection task (after admission — origin, auth, limit —
    /// succeeds). The guard decrements the count on drop, so the
    /// application does not need a manual `release` call.
    ///
    /// # Errors
    ///
    /// Returns [`RealtimeError::ConnectionLimit`] if the count is already
    /// at `max`, so an attacker cannot open unbounded connections
    /// (AGENTS.md §29). The application maps this to a 503 before the
    /// upgrade completes.
    pub fn acquire(&self, max: usize) -> Result<ConnectionGuard, RealtimeError> {
        // Saturating add: never exceed `max` via a race. We use a CAS loop so
        // a concurrent acquire does not both succeed past the cap.
        loop {
            let current = self.shared.live.load(Ordering::Relaxed);
            if current >= max {
                return Err(RealtimeError::ConnectionLimit);
            }
            if self
                .shared
                .live
                .compare_exchange(current, current + 1, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                return Ok(ConnectionGuard {
                    shared: Arc::clone(&self.shared),
                });
            }
            // CAS failed: another acquire raced; retry the cap check.
        }
    }

    /// Signal every live connection to drain and wait for the count to reach
    /// zero or the `bound` timeout to elapse.
    ///
    /// This is the self-contained shutdown seam the application calls from
    /// its own shutdown hook. It does not touch the global engine drain
    /// (`crates/arcature/src/application/shutdown.rs`); the master wires it
    /// in after the Production lane lands.
    ///
    /// # Errors
    ///
    /// Returns [`RealtimeError::Shutdown`] with the remaining live count if
    /// the bound elapses before the count reaches zero. The connections
    /// that did not drain are still live (their tasks will be dropped by
    /// the runtime when the process exits); the error is honest about it
    /// (AGENTS.md §9).
    pub async fn drain(&self, bound: std::time::Duration) -> Result<(), RealtimeError> {
        self.shared.drain.notify_waiters();
        // Wait for the count to reach zero, polling with a short sleep so we
        // can respect the bound. A `Notify`-per-decrement would be more
        // precise but adds a per-drop allocation; the count is small and
        // the bound is the correctness guarantee.
        let result = tokio::time::timeout(bound, async {
            loop {
                if self.shared.live.load(Ordering::Relaxed) == 0 {
                    return;
                }
                // Yield so other tasks (the connection tasks finishing their
                // cleanup and dropping their guards) get to run.
                tokio::task::yield_now().await;
            }
        })
        .await;
        match result {
            Ok(()) => Ok(()),
            Err(_) => Err(RealtimeError::Shutdown {
                remaining: self.shared.live.load(Ordering::Relaxed),
            }),
        }
    }

    /// A `Notify` the connection tasks select against to observe the drain
    /// signal. The connection task should `select!` between its receive
    /// loop and `registry.drain_signal().notified()`.
    #[must_use]
    pub fn drain_signal(&self) -> &Notify {
        &self.shared.drain
    }
}

impl Default for Registry {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for Registry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Registry")
            .field("live", &self.live_count())
            .finish()
    }
}

impl ConnectionGuard {
    /// The number of live connections including this one.
    #[must_use]
    pub fn live_count(&self) -> usize {
        self.shared.live.load(Ordering::Relaxed)
    }
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        // Saturating subtract via CAS floor of 0.
        loop {
            let current = self.shared.live.load(Ordering::Relaxed);
            if current == 0 {
                break;
            }
            if self
                .shared
                .live
                .compare_exchange(current, current - 1, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                // Notify the drain waiter so it observes the decrement
                // promptly (it also polls, but the notify is a courtesy).
                self.shared.drain.notify_one();
                break;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn acquire_and_drop_track_count() {
        let reg = Registry::new();
        assert_eq!(reg.live_count(), 0);
        let g = reg.acquire(8).expect("under cap");
        assert_eq!(reg.live_count(), 1);
        assert_eq!(g.live_count(), 1);
        drop(g);
        assert_eq!(reg.live_count(), 0);
    }

    #[test]
    fn acquire_enforces_cap() {
        let reg = Registry::new();
        let _g1 = reg.acquire(2).expect("first");
        let _g2 = reg.acquire(2).expect("second");
        let outcome = reg.acquire(2);
        assert!(
            matches!(outcome, Err(RealtimeError::ConnectionLimit)),
            "{outcome:?}"
        );
        // Dropping one frees a slot.
        drop(_g2);
        let _g3 = reg.acquire(2).expect("slot freed");
    }

    #[tokio::test]
    async fn drain_completes_when_connections_drop() {
        let reg = Registry::new();
        let _g = reg.acquire(8).expect("under cap");
        assert_eq!(reg.live_count(), 1);
        // Drain in a background task that will complete once the guard drops.
        let reg2 = reg.clone();
        let drain_task =
            tokio::spawn(async move { reg2.drain(std::time::Duration::from_secs(5)).await });
        // Give the drain task a moment to start waiting.
        tokio::task::yield_now().await;
        // Drop the guard; drain should complete.
        drop(_g);
        let result = tokio::time::timeout(std::time::Duration::from_secs(5), drain_task)
            .await
            .expect("drain task did not hang")
            .expect("drain task did not panic");
        assert!(result.is_ok(), "{result:?}");
        assert_eq!(reg.live_count(), 0);
    }

    #[tokio::test]
    async fn drain_times_out_with_remaining_count() {
        let reg = Registry::new();
        // Keep a guard alive across the drain so it must time out.
        let _g = reg.acquire(8).expect("under cap");
        let result = reg.drain(std::time::Duration::from_millis(50)).await;
        assert!(
            matches!(result, Err(RealtimeError::Shutdown { remaining: 1 })),
            "{result:?}"
        );
    }

    #[test]
    fn registry_is_send_sync_clone_for_appstate() {
        fn assert_send_sync_clone<T: Send + Sync + Clone>() {}
        assert_send_sync_clone::<Registry>();
    }
}