monocoque-rs-zmtp 0.3.0

//! PUSH socket implementation
//!
//! PUSH sockets are send-only endpoints in the pipeline pattern. They distribute
//! messages in a round-robin fashion to connected PULL sockets.
//!
//! # Characteristics
//!
//! - **Send-only**: Cannot receive messages
//! - **Load balancing**: Distributes work across PULL sockets
//! - **Non-blocking**: Never blocks on slow receivers (drops if HWM reached)
//! - **Pipeline pattern**: For distributing tasks to workers
//!
//! # Use Cases
//!
//! - Task distribution (ventilator pattern)
//! - Parallel pipeline processing
//! - Work queue distribution

use crate::base::SocketBase;
use crate::{handshake::perform_handshake_with_options, session::SocketType};
use bytes::Bytes;
use compio_io::{AsyncRead, AsyncWrite};
use monocoque_core::options::SocketOptions;
use monocoque_core::rt::TcpStream;
use std::io;
use tracing::{debug, trace};

/// PUSH socket for distributing messages in a pipeline.
///
/// PUSH sockets send messages to connected PULL sockets in a round-robin
/// fashion, providing load balancing for parallel processing.
pub struct PushSocket<S = TcpStream>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Base socket infrastructure (stream, buffers, options)
    base: SocketBase<S>,
}

impl<S> PushSocket<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Create a new PUSH socket from a stream with default buffer configuration.
    pub async fn new(stream: S) -> io::Result<Self> {
        Self::with_options(stream, SocketOptions::default()).await
    }

    /// Create a new PUSH socket with custom buffer configuration and socket options.
    pub async fn with_options(mut stream: S, options: SocketOptions) -> io::Result<Self> {
        debug!("[PUSH] Creating new PUSH socket");

        // Perform ZMTP handshake
        debug!("[PUSH] Performing ZMTP handshake...");
        let handshake_result = perform_handshake_with_options(
            &mut stream,
            SocketType::Push,
            options.routing_id.as_deref(),
            Some(options.handshake_timeout),
            &options,
        )
        .await
        .map_err(|e| io::Error::other(format!("Handshake failed: {}", e)))?;

        debug!(
            peer_identity = ?handshake_result.peer_identity,
            peer_socket_type = ?handshake_result.peer_socket_type,
            "[PUSH] Handshake complete"
        );

        debug!("[PUSH] Socket initialized");

        let mut base = SocketBase::new(stream, SocketType::Push, options);
        base.curve_cipher = handshake_result.curve_cipher;
        Ok(Self { base })
    }

    /// Send a message to a connected PULL socket.
    ///
    /// Messages are distributed in a round-robin fashion when multiple
    /// PULL sockets are connected (in a multi-connection scenario).
    ///
    /// By default each call writes to the kernel immediately (eager mode, one
    /// io_uring operation per message). For throughput-bound pipelines, enable write
    /// coalescing via [`SocketOptions::with_write_coalescing`] and call
    /// [`flush`](Self::flush) after the last send in each burst. In coalesced mode,
    /// bytes may remain in userspace until the 64 KB threshold fills or `flush()` is
    /// called explicitly.
    ///
    /// # Errors
    ///
    /// Returns an error if the socket is poisoned, disconnected, or if the write fails.
    pub async fn send(&mut self, msg: Vec<Bytes>) -> io::Result<()> {
        trace!("[PUSH] Sending {} frames", msg.len());

        if self.base.options.write_coalescing {
            self.base.send_coalesced(&msg).await?;
        } else if self.base.should_vectored_write(&msg) {
            // Large frame: write header + body as an iovec, skipping the copy
            // into the userspace send buffer.
            self.base.send_vectored(&msg).await?;
        } else {
            self.base.encode_message_to_write_buf(&msg)?;
            self.base.write_from_buf().await?;
        }

        // Check heartbeat: send PING if the connection has been idle too long
        if self.base.check_heartbeat()? {
            self.base.flush_send_buffer().await?;
        }

        trace!("[PUSH] Message sent successfully");
        Ok(())
    }

    /// Send a single-frame message without allocating a one-element `Vec`.
    ///
    /// This is equivalent to `send(vec![frame])`, but keeps the hot path for
    /// single-frame PUSH/PULL pipelines from measuring the caller's multipart
    /// container allocation.
    pub async fn send_one(&mut self, frame: Bytes) -> io::Result<()> {
        trace!("[PUSH] Sending 1 frame");

        if self.base.options.write_coalescing {
            if self.base.encode_one_coalesced(&frame)? {
                self.base.flush_send_buffer().await?;
            }
        } else {
            let msg = std::slice::from_ref(&frame);
            if self.base.should_vectored_write(msg) {
                self.base.send_vectored(msg).await?;
            } else {
                self.base.encode_message_to_write_buf(msg)?;
                self.base.write_from_buf().await?;
            }
        }

        if self.base.check_heartbeat()? {
            self.base.flush_send_buffer().await?;
        }

        trace!("[PUSH] Message sent successfully");
        Ok(())
    }

    /// Flush any messages still buffered by write coalescing.
    ///
    /// Call this after the last `send()` in a burst when `write_coalescing` is
    /// enabled to ensure all pending data is written to the kernel.
    pub async fn flush(&mut self) -> io::Result<()> {
        self.base.flush_send_buffer().await
    }

    /// Encode and send a batch of messages in a single kernel write.
    ///
    /// Encodes every message in `msgs` into the send buffer, then flushes once.
    /// This gives the same kernel-call efficiency as write coalescing but with
    /// explicit batch boundaries - no threshold check and no `flush()` required.
    ///
    /// Works independently of the `write_coalescing` option and can be mixed
    /// with `send()` calls freely.
    ///
    /// Returns the number of messages sent.
    pub async fn send_batch<I>(&mut self, msgs: I) -> io::Result<usize>
    where
        I: IntoIterator<Item = Vec<Bytes>>,
    {
        let mut count = 0;
        for msg in msgs {
            trace!("[PUSH] Buffering batch message {}", count);
            self.base.encode_message_to_send_buf(&msg)?;
            count += 1;
        }
        if count > 0 {
            self.base.flush_send_buffer().await?;
        }
        if self.base.check_heartbeat()? {
            self.base.flush_send_buffer().await?;
        }
        trace!("[PUSH] Batch of {} messages sent", count);
        Ok(count)
    }

    /// Close the socket gracefully by shutting down the underlying stream.
    pub async fn close(mut self) -> io::Result<()> {
        trace!("[PUSH] Closing socket");
        self.base.close().await
    }

    /// Get a reference to the socket options.
    #[inline]
    pub const fn options(&self) -> &SocketOptions {
        &self.base.options
    }

    /// Get a mutable reference to the socket options.
    #[inline]
    pub fn options_mut(&mut self) -> &mut SocketOptions {
        &mut self.base.options
    }

    /// Set socket options (builder-style).
    #[inline]
    pub fn set_options(&mut self, options: SocketOptions) {
        self.base.set_options(options);
    }
}

// Specialized implementation for TCP streams to enable TCP_NODELAY
impl PushSocket<TcpStream> {
    /// Create a new PUSH socket from a TCP stream with TCP_NODELAY enabled.
    pub async fn from_tcp(stream: TcpStream) -> io::Result<Self> {
        Self::from_tcp_with_options(stream, SocketOptions::default()).await
    }

    /// Create a new PUSH socket from a TCP stream with TCP_NODELAY and custom options.
    pub async fn from_tcp_with_options(
        stream: TcpStream,
        options: SocketOptions,
    ) -> io::Result<Self> {
        // Configure TCP optimizations including keepalive
        crate::utils::configure_tcp_stream(&stream, &options, "PUSH")?;
        Self::with_options(stream, options).await
    }

    /// Connect to a remote PUSH socket, storing the endpoint for automatic reconnection.
    pub async fn connect(addr: impl monocoque_core::rt::ToSocketAddrs) -> io::Result<Self> {
        Self::connect_with_options(addr, SocketOptions::default()).await
    }

    /// Connect with custom options, storing the endpoint for reconnection.
    pub async fn connect_with_options(
        addr: impl monocoque_core::rt::ToSocketAddrs,
        options: SocketOptions,
    ) -> io::Result<Self> {
        let stream = TcpStream::connect(addr).await?;
        let peer_addr = stream.peer_addr()?;
        crate::utils::configure_tcp_stream(&stream, &options, "PUSH")?;

        let mut stream = stream;
        let handshake_result = perform_handshake_with_options(
            &mut stream,
            SocketType::Push,
            options.routing_id.as_deref(),
            Some(options.handshake_timeout),
            &options,
        )
        .await
        .map_err(|e| io::Error::other(format!("Handshake failed: {}", e)))?;

        debug!(
            peer_identity = ?handshake_result.peer_identity,
            peer_socket_type = ?handshake_result.peer_socket_type,
            "[PUSH] Connected to {} (endpoint stored for reconnection)",
            peer_addr
        );

        let endpoint = monocoque_core::endpoint::Endpoint::Tcp(peer_addr);
        let mut base =
            crate::base::SocketBase::with_endpoint(stream, SocketType::Push, endpoint, options);
        base.curve_cipher = handshake_result.curve_cipher;
        Ok(Self { base })
    }

    /// Check if the socket is currently connected.
    #[inline]
    pub fn is_connected(&self) -> bool {
        self.base.is_connected()
    }

    /// Try to reconnect to the stored endpoint.
    pub async fn try_reconnect(&mut self) -> io::Result<()> {
        self.base.try_reconnect(SocketType::Push).await
    }

    /// Send a message with automatic reconnection on network error.
    ///
    /// On BrokenPipe / ConnectionReset, `write_from_buf()` already sets
    /// `stream = None`, so the next loop iteration reconnects automatically.
    ///
    /// Respects `max_reconnect_attempts`  -  returns `NotConnected` when exhausted.
    pub async fn send_with_reconnect(&mut self, msg: Vec<Bytes>) -> io::Result<()> {
        let max = self.base.options.max_reconnect_attempts;
        let mut attempts = 0u32;

        loop {
            if self.base.stream.is_none() {
                if let Some(limit) = max
                    && attempts >= limit
                {
                    return Err(io::Error::new(
                        io::ErrorKind::NotConnected,
                        format!("Max {} reconnection attempts exceeded", limit),
                    ));
                }
                attempts += 1;
                trace!(
                    "[PUSH] Stream disconnected, reconnecting (attempt {})",
                    attempts
                );
                self.try_reconnect().await?;
            }

            match self.send(msg.clone()).await {
                Ok(()) => return Ok(()),
                Err(_) if self.base.stream.is_none() => {
                    // write_from_buf set stream = None → network error, retry
                    debug!("[PUSH] Send failed (stream lost), will reconnect");
                }
                Err(e) => return Err(e),
            }
        }
    }
}

crate::impl_socket_trait!(PushSocket<S>, SocketType::Push);

#[cfg(all(test, unix))]
mod tests {
    use super::*;
    use monocoque_core::options::SocketOptions;
    use monocoque_core::rt::{LocalRuntime, TcpListener};
    use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
    use std::sync::mpsc;
    use std::thread;

    /// Read `TCP_NODELAY` from a live fd without taking ownership of it.
    fn fd_nodelay(fd: RawFd) -> bool {
        let sock = unsafe { socket2::Socket::from_raw_fd(fd) };
        let nd = sock.nodelay().expect("query TCP_NODELAY");
        std::mem::forget(sock); // borrowed fd - do not close it
        nd
    }

    /// A reconnect opens a brand-new fd that starts at the kernel default
    /// (Nagle on), so the socket must re-apply `TCP_NODELAY`; otherwise latency
    /// silently degrades after any automatic reconnect. This drives a real
    /// connect followed by a forced reconnect and checks the live socket fd.
    #[test]
    fn tcp_nodelay_survives_reconnect() {
        let (port_tx, port_rx) = mpsc::channel::<u16>();
        let (done_tx, done_rx) = mpsc::channel::<()>();

        // Server accepts twice - the initial connection and the forced
        // reconnect - completing the ZMTP handshake each time with a real PULL
        // peer, and holds both open until the client has inspected its socket.
        let server = thread::spawn(move || {
            let rt = LocalRuntime::new().unwrap();
            rt.block_on(async move {
                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
                port_tx.send(listener.local_addr().unwrap().port()).unwrap();

                let (s1, _) = listener.accept().await.unwrap();
                let _peer1 = crate::pull::PullSocket::new(s1).await.unwrap();

                let (s2, _) = listener.accept().await.unwrap();
                let _peer2 = crate::pull::PullSocket::new(s2).await.unwrap();

                done_rx.recv().unwrap();
            });
        });

        let port = port_rx.recv().unwrap();
        let client = thread::spawn(move || {
            let rt = LocalRuntime::new().unwrap();
            rt.block_on(async move {
                let mut push =
                    PushSocket::connect_with_options(("127.0.0.1", port), SocketOptions::default())
                        .await
                        .unwrap();

                // The initial connection sets NODELAY (existing behavior).
                let fd0 = push.base.stream.as_ref().unwrap().as_raw_fd();
                assert!(fd_nodelay(fd0), "initial connect must set TCP_NODELAY");

                // Force a reconnect: a fresh fd that defaults to Nagle-on.
                push.try_reconnect().await.unwrap();

                let fd1 = push.base.stream.as_ref().unwrap().as_raw_fd();
                assert!(
                    fd_nodelay(fd1),
                    "TCP_NODELAY must be re-applied on the reconnected socket",
                );

                done_tx.send(()).unwrap();
            });
        });

        client.join().unwrap();
        server.join().unwrap();
    }
}