arcly_http/realtime/

connection.rs

//! Lock-free connection & broadcasting registry.
//!
//! ## Topology
//!
//! The hot path for real-time servers is *broadcast*: one inbound event fans
//! out to thousands of sockets. A naive `Mutex<HashMap<Id, Sender>>` serialises
//! every send behind one lock — the classic C10K/C100K bottleneck.
//!
//! `arcly-http` instead uses [`dashmap::DashMap`], which shards the key space
//! across N independent locks (N = `4 * num_cpus` by default). Two broadcasts
//! touching different shards never contend. Each connection owns a **bounded**
//! `mpsc::Sender`; sending is a wait-free `try_send` that never blocks the
//! frame-processing task. The per-socket writer drains its own receiver, and a
//! client that can't keep up is **evicted** when its queue fills — a slow
//! consumer costs at most `ws_outbound_buffer` frames of memory, never an
//! unbounded buffer.
//!
//! ## Encapsulation
//!
//! No `axum`/`tower` types appear here. The registry speaks [`WsMessage`], an
//! arcly-owned enum. The websocket boundary in [`super::ws`] is the single
//! place that translates `WsMessage` ↔ `axum::extract::ws::Message`.

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

use dashmap::{DashMap, DashSet};
use serde::Serialize;
use tokio::sync::mpsc;

use crate::web::context::Claims;

/// Monotonic, process-unique connection identifier.
pub type ConnId = u64;

/// An arcly-owned outbound frame. The websocket boundary maps this onto the
/// concrete transport frame type — keeping `axum` out of the public surface.
///
/// `Text` holds `Arc<str>` so broadcast can enqueue N pointer copies from a
/// single allocation rather than cloning the string body N times.
#[derive(Clone, Debug)]
pub enum WsMessage {
    /// UTF-8 text frame (our event envelope is always JSON text).
    Text(Arc<str>),
    /// Server-initiated keep-alive ping (liveness probe through NATs/proxies).
    Ping,
    /// Graceful close.
    Close,
}

/// Per-connection registry row. `tx` is the head of that socket's **bounded**
/// outbound queue; `member_rooms` tracks which rooms to evict the id from on
/// disconnect; `last_seen` (unix secs) is touched by the reader on every
/// inbound frame and consulted by the idle sweeper.
struct ConnEntry {
    tx: mpsc::Sender<WsMessage>,
    member_rooms: DashSet<String>,
    last_seen: AtomicU64,
    #[allow(dead_code)]
    claims: Option<Arc<Claims>>,
}

#[inline]
fn unix_now() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

/// Sharded, lock-free-on-the-hot-path connection registry.
///
/// One instance is created at launch, leaked to `&'static`, and shared by every
/// gateway route. All mutating operations are `O(1)` amortised and touch only a
/// single shard; `broadcast` is `O(connections)` with no global lock.
pub struct ConnectionRegistry {
    next_id: AtomicU64,
    conns: DashMap<ConnId, ConnEntry>,
    rooms: DashMap<String, DashSet<ConnId>>,
}

impl ConnectionRegistry {
    pub fn new() -> Self {
        Self {
            next_id: AtomicU64::new(1),
            conns: DashMap::new(),
            rooms: DashMap::new(),
        }
    }

    /// Register a freshly-upgraded socket. Returns its assigned [`ConnId`].
    ///
    /// `tx` must be the sending half of a **bounded** channel — the queue
    /// depth is the slow-client memory ceiling (`LaunchConfig::ws_outbound_buffer`).
    pub fn register(&self, tx: mpsc::Sender<WsMessage>, claims: Option<Arc<Claims>>) -> ConnId {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        self.conns.insert(
            id,
            ConnEntry {
                tx,
                member_rooms: DashSet::new(),
                last_seen: AtomicU64::new(unix_now()),
                claims,
            },
        );
        metrics::gauge!("ws_connections").increment(1.0);
        id
    }

    /// Record inbound activity for the idle sweeper. Called by the reader on
    /// every frame (including pongs); one relaxed store.
    #[inline]
    pub fn touch(&self, id: ConnId) {
        if let Some(entry) = self.conns.get(&id) {
            entry.last_seen.store(unix_now(), Ordering::Relaxed);
        }
    }

    /// Enqueue to one entry; on a **full queue** the client is too slow to
    /// keep up — evict it (drop the entry; the writer drains the backlog,
    /// sees the closed channel, and tears the socket down). Unbounded growth
    /// in server memory is never an option.
    ///
    /// Returns `false` when the entry was evicted (caller loops may skip it).
    fn enqueue_or_evict(&self, id: ConnId, entry: &ConnEntry, msg: WsMessage) -> bool {
        match entry.tx.try_send(msg) {
            Ok(()) => {
                metrics::counter!("ws_messages_out_total").increment(1);
                true
            }
            Err(mpsc::error::TrySendError::Full(_)) => {
                metrics::counter!("ws_slow_client_evictions_total").increment(1);
                tracing::warn!(conn = id, "WS outbound queue full — evicting slow client");
                false
            }
            // Receiver already gone — reader will unregister shortly.
            Err(mpsc::error::TrySendError::Closed(_)) => true,
        }
    }

    /// Send a Close frame to **every** live socket. Used by graceful
    /// shutdown: WS connections otherwise keep axum's drain phase open
    /// indefinitely (until clients disconnect or the supervisor SIGKILLs),
    /// which would mean plugin `on_shutdown` hooks never run in production.
    /// Each socket's writer task sends the frame and exits; readers follow.
    pub fn close_all(&self) -> usize {
        let mut n = 0;
        for entry in self.conns.iter() {
            // try_send: a full queue can't take the Close, but eviction by
            // the next broadcast — or the drain deadline — handles it.
            let _ = entry.value().tx.try_send(WsMessage::Close);
            n += 1;
        }
        n
    }

    /// Close connections with no inbound activity for `max_idle_secs`.
    /// Dead TCP links (NAT timeouts, vanished mobile clients) never send a
    /// Close frame — without reaping they linger in the registry forever,
    /// eating queue memory on every broadcast. Returns the ids reaped.
    ///
    /// Pair with a server ping (`ws_ping_interval`): pings provoke pongs,
    /// pongs touch `last_seen`, so only truly dead links exceed the window.
    pub fn sweep_idle(&self, max_idle_secs: u64) -> Vec<ConnId> {
        let now = unix_now();
        let stale: Vec<ConnId> = self
            .conns
            .iter()
            .filter(|e| {
                now.saturating_sub(e.value().last_seen.load(Ordering::Relaxed)) > max_idle_secs
            })
            .map(|e| *e.key())
            .collect();
        for id in &stale {
            metrics::counter!("ws_idle_reaped_total").increment(1);
            if let Some(entry) = self.conns.get(id) {
                let _ = entry.tx.try_send(WsMessage::Close);
            }
            // If the queue was full the Close never lands — force the writer
            // down by dropping the entry (channel closes, writer exits).
            self.unregister(*id);
        }
        stale
    }

    /// Remove a socket and evict it from every room it had joined.
    pub fn unregister(&self, id: ConnId) {
        if let Some((_, entry)) = self.conns.remove(&id) {
            metrics::gauge!("ws_connections").decrement(1.0);
            for room in entry.member_rooms.iter() {
                let key = room.key();
                if let Some(set) = self.rooms.get(key) {
                    set.remove(&id);
                    // Prune the room key when empty to prevent unbounded growth.
                    // Drop the Ref before acquiring the write lock inside remove_if.
                    if set.is_empty() {
                        drop(set);
                        self.rooms.remove_if(key, |_, s| s.is_empty());
                    }
                }
            }
        }
    }

    /// Number of live connections. O(1).
    #[inline]
    pub fn connection_count(&self) -> usize {
        self.conns.len()
    }

    /// Add a connection to a room (Socket.IO-style logical channel).
    ///
    /// ## TOCTOU safety
    ///
    /// Both `member_rooms` and `rooms` are updated while holding the DashMap
    /// shard read-lock on `conns` (the `Ref` keeps the shard locked until it is
    /// dropped). This prevents a concurrent `unregister()` — which needs the
    /// shard *write*-lock — from removing the entry between the two writes,
    /// which would otherwise leave a stale `ConnId` in `rooms` forever.
    pub fn join_room(&self, id: ConnId, room: String) {
        if let Some(entry) = self.conns.get(&id) {
            // member_rooms updated while shard read-lock is held.
            entry.member_rooms.insert(room.clone());
            // rooms secondary index updated while same lock is still live.
            self.rooms.entry(room).or_default().insert(id);
            // Ref dropped here → shard lock released.
        }
        // If the conn was not found the socket is already disconnected; no-op.
    }

    /// Remove a connection from a room.
    ///
    /// Mirrors `join_room`'s lock order: hold the `conns` shard read-lock first
    /// so a concurrent `unregister()` (which needs the write-lock) cannot remove
    /// the entry between the two writes and leave a stale id in `rooms`.
    pub fn leave_room(&self, id: ConnId, room: &str) {
        if let Some(entry) = self.conns.get(&id) {
            entry.member_rooms.remove(room);
            if let Some(set) = self.rooms.get(room) {
                set.remove(&id);
                if set.is_empty() {
                    drop(set);
                    self.rooms.remove_if(room, |_, s| s.is_empty());
                }
            }
        }
    }

    /// Enqueue a text frame to one connection. Non-blocking; a full queue
    /// evicts the slow client, a dead receiver is dropped silently.
    #[inline]
    pub fn send_text(&self, id: ConnId, text: Arc<str>) {
        let evict = match self.conns.get(&id) {
            Some(entry) => !self.enqueue_or_evict(id, &entry, WsMessage::Text(text)),
            None => false,
        }; // Ref dropped before any removal — same-shard deadlock safety.
        if evict {
            self.unregister(id);
        }
    }

    /// Fan a text frame out to every live connection. Creates one `Arc<str>`
    /// allocation then enqueues N pointer copies — no per-connection memcpy.
    /// Clients whose queues are full are evicted **after** the sweep
    /// (removal inside iteration would deadlock the DashMap shard).
    pub fn broadcast_text(&self, text: &str) {
        let arc: Arc<str> = Arc::from(text);
        let mut evict: Vec<ConnId> = Vec::new();
        for entry in self.conns.iter() {
            if !self.enqueue_or_evict(
                *entry.key(),
                entry.value(),
                WsMessage::Text(Arc::clone(&arc)),
            ) {
                evict.push(*entry.key());
            }
        }
        for id in evict {
            self.unregister(id);
        }
    }

    /// Fan a text frame out to the members of one room.
    pub fn broadcast_room_text(&self, room: &str, text: &str) {
        let Some(members) = self.rooms.get(room) else {
            return;
        };
        let arc: Arc<str> = Arc::from(text);
        let mut evict: Vec<ConnId> = Vec::new();
        for id in members.iter() {
            if let Some(entry) = self.conns.get(&id) {
                if !self.enqueue_or_evict(*id, &entry, WsMessage::Text(Arc::clone(&arc))) {
                    evict.push(*id);
                }
            }
        }
        drop(members); // release the rooms shard before unregister mutates it
        for id in evict {
            self.unregister(id);
        }
    }
}

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

// ─── Event envelope ───────────────────────────────────────────────────────

/// Serialise an `{ "event": <name>, "data": <payload> }` envelope. This is the
/// wire format for both the multiplexed WebSocket protocol and room broadcasts.
fn envelope<T: Serialize>(event: &str, payload: &T) -> String {
    #[derive(Serialize)]
    struct Envelope<'a, P> {
        event: &'a str,
        data: &'a P,
    }
    serde_json::to_string(&Envelope {
        event,
        data: payload,
    })
    .unwrap_or_else(|_| String::from(r#"{"event":"error","data":null}"#))
}

// ─── WsClient: the developer-facing connection handle ───────────────────────

/// A cheap, clonable handle to one connected client.
///
/// Passed to every gateway lifecycle hook and `#[Subscribe]` handler. All
/// methods are non-blocking and lock-free on the hot path: emitting just
/// enqueues onto a sharded channel. Holds a `&'static ConnectionRegistry`, so
/// cloning is a pointer copy — no `Arc` refcount traffic for the registry.
#[derive(Clone)]
pub struct WsClient {
    id: ConnId,
    reg: &'static ConnectionRegistry,
    claims: Option<Arc<Claims>>,
    tenant: Option<Arc<crate::web::tenant::TenantConfig>>,
}

impl WsClient {
    #[doc(hidden)]
    pub fn __new(
        id: ConnId,
        reg: &'static ConnectionRegistry,
        claims: Option<Arc<Claims>>,
        tenant: Option<Arc<crate::web::tenant::TenantConfig>>,
    ) -> Self {
        Self {
            id,
            reg,
            claims,
            tenant,
        }
    }

    /// This connection's unique id.
    #[inline]
    pub fn id(&self) -> ConnId {
        self.id
    }

    /// Session claims captured during the handshake (e.g. decoded JWT), if any.
    #[inline]
    pub fn claims(&self) -> Option<&Claims> {
        self.claims.as_deref()
    }

    /// The tenant resolved during the handshake — same registry, strategy,
    /// and suspension semantics as HTTP requests. Gateways enforce
    /// multi-tenancy with this exactly like controllers use `ctx.tenant()`.
    #[inline]
    pub fn tenant(&self) -> Option<&crate::web::tenant::TenantConfig> {
        self.tenant.as_deref()
    }

    /// Send an event to **this** client only.
    pub async fn emit<T: Serialize>(&self, event: &str, payload: T) {
        let text: Arc<str> = envelope(event, &payload).into();
        self.reg.send_text(self.id, text);
    }

    /// Broadcast an event to **every** connected client (including self).
    pub async fn broadcast<T: Serialize>(&self, event: &str, payload: T) {
        self.reg.broadcast_text(&envelope(event, &payload));
    }

    /// Broadcast an event to every client currently in `room`.
    pub async fn broadcast_to_room<T: Serialize>(&self, room: &str, event: &str, payload: T) {
        self.reg
            .broadcast_room_text(room, &envelope(event, &payload));
    }

    /// Join a logical room/channel. Subsequent room broadcasts reach this client.
    pub fn join_room(&self, room: impl Into<String>) {
        self.reg.join_room(self.id, room.into());
    }

    /// Leave a room.
    pub fn leave_room(&self, room: &str) {
        self.reg.leave_room(self.id, room);
    }
}

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

    #[cfg(test)]
    impl ConnectionRegistry {
        /// Test helper: backdate a connection's last activity.
        fn set_last_seen(&self, id: ConnId, unix_secs: u64) {
            if let Some(entry) = self.conns.get(&id) {
                entry.last_seen.store(unix_secs, Ordering::Relaxed);
            }
        }
    }

    fn reg_with_conn(buffer: usize) -> (ConnectionRegistry, ConnId, mpsc::Receiver<WsMessage>) {
        let reg = ConnectionRegistry::new();
        let (tx, rx) = mpsc::channel(buffer);
        let id = reg.register(tx, None);
        (reg, id, rx)
    }

    #[tokio::test]
    async fn slow_client_is_evicted_when_queue_fills() {
        let (reg, _id, _rx) = reg_with_conn(1);
        assert_eq!(reg.connection_count(), 1);

        // First frame fills the (never-drained) queue…
        reg.broadcast_text("one");
        assert_eq!(reg.connection_count(), 1);
        // …the second finds it full → the slow client is evicted.
        reg.broadcast_text("two");
        assert_eq!(reg.connection_count(), 0, "slow client must be evicted");
    }

    #[tokio::test]
    async fn fast_client_receives_broadcasts() {
        let (reg, _id, mut rx) = reg_with_conn(8);
        reg.broadcast_text("hello");
        match rx.recv().await {
            Some(WsMessage::Text(t)) => assert_eq!(&*t, "hello"),
            other => panic!("expected text frame, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn room_membership_and_cleanup() {
        let (reg, id, mut rx) = reg_with_conn(8);
        reg.join_room(id, "alpha".into());

        reg.broadcast_room_text("alpha", "in-room");
        assert!(matches!(rx.recv().await, Some(WsMessage::Text(t)) if &*t == "in-room"));

        reg.leave_room(id, "alpha");
        reg.broadcast_room_text("alpha", "after-leave");
        assert!(
            rx.try_recv().is_err(),
            "must not receive after leaving the room"
        );

        // Unregister evicts from rooms and prunes empty room keys.
        reg.join_room(id, "beta".into());
        reg.unregister(id);
        assert_eq!(reg.connection_count(), 0);
        reg.broadcast_room_text("beta", "to-nobody"); // must not panic
    }

    #[tokio::test]
    async fn sweep_idle_reaps_only_stale_connections() {
        let (reg, stale, _rx1) = reg_with_conn(8);
        let (tx2, mut rx2) = mpsc::channel(8);
        let fresh = reg.register(tx2, None);

        reg.set_last_seen(stale, unix_now() - 3600);
        let reaped = reg.sweep_idle(60);

        assert_eq!(reaped, vec![stale]);
        assert_eq!(reg.connection_count(), 1);
        // The fresh connection is untouched and still receives traffic.
        reg.send_text(fresh, Arc::from("still-alive"));
        assert!(matches!(rx2.recv().await, Some(WsMessage::Text(t)) if &*t == "still-alive"));
    }

    #[tokio::test]
    async fn close_all_enqueues_close_frames() {
        let (reg, _id, mut rx) = reg_with_conn(8);
        assert_eq!(reg.close_all(), 1);
        assert!(matches!(rx.recv().await, Some(WsMessage::Close)));
    }
}