infinity_bridge_host/

hub.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use infinity_bridge_wire::{
    AckPayload, BridgeError, CmdPayload, EventPayload, HelloPayload, WireMsg,
};
use serde_json::Value;
use tokio::sync::{Mutex, broadcast, mpsc, oneshot, watch};
use uuid::Uuid;

use crate::client::Client;

#[derive(Debug, Clone)]
pub struct ClientInfo {
    pub id: u64,
    pub hello: Option<HelloPayload>,
}

pub(crate) struct Hub {
    next_id: AtomicU64,
    clients: Mutex<HashMap<u64, Client>>,
    pending: Mutex<HashMap<String, oneshot::Sender<AckPayload>>>,

    event_tx: broadcast::Sender<EventPayload>,

    connection_tx: watch::Sender<bool>,
    connection_rx: watch::Receiver<bool>,

    connect_notify: tokio::sync::Notify,
}

impl Hub {
    pub fn new(event_capacity: usize) -> Arc<Self> {
        let (event_tx, _) = broadcast::channel(event_capacity);
        let (connection_tx, connection_rx) = watch::channel(false);

        Arc::new(Self {
            next_id: AtomicU64::new(1),
            clients: Mutex::new(HashMap::new()),
            pending: Mutex::new(HashMap::new()),
            event_tx,
            connection_tx,
            connection_rx,
            connect_notify: tokio::sync::Notify::new(),
        })
    }

    pub async fn register_client(&self, tx: mpsc::UnboundedSender<String>) -> u64 {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let mut clients = self.clients.lock().await;
        clients.insert(
            id,
            Client {
                tx,
                hello: None,
                last_seen: tokio::time::Instant::now(),
                ready: false,
            },
        );
        self.connection_tx.send_replace(true);
        self.connect_notify.notify_waiters();
        id
    }

    pub async fn set_client_hello(&self, id: u64, hello: HelloPayload) {
        let mut clients = self.clients.lock().await;
        if let Some(c) = clients.get_mut(&id) {
            c.hello = Some(hello);
        }
    }

    pub async fn unregister_client(&self, id: u64) {
        let mut clients = self.clients.lock().await;
        clients.remove(&id);
        let connected = !clients.is_empty();
        self.connection_tx.send_replace(connected);
    }

    pub async fn touch_client(&self, id: u64) {
        let mut clients = self.clients.lock().await;
        if let Some(c) = clients.get_mut(&id) {
            c.last_seen = tokio::time::Instant::now();
        }
    }

    pub async fn reap_dead_clients(&self, timeout: Duration) -> Vec<u64> {
        let now = tokio::time::Instant::now();
        let mut clients = self.clients.lock().await;
        let dead: Vec<u64> = clients
            .iter()
            .filter(|(_, c)| now.duration_since(c.last_seen) > timeout)
            .map(|(&id, _)| id)
            .collect();

        for &id in &dead {
            clients.remove(&id);
        }

        if !dead.is_empty() {
            let connected = !clients.is_empty();
            self.connection_tx.send_replace(connected);
        }

        dead
    }

    /// Record a client's downstream readiness (see
    /// [`infinity_bridge_wire::READY_EVENT`]).
    pub async fn set_client_ready(&self, id: u64, ready: bool) {
        let mut clients = self.clients.lock().await;
        if let Some(c) = clients.get_mut(&id) {
            c.ready = ready;
        }
    }

    pub async fn is_connected(&self) -> bool {
        !self.clients.lock().await.is_empty()
    }

    /// `true` when at least one client has reported its downstream link bound.
    ///
    /// Distinct from [`Self::is_connected`]: a relay gauge can hold an open
    /// socket for the whole session while the module it fronts is still
    /// loading, absent, or has lost its IPC binding.
    pub async fn is_ready(&self) -> bool {
        self.clients.lock().await.values().any(|c| c.ready)
    }

    pub async fn wait_connected(&self) {
        loop {
            if self.is_connected().await {
                return;
            }
            self.connect_notify.notified().await;
        }
    }

    pub fn subscribe_connection_status(&self) -> watch::Receiver<bool> {
        self.connection_rx.clone()
    }

    pub async fn connected_clients(&self) -> Vec<ClientInfo> {
        let clients = self.clients.lock().await;
        clients
            .iter()
            .map(|(&id, c)| ClientInfo {
                id,
                hello: c.hello.clone(),
            })
            .collect()
    }

    pub fn subscribe_events(&self) -> broadcast::Receiver<EventPayload> {
        self.event_tx.subscribe()
    }

    pub fn dispatch_event(&self, event: EventPayload) {
        let _ = self.event_tx.send(event);
    }

    pub async fn emit(&self, name: impl Into<String>, data: Value) -> Result<(), BridgeError> {
        let msg = WireMsg::Event(EventPayload::new(name, data));
        let json = msg.to_json()?;

        let clients = self.clients.lock().await;
        if clients.is_empty() {
            return Err(BridgeError::no_clients(
                "no gauges connected — event dropped",
            ));
        }

        let mut send_failures = 0u32;
        for client in clients.values() {
            if client.tx.send(json.clone()).is_err() {
                send_failures += 1;
            }
        }

        if send_failures > 0 && send_failures as usize == clients.len() {
            return Err(BridgeError::transport(
                "all gauge connections failed to accept event",
            ));
        }

        Ok(())
    }

    pub async fn command(
        &self,
        name: Option<&str>,
        payload: Value,
        timeout: Duration,
    ) -> Result<Value, BridgeError> {
        let id = Uuid::new_v4().to_string();

        let cmd = match name {
            Some(n) => CmdPayload::named(id.clone(), n, payload),
            None => CmdPayload::new(id.clone(), payload),
        };
        let msg = WireMsg::Cmd(cmd);
        let json = msg.to_json()?;

        let (ack_tx, ack_rx) = oneshot::channel();
        self.pending.lock().await.insert(id.clone(), ack_tx);

        {
            let clients = self.clients.lock().await;
            if clients.is_empty() {
                self.pending.lock().await.remove(&id);
                return Err(BridgeError::no_clients(
                    "no gauges connected — cannot send command",
                ));
            }

            // Prefer clients that have reported their downstream link bound.
            // Broadcasting to a relay whose CommBus isn't attached burns the
            // full timeout on a socket that was never going to answer. When
            // nothing has ever reported ready the flag carries no information
            // (older relays don't send it), so fall back to every client.
            let any_ready = clients.values().any(|c| c.ready);
            let mut delivered = 0usize;
            for client in clients.values() {
                if any_ready && !client.ready {
                    continue;
                }
                if client.tx.send(json.clone()).is_ok() {
                    delivered += 1;
                }
            }

            // Every target's writer half is gone: the sockets are dead but not
            // yet unregistered. Fail now instead of waiting out the timeout —
            // the caller can retry into a fresh connection immediately.
            if delivered == 0 {
                self.pending.lock().await.remove(&id);
                return Err(BridgeError::transport(
                    "no gauge connection accepted the command",
                ));
            }
        }

        let ack_result = tokio::time::timeout(timeout, ack_rx).await;

        self.pending.lock().await.remove(&id);

        match ack_result {
            Ok(Ok(ack)) => {
                if ack.ok {
                    Ok(ack.response.unwrap_or(Value::Null))
                } else {
                    Err(BridgeError::application(
                        ack.error.unwrap_or_else(|| "unknown error".into()),
                    ))
                }
            }
            Ok(Err(_)) => Err(BridgeError::transport(
                "all gauge connections dropped before ack",
            )),
            Err(_) => Err(BridgeError::timeout(format!(
                "no ack received within {timeout:?}"
            ))),
        }
    }

    pub async fn dispatch_ack(&self, ack: AckPayload) {
        let tx = self.pending.lock().await.remove(&ack.id);
        if let Some(tx) = tx {
            let _ = tx.send(ack);
        }
    }

    pub async fn send_to(&self, client_id: u64, json: String) -> Result<(), BridgeError> {
        let clients = self.clients.lock().await;
        let client = clients
            .get(&client_id)
            .ok_or_else(|| BridgeError::transport(format!("client {client_id} not found")))?;
        client
            .tx
            .send(json)
            .map_err(|_| BridgeError::transport(format!("client {client_id} channel closed")))
    }

    pub async fn broadcast(&self, json: String) -> Result<(), BridgeError> {
        let clients = self.clients.lock().await;
        if clients.is_empty() {
            return Err(BridgeError::no_clients("no gauges connected"));
        }
        for client in clients.values() {
            let _ = client.tx.send(json.clone());
        }
        Ok(())
    }
}

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

    /// Long enough to prove nothing was delivered, short enough to keep the
    /// suite quick — no ack is ever sent in these tests.
    const NO_ACK: Duration = Duration::from_millis(50);

    #[tokio::test]
    async fn readiness_is_tracked_per_client() {
        let hub = Hub::new(16);
        let (tx_a, _rx_a) = mpsc::unbounded_channel();
        let (tx_b, _rx_b) = mpsc::unbounded_channel();
        let a = hub.register_client(tx_a).await;
        let _b = hub.register_client(tx_b).await;

        assert!(hub.is_connected().await);
        assert!(!hub.is_ready().await, "a fresh client is not yet ready");

        hub.set_client_ready(a, true).await;
        assert!(hub.is_ready().await);

        hub.set_client_ready(a, false).await;
        assert!(!hub.is_ready().await);
        assert!(hub.is_connected().await, "readiness is not connectedness");
    }

    #[tokio::test]
    async fn a_command_skips_clients_that_are_not_ready() {
        let hub = Hub::new(16);
        let (tx_stale, mut rx_stale) = mpsc::unbounded_channel();
        let (tx_live, mut rx_live) = mpsc::unbounded_channel();
        let _stale = hub.register_client(tx_stale).await;
        let live = hub.register_client(tx_live).await;
        hub.set_client_ready(live, true).await;

        let err = hub
            .command(Some("ping"), Value::Null, NO_ACK)
            .await
            .expect_err("nothing acks in this test");
        assert_eq!(err.kind(), ErrorKind::Timeout);

        assert!(
            rx_live.try_recv().is_ok(),
            "the ready client got the command"
        );
        assert!(
            rx_stale.try_recv().is_err(),
            "a relay that never reported ready must not absorb the timeout budget"
        );
    }

    #[tokio::test]
    async fn a_command_goes_everywhere_while_readiness_is_unknown() {
        // Relays predating the ready event never send one. Their silence must
        // read as "unknown", not "unreachable", or they stop working entirely.
        let hub = Hub::new(16);
        let (tx_a, mut rx_a) = mpsc::unbounded_channel();
        let (tx_b, mut rx_b) = mpsc::unbounded_channel();
        hub.register_client(tx_a).await;
        hub.register_client(tx_b).await;

        let _ = hub.command(Some("ping"), Value::Null, NO_ACK).await;

        assert!(rx_a.try_recv().is_ok());
        assert!(rx_b.try_recv().is_ok());
    }

    #[tokio::test]
    async fn a_command_fails_fast_when_every_connection_is_gone() {
        // Sockets can be dead before the reaper has unregistered them. Waiting
        // out the full timeout for a send that provably went nowhere is time
        // the caller could have spent retrying into a fresh connection.
        let hub = Hub::new(16);
        let (tx, rx) = mpsc::unbounded_channel();
        hub.register_client(tx).await;
        drop(rx);

        let start = tokio::time::Instant::now();
        let err = hub
            .command(Some("ping"), Value::Null, Duration::from_secs(30))
            .await
            .expect_err("the only receiver is gone");
        assert_eq!(err.kind(), ErrorKind::Transport);
        assert!(
            start.elapsed() < Duration::from_secs(1),
            "should not have waited out the timeout"
        );
    }

    #[tokio::test]
    async fn a_timed_out_command_leaves_no_pending_entry() {
        let hub = Hub::new(16);
        let (tx, _rx) = mpsc::unbounded_channel();
        hub.register_client(tx).await;

        let _ = hub.command(Some("ping"), Value::Null, NO_ACK).await;
        assert!(hub.pending.lock().await.is_empty());
    }

    #[tokio::test(start_paused = true)]
    async fn a_silent_client_is_reaped_and_named() {
        let hub = Hub::new(16);
        let (tx, _rx) = mpsc::unbounded_channel();
        let id = hub.register_client(tx).await;

        tokio::time::advance(Duration::from_secs(31)).await;
        let dead = hub.reap_dead_clients(Duration::from_secs(30)).await;

        assert_eq!(dead, vec![id], "the reaper names who it dropped");
        assert!(!hub.is_connected().await);
    }
}