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 an
//! `mpsc::UnboundedSender`; sending is a wait-free enqueue (`try`-style) that
//! never blocks the websocket frame-processing task. The per-socket writer task
//! drains its own receiver — so a slow client back-pressures only itself.
//!
//! ## 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>),
    /// Graceful close.
    Close,
}

/// Per-connection registry row. `tx` is the head of that socket's outbound
/// queue; `member_rooms` tracks which rooms to evict the id from on disconnect.
struct ConnEntry {
    tx: mpsc::UnboundedSender<WsMessage>,
    member_rooms: DashSet<String>,
    #[allow(dead_code)]
    claims: Option<Arc<Claims>>,
}

/// 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`].
    pub fn register(
        &self,
        tx: mpsc::UnboundedSender<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(),
                claims,
            },
        );
        id
    }

    /// 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() {
            let _ = entry.value().tx.send(WsMessage::Close);
            n += 1;
        }
        n
    }

    /// 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) {
            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 dead receiver
    /// is silently dropped (the reader loop will unregister it shortly).
    #[inline]
    pub fn send_text(&self, id: ConnId, text: Arc<str>) {
        if let Some(entry) = self.conns.get(&id) {
            let _ = entry.tx.send(WsMessage::Text(text));
        }
    }

    /// Fan a text frame out to every live connection. Creates one `Arc<str>`
    /// allocation then enqueues N pointer copies — no per-connection memcpy.
    pub fn broadcast_text(&self, text: &str) {
        let arc: Arc<str> = Arc::from(text);
        for entry in self.conns.iter() {
            let _ = entry.value().tx.send(WsMessage::Text(Arc::clone(&arc)));
        }
    }

    /// 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);
        for id in members.iter() {
            if let Some(entry) = self.conns.get(&id) {
                let _ = entry.tx.send(WsMessage::Text(Arc::clone(&arc)));
            }
        }
    }
}

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>>,
}

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

    /// 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()
    }

    /// 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);
    }
}