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

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

    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",
                ));
            }
            for client in clients.values() {
                let _ = client.tx.send(json.clone());
            }
        }

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