hydra_sync/

server.rs

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

/// `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);
///     tokio::spawn(async move { server.run().await });
/// }
/// ```
pub struct HydraServer {
    listener: TcpListener,
    config: Arc<HydraConfig>,
    sessions: Arc<SessionMap>,
}

/// Configuration for the `HydraServer` features.
pub struct HydraConfig {
    /// Track the number of active connections to the server.
    curr_connections: Arc<AtomicU64>,
    /// 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,
    /// 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,
        overflow_mode: ChannelOverflowStrategy,
    ) -> Self {
        if fixed_packet_length == 0 {
            panic!("fixed_packet_length must be greater than 0");
        }
        Self {
            curr_connections: Arc::new(AtomicU64::new(0)),
            max_connections,
            accept_timeout,
            fixed_packet_length: fixed_packet_length + (NONCE_LEN + TAG_LEN) as u64, // add aes-gcm overhead
            channel_capacity,
            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
    /// - overflow_mode: DropPacket
    fn default() -> Self {
        Self {
            curr_connections: Arc::new(AtomicU64::new(0)),
            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,
            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<(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<(Self, SocketAddr)> {
        let server_tcp = TcpListener::bind(socket_addr).await?;
        let server_addr = server_tcp.local_addr()?;
        Ok((
            Self {
                listener: server_tcp,
                config: Arc::new(hydra_config),
                sessions: Arc::new(SessionMap::init()),
            },
            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<()> {
        loop {
            if self.config.curr_connections.fetch_add(1, Ordering::Relaxed)
                >= self.config.max_connections
            {
                self.config.curr_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;
            }

            match self.listener.accept().await {
                Ok((stream, peer_addr)) => {
                    stream.set_nodelay(true).ok();
                    let connections = Arc::clone(&self.config.curr_connections);
                    let sessions = Arc::clone(&self.sessions);
                    let read_write_len = self.config.fixed_packet_length;
                    let channel_capacity = self.config.channel_capacity;
                    let overflow_mode = self.config.overflow_mode;
                    tokio::spawn(async move {
                        trace!("Accepted connection from: {}", peer_addr);
                        if let Err(e) = Self::handle_connection(
                            stream,
                            read_write_len,
                            channel_capacity,
                            overflow_mode,
                            sessions,
                        )
                        .await
                        {
                            error!("Connection handling error: {} from: {}", e, peer_addr);
                        }
                        connections.fetch_sub(1, Ordering::Release);
                    });
                }
                Err(e) => {
                    self.config.curr_connections.fetch_sub(1, Ordering::Release);
                    error!("Connection accepting error: {}", e);
                }
            }
        }
    }

    /// Handles an individual client connection.
    async fn handle_connection(
        mut stream: TcpStream,
        read_write_len: u64,
        channel_capacity: u64,
        on_channel_overflow: ChannelOverflowStrategy,
        sessions: Arc<SessionMap>,
    ) -> 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(read_write_len 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.session_id)[..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::Admin => unimplemented!(),
            Role::Producer => {
                // create session for producer, if already exists, return error
                if !sessions.try_create_session(join_header.session_id) {
                    warn!(
                        "Producer (ip: {}) failed to create session_uuid: {}...",
                        client_addr,
                        &encode(join_header.session_id)[..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.session_id)[..12]
                );
                // 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,
                    read_write_len - (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(read_write_len 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.session_id)
                            {
                                // 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
                                    match entry.value_mut().push(data.clone()) {
                                        Ok(()) => {}
                                        Err(FullError) => match on_channel_overflow {
                                            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;
                                            }
                                            // TODO: clone inside a loop?!
                                            match p.push(data.clone()) {
                                                Ok(()) => Step::Pushed, // set 'we are done, exit outside'
                                                Err(FullError) => Step::Wait(p.space_notify()), // set 'we wait outside'
                                            }
                                        };
                                        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(read_write_len as usize, 0);
                        }
                        Err(_) => {
                            error!(
                                "Producer (ip: {}) read error, closing connection for session_uuid: {}...",
                                client_addr,
                                &encode(join_header.session_id)[..12]
                            );
                            break;
                        }
                    }
                }

                info!(
                    "Removing session_uuid: {}...",
                    &encode(join_header.session_id)[..12]
                );

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

            Role::Consumer => {
                // check session exists, if not, return error
                let session = match sessions.try_get_session(join_header.session_id) {
                    Some(s) => s,
                    None => {
                        warn!(
                            "Consumer (ip: {}) failed to join session_uuid: {}...",
                            client_addr,
                            &encode(join_header.session_id)[..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.session_id)[..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,
                    read_write_len - (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(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.shutdown().await;
                                        error!(
                                            "Consumer (ip: {}) write error, closing connection for session: {}...",
                                            client_addr,
                                            &encode(join_header.session_id)[..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.session_id)[..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.session_id)[..12]
                                    );
                                    break // if read fails or connection closed, break
                                },
                                _ => {}
                            }
                        }
                    }
                }

                // remove own queue from the session so it doesn't linger as a zombie
                if let Some(session) = sessions.try_get_session(join_header.session_id) {
                    session.consumer_local_buffers.remove(&consumer_id);
                }
            }
        }

        // TODO; consumer queue and idx still remains in the session map

        Ok(())
    }

    pub async fn shutdown(&self) -> Result<()> {
        unimplemented!()
    }
}