hydra_sync/

server.rs

//! Core implementation for the `HydraSync` server.
use crate::channel::{ClosedError, FullError};
use crate::client::ServerMetrics;
use crate::crypto::{NONCE_LEN, TAG_LEN};
use crate::protocol::StatusCode;
use crate::protocol::{
    Role, perform_server_handshake, read_join_header, write_server_metrics,
    write_server_read_write_len, write_status_code,
};
use crate::session::SessionMap;
use crate::{ChannelOverflowStrategy, START_TIME, error, get_uptime_hrs, info, trace, warn};
use anyhow::Result;
use bytes::BytesMut;
use hex::encode;
use std::net::SocketAddr;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::Notify;

/// Padded atomic64 to avoid false sharing on hot counters (64B cache line).
#[repr(C, align(64))]
struct PaddedAtomicU64(AtomicU64);
impl PaddedAtomicU64 {
    const fn pad(v: u64) -> Self {
        Self(AtomicU64::new(v))
    }
}
impl Deref for PaddedAtomicU64 {
    type Target = AtomicU64;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
impl DerefMut for PaddedAtomicU64 {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

/// `HydraServer`- A light-weight, E2E `(AES-GCM)`, multi-threaded SPMC broadcast server using `TCP`.
/// It implements internal _ring buffer_ configured by [`HydraConfig`], uses **little memory** and **minimal copy** as possible.
///
/// ```no_run
/// use hydra_sync::server::HydraServer;
///
/// #[tokio::main]
/// async fn main() {
///     let (server, addr) = HydraServer::bind_default().await.unwrap();
///     println!("Server running on: {}", addr);
///     let server_runner = server.clone();
///     tokio::spawn(async move { server_runner.run().await }); // run in background
///
///     server.shutdown().await.unwrap(); // shutdown gracefully
/// }
/// ```
pub struct HydraServer {
    listener: TcpListener,
    config: Arc<HydraConfig>,
    /*
    field below are used for collecting server metrics
    */
    /// Global session concurrent map.
    session_map: Arc<SessionMap>,
    /// Track the number of active connections to the server.
    active_connections: Arc<AtomicU64>,
    /// Total number of sessions ever created.
    total_sessions: Arc<PaddedAtomicU64>,
    /// Total number of bytes relayed through the server.
    total_bytes: Arc<PaddedAtomicU64>,
    /// Flag to indicate if the server is shutting down gracefully
    /// (server will stop accepting new connections and wait for existing sessions to finish).
    shutdown_flag: Arc<AtomicBool>,
}

/// Configuration for the `HydraServer` features.
pub struct HydraConfig {
    /// Maximum number of concurrent connections allowed.
    pub max_connections: u64,
    /// Timeout in milliseconds for accepting new connections when capacity is reached.
    pub accept_timeout: u64,
    /// Fixed length of a single packet/frame in bytes to read/write `(excluding AES-GCM overhead)`.
    /// This length is used to `read/write EXACT` data from/to the network stream.
    pub fixed_packet_length: u64,
    /// Capacity of the internal channel for each session.
    pub channel_capacity: u64,
    /// Token for authenticating observer clients for getting [`ServerMetrics`] from server ring's buffer.
    pub observer_token: Option<[u8; 64]>,
    /// Strategy for handling channel overflow when consumers cannot keep up with producers.
    pub overflow_mode: ChannelOverflowStrategy,
}

impl HydraConfig {
    /// Creates `HydraConfig` with the specified parameters.
    /// > NOTE: `fixed_packet_length` is the length of the plaintext data, the actual read/write length will be `fixed_packet_length + AES-GCM overhead (28 bytes)`.
    /// > So this function adds the AES-GCM overhead to the `fixed_packet_length` (make sure to consider for it, if `HydraConfig` is constructed without calling this function).
    pub fn config(
        max_connections: u64,
        accept_timeout: u64,
        fixed_packet_length: u64,
        channel_capacity: u64,
        observer_token: Option<[u8; 64]>,
        overflow_mode: ChannelOverflowStrategy,
    ) -> Self {
        if fixed_packet_length == 0 {
            panic!("fixed_packet_length must be greater than 0");
        }
        Self {
            max_connections,
            accept_timeout,
            fixed_packet_length: fixed_packet_length + (NONCE_LEN + TAG_LEN) as u64, // add aes-gcm overhead
            channel_capacity,
            observer_token,
            overflow_mode,
        }
    }
}

impl Default for HydraConfig {
    /// Default HydraConfig
    /// - max_connections: 32
    /// - accept_timeout: 500 ms
    /// - fixed_packet_length: 4 MB + `AES-GCM overhead (28 bytes)`
    /// - channel_capacity: 256
    /// - observer_token: randomly generated
    /// - overflow_mode: DropPacket
    fn default() -> Self {
        Self {
            max_connections: 32,
            accept_timeout: 500,
            fixed_packet_length: (4 << 20) + (NONCE_LEN + TAG_LEN) as u64, // 4 MB + AES-GCM overhead
            channel_capacity: 256,
            observer_token: Some(rand::random::<[u8; 64]>()),
            overflow_mode: ChannelOverflowStrategy::DropPacket,
        }
    }
}

impl HydraServer {
    /// Binds the server with random local port and default [`HydraConfig`], returns the port.
    pub async fn bind_default() -> Result<(Arc<Self>, SocketAddr)> {
        let addr = "127.0.0.1:0".parse::<SocketAddr>()?;
        let (server, server_addr) = HydraServer::bind(addr, HydraConfig::default()).await?;
        Ok((server, server_addr))
    }

    /// Binds the server to the specified [`SocketAddr`] with the provided [`HydraConfig`].
    pub async fn bind(
        socket_addr: SocketAddr,
        hydra_config: HydraConfig,
    ) -> Result<(Arc<Self>, SocketAddr)> {
        let server_tcp = TcpListener::bind(socket_addr).await?;
        let server_addr = server_tcp.local_addr()?;
        Ok((
            Arc::new(Self {
                listener: server_tcp,
                config: Arc::new(hydra_config),
                session_map: Arc::new(SessionMap::init()),
                active_connections: Arc::new(AtomicU64::new(0)),
                total_sessions: Arc::new(PaddedAtomicU64::pad(0)),
                total_bytes: Arc::new(PaddedAtomicU64::pad(0)),
                shutdown_flag: Arc::new(AtomicBool::new(false)),
            }),
            server_addr,
        ))
    }

    /// Starts the server, accepting incoming connections and handling them concurrently.
    /// Overview:
    /// * Accept incoming `TCP connections`, spawn handler tasks and `allocate read/write buffers` or reject if max connections reached.
    /// * Perform handshake, read/decrypt `JoinHeader`, determine role (Producer/Consumer/Admin).
    /// * For Producers: Create valid session, read data (exactly **fixed_packet_length**, including `AES-GCM` (28 bytes) overhead, with `read_exact()`, so it may truncate)
    ///   and broadcast `(cheap-copy)` to all Consumers in the session (Write behavior is configured by [`ChannelOverflowStrategy`]).
    /// * For Consumers: Validate session, read from `local ring buffer` from producer handle, write to TCP stream.
    /// * `LOG_LEVEL` & `LOG_FILE_DIR` env vars can be set to control `logging verbosity` and output file (defaults to `info` level and stdout and file).
    /// * EOF check are gracefully handled by closing the connection without logging an error.
    pub async fn run(&self) -> Result<()> {
        START_TIME.set(chrono::Local::now()).ok();

        loop {
            // check shutdown token
            if self.shutdown_flag.load(Ordering::Relaxed) {
                // return if no active connections
                if self.active_connections.load(Ordering::Relaxed) == 0 {
                    return Ok(());
                }
                loop {
                    // check current connections, every 3 seconds, wait for them to finish and return
                    tokio::time::sleep(Duration::from_secs(3)).await;
                    if self.active_connections.load(Ordering::Relaxed) > 0 {
                        info!(
                            "Server is shutting down, waiting for {} active connections to finish...",
                            self.active_connections.load(Ordering::Relaxed)
                        );
                    } else {
                        info!("Server has stopped.");
                        return Ok(());
                    }
                }
            }

            if self.active_connections.fetch_add(1, Ordering::Relaxed)
                >= self.config.max_connections
            {
                self.active_connections.fetch_sub(1, Ordering::Relaxed);
                warn!(
                    "Max connections reached: {}, waiting {} ms before accepting new connections",
                    self.config.max_connections, self.config.accept_timeout
                );
                tokio::time::sleep(Duration::from_millis(self.config.accept_timeout)).await;
                continue;
            }

            tokio::select! {
                // accept new connection
                result = self.listener.accept() => {
                    match result {
                        Ok((stream, peer_addr)) => {
                            stream.set_nodelay(true).ok();
                            let config = Arc::clone(&self.config);
                            let session_map = Arc::clone(&self.session_map);
                            let connections = Arc::clone(&self.active_connections);
                            let total_sessions = Arc::clone(&self.total_sessions);
                            let total_bytes = Arc::clone(&self.total_bytes);
                            let is_shutting_down = Arc::clone(&self.shutdown_flag);
                            // spawn client_handler task
                            tokio::spawn(async move {
                                trace!("Accepted connection from: {}", peer_addr);
                                if let Err(e) = Self::handle_connection(
                                    stream,
                                    session_map,
                                    config,
                                    total_sessions,
                                    total_bytes,
                                    is_shutting_down,
                                )
                                .await
                                {
                                    error!("Connection handling error: {} from: {}", e, peer_addr);
                                }
                                connections.fetch_sub(1, Ordering::Release);
                            });
                        }
                        Err(e) => {
                            self.active_connections.fetch_sub(1, Ordering::Release);
                            error!("Connection accepting error: {}", e);
                        }
                    }
                }
                // shutdown signal, then break blocking accept()
                _ = async {
                    while !self.shutdown_flag.load(Ordering::Relaxed) {
                        tokio::time::sleep(Duration::from_millis(100)).await;
                    }
                } => {
                    self.active_connections.fetch_sub(1, Ordering::Release);
                    continue;
                }
            }
        }
    }

    /// Handles an individual client connection.
    async fn handle_connection(
        mut stream: TcpStream,
        sessions: Arc<SessionMap>,
        config: Arc<HydraConfig>,
        total_sessions: Arc<PaddedAtomicU64>,
        total_bytes: Arc<PaddedAtomicU64>,
        is_shutting_down: Arc<AtomicBool>,
    ) -> Result<()> {
        stream.set_nodelay(true).ok();
        let client_addr = stream.peer_addr()?;
        let (mut read_h, mut writer_h) = stream.split();
        let mut read_buffer = BytesMut::zeroed(config.fixed_packet_length as usize);
        let transport_key = perform_server_handshake(&mut read_h, &mut writer_h).await?;
        let join_header = read_join_header(&mut read_h, &transport_key, &mut read_buffer).await?;

        // determine role and handle accordingly
        let role = match Role::from_u8(join_header.role) {
            Ok(role) => {
                info!(
                    "Client (ip: {}) joined with role: {}, session_uuid: {}...",
                    client_addr,
                    role.to_string(),
                    &encode(join_header.uuid_or_token)[..12]
                );
                role
            }
            Err(_) => {
                warn!(
                    "Client (ip: {}) joining with invalid role: {:#04x}",
                    client_addr, join_header.role
                );
                write_status_code(
                    &mut writer_h,
                    StatusCode::ErrInvalidRole,
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;
                return Ok(());
            }
        };

        // handle roles
        match role {
            Role::Producer => {
                // create session for producer, if already exists, return error
                if !sessions.try_create_session(join_header.uuid_or_token) {
                    warn!(
                        "Producer (ip: {}) failed to create session_uuid: {}...",
                        client_addr,
                        &encode(join_header.uuid_or_token)[..12]
                    );

                    // send error code
                    write_status_code(
                        &mut writer_h,
                        StatusCode::ErrSessionAlreadyOccupied,
                        &transport_key,
                        &mut read_buffer,
                    )
                    .await?;
                    return Ok(());
                }

                info!(
                    "Producer (ip: {}) created session_uuid: {}...",
                    client_addr,
                    &encode(join_header.uuid_or_token)[..12]
                );

                total_sessions.fetch_add(1, Ordering::Relaxed);
                // on successful session creation, send read_len to producer
                write_status_code(
                    &mut writer_h,
                    StatusCode::Success,
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;
                write_server_read_write_len(
                    &mut writer_h,
                    config.fixed_packet_length - (NONCE_LEN + TAG_LEN) as u64, // announce raw payload capacity (encryption overhead is internal)
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;

                // read from producer and write to all consumers with fixed read_write len
                let mut write_buffer = BytesMut::zeroed(config.fixed_packet_length as usize);
                let mut dead_consumers = vec![]; // 'client to drop' collection
                let mut pending_consumers = vec![]; // 'slow client' collection

                loop {
                    // react_exact to fill buf, this does not extend buf anyway,
                    // this will return an error if the stream is closed or if it cannot fill the buffer exactly
                    // 'reads' may get truncated and may send corrupted data to consumers
                    match read_h.read_exact(&mut write_buffer).await {
                        Ok(_) => {
                            let data = write_buffer.split().freeze(); // zero-copy 'freeze' to Bytes

                            // get session ref (TODO; get read lock, maybe block other map shards? idk)
                            if let Some(session) =
                                sessions.try_get_session(join_header.uuid_or_token)
                            {
                                // non-blocking pass over every consumer ring,
                                // fast consumers are served immediately, regardless of where a
                                // slow one sits in the (arbitrary) DashMap iteration order
                                for mut entry in session.consumer_local_buffers.iter_mut() {
                                    // rm consumers that disconnected and their rings
                                    if entry.value_mut().is_closed() {
                                        dead_consumers.push(*entry.key());
                                        continue;
                                    }
                                    // broadcast away
                                    let data_len = data.len() as u64;
                                    match entry.value_mut().push(data.clone()) {
                                        Ok(()) => {
                                            total_bytes.fetch_add(data_len, Ordering::Relaxed);
                                        }
                                        Err(FullError) => match config.overflow_mode {
                                            ChannelOverflowStrategy::DropPacket => {
                                                // skip, don't push latest packet to local consumer queue
                                            }
                                            ChannelOverflowStrategy::DropClient => {
                                                // collect the 'lagged' client
                                                dead_consumers.push(*entry.key());
                                            }
                                            ChannelOverflowStrategy::BackPressure => {
                                                // collect all slow consumer
                                                pending_consumers.push(*entry.key());
                                            }
                                        },
                                    }
                                }

                                // blocking delivery pass for full rings only,
                                // parks until space frees up; woken by pop() or by the
                                // consumer disconnecting by ClosedError
                                for id in pending_consumers.drain(..) {
                                    loop {
                                        enum Step {
                                            Pushed,
                                            Remove,
                                            Wait(Arc<Notify>),
                                        }
                                        let step: Step = 'lock: {
                                            // get lock on the consumer ring, or remove the entry
                                            let Some(mut entry) =
                                                session.consumer_local_buffers.get_mut(&id)
                                            else {
                                                break 'lock Step::Remove;
                                            };
                                            // check if closed, or remove
                                            let p = entry.value_mut();
                                            if p.is_closed() {
                                                break 'lock Step::Remove;
                                            }
                                            match p.push(data.clone()) {
                                                Ok(()) => {
                                                    total_bytes.fetch_add(
                                                        data.len() as u64,
                                                        Ordering::Relaxed,
                                                    );
                                                    Step::Pushed
                                                }
                                                Err(FullError) => Step::Wait(p.space_notify()),
                                            }
                                        };
                                        match step {
                                            Step::Pushed => break,
                                            Step::Remove => {
                                                dead_consumers.push(id);
                                                break;
                                            }
                                            // no lock held here, the guard was dropped at 'lock
                                            Step::Wait(notify) => notify.notified().await,
                                        }
                                    }
                                }

                                // remove dead consumer queues
                                for id in dead_consumers.drain(..) {
                                    session.consumer_local_buffers.remove(&id);
                                }
                            }

                            write_buffer.resize(config.fixed_packet_length as usize, 0);
                        }
                        Err(_) => {
                            error!(
                                "Producer (ip: {}) read error, closing connection for session_uuid: {}...",
                                client_addr,
                                &encode(join_header.uuid_or_token)[..12]
                            );
                            break;
                        }
                    }
                }

                info!(
                    "Producer (ip: {}) disconnected, removing session_uuid: {}...",
                    client_addr,
                    &encode(join_header.uuid_or_token)[..12]
                );

                // remove session and return (close stream) on disconnect or error
                sessions.remove_session(join_header.uuid_or_token);
            }

            Role::Consumer => {
                // check session exists, if not, return error
                let session = match sessions.try_get_session(join_header.uuid_or_token) {
                    Some(s) => s,
                    None => {
                        warn!(
                            "Consumer (ip: {}) failed to join session_uuid: {}...",
                            client_addr,
                            &encode(join_header.uuid_or_token)[..12]
                        );

                        write_status_code(
                            &mut writer_h,
                            StatusCode::ErrSessionNotFound,
                            &transport_key,
                            &mut read_buffer,
                        )
                        .await?;
                        return Ok(());
                    }
                };

                info!(
                    "Consumer (ip: {}) joined session_uuid: {}...",
                    client_addr,
                    &encode(join_header.uuid_or_token)[..12]
                );

                // on success, send write_len to consumer
                write_status_code(
                    &mut writer_h,
                    StatusCode::Success,
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;
                write_server_read_write_len(
                    &mut writer_h,
                    config.fixed_packet_length - (NONCE_LEN + TAG_LEN) as u64, // announce raw payload capacity (encryption overhead is internal)
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;

                // add new consumer to session, this adds local ring buffer for this consumer, and returns a consumer handle
                let (consumer_id, mut consumer) =
                    session.add_consumer(config.channel_capacity as usize)?;

                let mut peek = [0u8; 1]; // for checking conn alive
                loop {
                    tokio::select! {
                        // block (async) until a packet arrives or peer drops;
                        // no busy waiting, woken exactly when the ring has data
                        popped = consumer.pop_async() => {
                            match popped {
                                Ok(data) => {
                                    if writer_h.write_all(&data).await.is_err() {
                                        let _ = writer_h.flush().await;
                                        let _ = writer_h.shutdown().await;
                                        error!(
                                            "Consumer (ip: {}) write error, closing connection for session: {}...",
                                            client_addr,
                                            &encode(join_header.uuid_or_token)[..12]
                                        );
                                        break // if write fails (broken pipe and such issues)
                                    }
                                }
                                Err(ClosedError) => {
                                    let _ = writer_h.flush().await;
                                    let _ = writer_h.shutdown().await;
                                    info!(
                                        "Consumer (ip: {}) session closed, closing connection for session: {}...",
                                        client_addr,
                                        &encode(join_header.uuid_or_token)[..12]
                                    );
                                    break // producer gone / session removed
                                },
                            }
                        }
                        result = read_h.read(&mut peek) => {
                            match result {
                                Ok(0) => break, // EOF check
                                Err(_) => {
                                    let _ = writer_h.flush().await;
                                    let _ = writer_h.shutdown().await;
                                    info!(
                                        "Consumer (ip: {}) peer read error, closing connection for session: {}...",
                                        client_addr,
                                        &encode(join_header.uuid_or_token)[..12]
                                    );
                                    break // if read fails or connection closed, break
                                },
                                _ => {}
                            }
                        }
                    }
                }

                info!(
                    "Consumer (ip: {}) disconnected, closing connection for session_uuid: {}...",
                    client_addr,
                    &encode(join_header.uuid_or_token)[..12]
                );

                // remove own queue from the session, its dead weight
                if let Some(session) = sessions.try_get_session(join_header.uuid_or_token) {
                    session.consumer_local_buffers.remove(&consumer_id);
                }
            }

            Role::Observer => {
                // check token, if not present or invalid, return error
                if let Some(token) = config.observer_token
                    && token != join_header.uuid_or_token
                {
                    warn!(
                        "Observer (ip: {}) failed to authenticate with token: {}...",
                        client_addr,
                        &encode(join_header.uuid_or_token)[..12]
                    );

                    write_status_code(
                        &mut writer_h,
                        StatusCode::ErrInvalidToken,
                        &transport_key,
                        &mut read_buffer,
                    )
                    .await?;
                    return Ok(());
                }

                info!("Observer (ip: {}) token approved", client_addr,);
                write_status_code(
                    &mut writer_h,
                    StatusCode::Success,
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;
                // send this anyway
                write_server_read_write_len(
                    &mut writer_h,
                    config.fixed_packet_length - (NONCE_LEN + TAG_LEN) as u64,
                    &transport_key,
                    &mut read_buffer,
                )
                .await?;

                // client sends 1 dumb byte, server replies with latest metrics
                // client polls manually
                let mut signal = [0u8; 1];
                loop {
                    // instant break server shutting down, no wait
                    if is_shutting_down.load(Ordering::Relaxed) {
                        break;
                    }
                    if read_h.read_exact(&mut signal).await.is_err() {
                        break; // client closed / read error
                    }
                    let metrics = ServerMetrics {
                        uptime_hrs: get_uptime_hrs(),
                        total_sessions: total_sessions.load(Ordering::Relaxed),
                        active_sessions: sessions.map.len() as u64,
                        total_network_bandwidth: total_bytes.load(Ordering::Relaxed),
                    };
                    if write_server_metrics(
                        &mut writer_h,
                        &metrics,
                        &transport_key,
                        &mut read_buffer,
                    )
                    .await
                    .is_err()
                    {
                        break; // broken pipe
                    }
                }
                info!(
                    "Observer (ip: {}) disconnected, closing connection...",
                    client_addr,
                );
                return Ok(());
            }
        }

        Ok(())
    }

    /// Returns the number of `active connections` to the server.
    pub fn get_active_connections(&self) -> u64 {
        self.active_connections.load(Ordering::Relaxed)
    }

    /// Returns the number of `active sessions` to the server.
    pub fn get_active_sessions(&self) -> u64 {
        self.session_map.map.len() as u64
    }

    /// Returns the total number of sessions ever created on the server.
    pub fn get_total_sessions(&self) -> u64 {
        self.total_sessions.load(Ordering::Relaxed)
    }

    /// Returns the observer token.
    pub fn get_observer_token(&self) -> Option<[u8; 64]> {
        self.config.observer_token
    }

    /// Returns the server uptime in seconds.
    pub fn get_uptime_secs(&self) -> f64 {
        if let Some(start_time) = START_TIME.get() {
            let now = chrono::Local::now();
            let duration = now.signed_duration_since(*start_time);
            duration.as_seconds_f64()
        } else {
            0.0
        }
    }

    /// Gracefully shuts down the server, stopping it from accepting `new connections` and waiting for `existing sessions` to finish.
    pub async fn shutdown(&self) -> Result<()> {
        self.shutdown_flag.store(true, Ordering::Relaxed);
        Ok(())
    }
}

#[allow(unused)]
async fn handle_producer() -> Result<()> {
    Ok(())
}
#[allow(unused)]
async fn handle_consumer() -> Result<()> {
    Ok(())
}
#[allow(unused)]
async fn handle_observer() -> Result<()> {
    Ok(())
}