quicproquo/crates/quicprochat-rpc/src/push.rs

//! Server-push infrastructure — manages push connections per identity key.

use bytes::Bytes;
use dashmap::DashMap;
use quinn::Connection;
use tracing::{debug, warn};

use crate::error::RpcError;

/// Manages push connections per identity key.
pub struct PushBroker {
    /// Map from identity_key to active QUIC connections.
    connections: DashMap<Vec<u8>, Vec<Connection>>,
}

impl PushBroker {
    pub fn new() -> Self {
        Self {
            connections: DashMap::new(),
        }
    }

    /// Register a connection for an identity.
    pub fn register(&self, identity_key: Vec<u8>, connection: Connection) {
        self.connections
            .entry(identity_key)
            .or_default()
            .push(connection);
    }

    /// Remove closed connections from the registry.
    pub fn gc(&self) {
        self.connections.alter_all(|_, mut conns| {
            conns.retain(|c| c.close_reason().is_none());
            conns
        });
        self.connections.retain(|_, conns| !conns.is_empty());
    }

    /// Send a push event to all connections for an identity.
    /// Returns the number of successful sends.
    pub async fn send_to(
        &self,
        identity_key: &[u8],
        event_type: u16,
        payload: Bytes,
    ) -> usize {
        let conns = match self.connections.get(identity_key) {
            Some(entry) => entry.clone(),
            None => return 0,
        };

        let mut sent = 0usize;
        for conn in &conns {
            match crate::server::send_push(conn, event_type, payload.clone()).await {
                Ok(()) => {
                    sent += 1;
                    debug!("push sent to connection {}", conn.remote_address());
                }
                Err(RpcError::Connection(e)) => {
                    warn!("push send failed to {}: {e}", conn.remote_address());
                }
                Err(e) => {
                    warn!("push send error: {e}");
                }
            }
        }
        sent
    }

    /// Send a push event to all members of a channel.
    /// `member_keys` is the list of identity keys in the channel.
    /// Returns the total number of successful sends.
    pub async fn send_to_channel(
        &self,
        member_keys: &[Vec<u8>],
        event_type: u16,
        payload: Bytes,
    ) -> usize {
        let mut total = 0usize;
        for key in member_keys {
            total += self.send_to(key, event_type, payload.clone()).await;
        }
        total
    }

    /// Number of identities with registered connections.
    pub fn identity_count(&self) -> usize {
        self.connections.len()
    }
}

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

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

    #[test]
    fn new_broker_is_empty() {
        let broker = PushBroker::new();
        assert_eq!(broker.identity_count(), 0);
    }

    #[test]
    fn default_broker_is_empty() {
        let broker = PushBroker::default();
        assert_eq!(broker.identity_count(), 0);
    }
}