monocoque-rs-zmtp 0.4.0

//! Base socket infrastructure shared by all ZMQ socket types.
//!
//! This module provides `SocketBase<S>` which contains all common fields and
//! low-level I/O operations used by DEALER, ROUTER, REQ, REP, PUB, SUB sockets.
//!
//! # Design Philosophy
//!
//! - **Zero-cost abstraction**: Composition-based, no vtables or dynamic dispatch
//! - **Single source of truth**: Common logic implemented once
//! - **Type safety**: Generic over stream type `S`
//! - **Protocol safety**: PoisonGuard integration for cancellation safety
//! - **Reconnection support**: Optional endpoint storage and backoff logic

use bytes::{BufMut, Bytes, BytesMut};
use compio_io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use monocoque_core::buffer::SegmentedBuffer;
use monocoque_core::endpoint::Endpoint;
use monocoque_core::io::take_read_buffer;
use monocoque_core::options::SocketOptions;
use monocoque_core::poison::PoisonGuard;
use monocoque_core::reconnect::ReconnectState;
use monocoque_core::rt::TcpStream;
use std::fmt;
use std::io;
use std::time::Instant;
use tracing::{debug, trace, warn};

use crate::codec::ZmtpDecoder;
use crate::handshake::perform_handshake_with_options;
use crate::session::SocketType;

// ─────────────────────────────────────────────────────────────────────────────
// ZMTP heartbeat helpers (RFC 23 / ZMTP 3.1)
// ─────────────────────────────────────────────────────────────────────────────

/// ZMTP PING command name (1-byte length prefix + "PING").
const PING_CMD: &[u8] = b"\x04PING";
/// ZMTP PONG command name (1-byte length prefix + "PONG").
const PONG_CMD: &[u8] = b"\x04PONG";

/// Build a ZMTP PING command frame.
///
/// Wire format (ZMTP command):
/// - `0x04` flag byte (COMMAND, short frame)
/// - 1-byte body length
/// - Body: `\x04PING` followed by 2-byte big-endian TTL in tenths of a second
///
/// The TTL tells the peer how long (in tenths of a second) it should wait
/// before considering the connection dead if it receives no traffic from us.
pub fn build_ping_frame(ttl_tenths: u16) -> Bytes {
    let mut body = BytesMut::with_capacity(PING_CMD.len() + 2);
    body.extend_from_slice(PING_CMD);
    body.put_u16(ttl_tenths);
    let body = body.freeze();
    crate::utils::encode_frame(crate::utils::FLAG_COMMAND, &body)
}

/// Build a ZMTP PONG command frame (reply to a received PING).
///
/// Wire format:
/// - `0x04` flag byte (COMMAND, short frame)
/// - 1-byte body length
/// - Body: `\x04PONG`
pub fn build_pong_frame() -> Bytes {
    let body = Bytes::from_static(PONG_CMD);
    crate::utils::encode_frame(crate::utils::FLAG_COMMAND, &body)
}

/// Return `true` if the decoded command payload begins with the PING name.
pub fn is_ping_payload(payload: &[u8]) -> bool {
    payload.starts_with(PING_CMD) && payload.len().saturating_sub(PING_CMD.len()) <= 18
}

/// Return `true` if the decoded command payload begins with the PONG name.
pub fn is_pong_payload(payload: &[u8]) -> bool {
    payload.starts_with(PONG_CMD) && payload.len().saturating_sub(PONG_CMD.len()) <= 16
}

/// Base socket infrastructure shared by all ZMQ socket types.
///
/// Contains all common fields and low-level I/O operations. Each socket type
/// (DEALER, ROUTER, REQ, REP, etc.) composes this struct and adds socket-specific
/// logic on top.
///
/// # Fields
///
/// - **Connection state**: `stream`, `endpoint`, `reconnect`
/// - **Buffers**: `recv`, `send_buffer`, `write_buf`, `read_buf`
/// - **Protocol**: `decoder`, `is_poisoned`
/// - **Configuration**: `config`, `options`
///
/// # Zero-Cost Abstraction
///
/// This is a plain struct with no vtable. The compiler can inline all methods,
/// resulting in zero runtime overhead compared to duplicating code in each socket.
pub struct SocketBase<S = TcpStream>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Underlying stream (TCP or Unix socket) - None when disconnected
    pub(crate) stream: Option<S>,

    /// Optional endpoint for automatic reconnection
    pub(crate) endpoint: Option<Endpoint>,

    /// Reconnection state tracker (exponential backoff)
    pub(crate) reconnect: Option<ReconnectState>,

    /// ZMTP frame decoder
    pub(crate) decoder: ZmtpDecoder,

    /// Segmented read buffer for incoming data
    pub(crate) recv: SegmentedBuffer,

    /// Reusable read slab; `take_read_buffer` carves each read off its tail.
    pub(crate) read_buf: BytesMut,

    /// Dedicated scratch buffer for reads that use a recv timeout. A timed read
    /// that elapses loses its buffer inside the cancelled future, so it must not
    /// come from the shared `read_buf` slab (that would waste the slab on every
    /// timeout). This buffer is reused across successful timed reads and only
    /// reallocated after an actual timeout drops it.
    pub(crate) recv_timeout_scratch: BytesMut,

    /// Absolute deadline for the logical recv currently in progress.
    ///
    /// `recv_timeout` bounds a whole `recv()` call, not each raw read. A single
    /// logical recv may loop over several `read_raw` reads (a multipart message
    /// arriving in pieces, or a partial frame); arming a fresh `timeout(dur, ..)`
    /// per read would let a peer that trickles one byte just under the timeout
    /// evade `recv_timeout` forever. Instead the first read of a logical recv
    /// arms this deadline once, every read in that recv races the *remaining*
    /// budget against it, and completing a message (`process_frame` returns a
    /// no-more-frames frame) clears it so the next recv starts fresh.
    pub(crate) recv_deadline: Option<Instant>,

    /// Reusable write buffer for outgoing data
    pub(crate) write_buf: BytesMut,

    /// Reusable iovec scratch for vectored writes (header/body entries),
    /// kept across calls so `send_vectored` allocates nothing on the hot path.
    pub(crate) iov: Vec<Bytes>,

    /// Send buffer for message batching
    pub(crate) send_buffer: BytesMut,

    /// Socket options (timeouts, limits, identity, buffer sizes)
    pub(crate) options: SocketOptions,

    /// Last connected/bound endpoint
    pub(crate) last_endpoint: Option<String>,

    /// Connection health flag (true if I/O was cancelled mid-operation)
    pub(crate) is_poisoned: bool,

    /// Number of messages currently buffered (for HWM enforcement)
    pub(crate) buffered_messages: usize,

    // ── Heartbeat state (ZMTP PING/PONG, RFC 23) ──────────────────────────
    //
    // Heartbeating keeps idle connections alive and detects dead peers.
    // When `options.heartbeat_ivl` is Some(dur):
    //   1. `last_recv_instant` is updated on every received frame.
    //   2. If `heartbeat_ivl` elapses with no received data, a PING command
    //      is sent and `ping_sent_at` records the transmission time.
    //   3. The peer must reply with a PONG within `heartbeat_timeout`
    //      (defaults to `heartbeat_ivl` when not set).  If it does not,
    //      `awaiting_pong` stays true and the next check can close the conn.
    //
    // Stub note: PING sending is wired into `check_heartbeat()`.
    // Timeout-based disconnection is noted below and left for a future PR.
    /// Instant of the last received frame on this connection.
    ///
    /// `None` before the first frame is received after the handshake, or
    /// when heartbeating is disabled.
    pub(crate) last_recv_instant: Option<Instant>,

    /// Instant at which the most recent PING was sent.
    ///
    /// `None` when no PING is outstanding.
    pub(crate) ping_sent_at: Option<Instant>,

    /// Whether we are waiting for a PONG reply to a sent PING.
    ///
    /// Set to `true` when a PING is transmitted; cleared when a matching
    /// PONG command frame is received.
    pub(crate) awaiting_pong: bool,

    /// Post-handshake CURVE cipher, if CURVE security is active.
    pub(crate) curve_cipher: Option<crate::security::curve::CurveMessageCipher>,

    /// Frames accumulated so far in the current logical (multipart) message.
    /// Reset when the MORE bit clears. Bounds an attacker who streams an endless
    /// run of MORE frames (`\x01\x01...`) to grow memory without bound, which the
    /// per-frame size cap does not catch.
    pub(crate) msg_frame_count: usize,
    /// Aggregate payload bytes accumulated in the current logical message.
    pub(crate) msg_byte_count: usize,
}

/// Maximum number of frames in one logical multipart message before the peer is
/// rejected. Legitimate multipart messages carry a handful of frames; a message
/// past this bound is treated as abusive.
const MAX_FRAMES_PER_MESSAGE: usize = 8192;
/// Maximum aggregate payload size of one logical multipart message.
const MAX_MESSAGE_AGGREGATE_BYTES: usize = 256 * 1024 * 1024;

pub fn append_zmtp_cmd_frame(buf: &mut BytesMut, body: &[u8]) {
    let len = body.len();
    if len <= 255 {
        buf.extend_from_slice(&[0x04, len as u8]);
    } else {
        buf.extend_from_slice(&[0x06]);
        buf.extend_from_slice(&(len as u64).to_be_bytes());
    }
    buf.extend_from_slice(body);
}

/// Result of processing one decoded ZMTP frame.
pub enum FrameResult {
    /// A data frame or decrypted CURVE MESSAGE: (more_flag, payload)
    Data(bool, Bytes),
    /// A command frame was handled internally (PING replied, PONG noted).
    /// The send_buffer may have been written; call flush_send_buffer if needed.
    CommandHandled,
    /// Need more data from the wire.
    NeedMore,
}

/// Apply equal jitter to a reconnect backoff delay, spreading the actual sleep
/// uniformly across `[delay/2, delay]`.
///
/// Without jitter, a fleet of clients that lost a restarted server would all
/// wake from the same backoff in lockstep and hammer it in synchronized waves.
/// Jitter decorrelates them while preserving the backoff's growth.
fn jittered_backoff(delay: std::time::Duration) -> std::time::Duration {
    use rand::Rng;
    if delay.is_zero() {
        return delay;
    }
    // Backoff delays are milliseconds to seconds, so the nanosecond count fits
    // in u64; use a checked conversion rather than a truncating cast so an
    // absurdly large delay saturates instead of wrapping.
    let half_ns = u64::try_from(delay.as_nanos() / 2).unwrap_or(u64::MAX);
    std::time::Duration::from_nanos(half_ns + rand::thread_rng().gen_range(0..=half_ns))
}

impl<S> SocketBase<S>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// Create a new SocketBase with the given stream and options.
    ///
    /// This is used when the socket is created from an existing stream
    /// (e.g., `from_tcp`, `from_unix_stream`). No endpoint or reconnection
    /// state is stored.
    ///
    /// Buffer sizes are taken from `options.read_buffer_size` and `options.write_buffer_size`.
    pub fn new(stream: S, _socket_type: SocketType, options: SocketOptions) -> Self {
        let write_capacity = options.write_buffer_size;
        let decoder = if let Some(max) = options.max_msg_size {
            ZmtpDecoder::with_max_frame_size(max)
        } else {
            ZmtpDecoder::new()
        };
        Self {
            stream: Some(stream),
            endpoint: None,
            reconnect: None,
            decoder,
            recv: SegmentedBuffer::new(),
            // Lazily allocated on the first read (matches the old arena, which
            // allocated no page until first use), so an idle socket holds none.
            read_buf: BytesMut::new(),
            recv_timeout_scratch: BytesMut::new(),
            recv_deadline: None,
            write_buf: BytesMut::with_capacity(write_capacity),
            iov: Vec::new(),
            send_buffer: BytesMut::with_capacity(write_capacity),
            options,
            last_endpoint: None,
            is_poisoned: false,
            buffered_messages: 0,
            // Heartbeat fields  -  initialised to idle state
            last_recv_instant: None,
            ping_sent_at: None,
            awaiting_pong: false,
            curve_cipher: None,
            msg_frame_count: 0,
            msg_byte_count: 0,
        }
    }

    /// Create a new SocketBase with endpoint storage for reconnection.
    ///
    /// This is used when the socket is created via `connect(endpoint)` and
    /// automatic reconnection is desired.
    ///
    /// Buffer sizes are taken from `options.read_buffer_size` and `options.write_buffer_size`.
    pub fn with_endpoint(
        stream: S,
        _socket_type: SocketType,
        endpoint: Endpoint,
        options: SocketOptions,
    ) -> Self {
        let endpoint_str = endpoint.to_string();
        let write_capacity = options.write_buffer_size;
        let decoder = if let Some(max) = options.max_msg_size {
            ZmtpDecoder::with_max_frame_size(max)
        } else {
            ZmtpDecoder::new()
        };
        Self {
            stream: Some(stream),
            endpoint: Some(endpoint),
            reconnect: Some(ReconnectState::new(&options)),
            decoder,
            recv: SegmentedBuffer::new(),
            // Lazily allocated on the first read (matches the old arena, which
            // allocated no page until first use), so an idle socket holds none.
            read_buf: BytesMut::new(),
            recv_timeout_scratch: BytesMut::new(),
            recv_deadline: None,
            write_buf: BytesMut::with_capacity(write_capacity),
            iov: Vec::new(),
            send_buffer: BytesMut::with_capacity(write_capacity),
            options,
            last_endpoint: Some(endpoint_str),
            is_poisoned: false,
            buffered_messages: 0,
            // Heartbeat fields  -  initialised to idle state
            last_recv_instant: None,
            ping_sent_at: None,
            awaiting_pong: false,
            curve_cipher: None,
            msg_frame_count: 0,
            msg_byte_count: 0,
        }
    }

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

    /// Check if the socket is poisoned (I/O was cancelled mid-operation).
    #[inline]
    pub const fn is_poisoned(&self) -> bool {
        self.is_poisoned
    }

    /// Get the number of buffered messages.
    #[inline]
    pub const fn buffered_messages(&self) -> usize {
        self.buffered_messages
    }

    /// Get the number of buffered bytes.
    #[inline]
    pub fn buffered_bytes(&self) -> usize {
        self.send_buffer.len()
    }

    /// Update live socket options and keep derived decoder state in sync.
    pub(crate) fn set_options(&mut self, options: SocketOptions) {
        self.decoder.set_max_body_len(options.max_msg_size);
        self.options = options;
    }

    /// Check if send HWM has been reached.
    #[inline]
    pub const fn hwm_reached(&self) -> bool {
        self.options.send_hwm != 0 && self.buffered_messages >= self.options.send_hwm
    }

    /// Get the endpoint this socket is connected/bound to, if any.
    ///
    /// Returns `None` if the socket was created from a raw stream without
    /// endpoint information.
    ///
    /// # ZeroMQ Compatibility
    ///
    /// Corresponds to `ZMQ_LAST_ENDPOINT` (32) option.
    #[inline]
    pub const fn last_endpoint(&self) -> Option<&Endpoint> {
        self.endpoint.as_ref()
    }

    /// Get the last endpoint as a string.
    ///
    /// # ZeroMQ Compatibility
    ///
    /// Corresponds to `ZMQ_LAST_ENDPOINT` (32) option.
    #[inline]
    pub fn last_endpoint_string(&self) -> Option<&str> {
        self.last_endpoint.as_deref()
    }

    /// Get the socket type.
    ///
    /// # ZeroMQ Compatibility
    ///
    /// Corresponds to `ZMQ_TYPE` (16) option.
    #[inline]
    // socket_type field is used internally for reconnection logic
    // Each socket implementation provides its own public socket_type() method
    /// Check if more message frames are expected (for multipart messages).
    ///
    /// This indicates whether the last received message has more frames
    /// coming after it. Always returns `false` for single-frame messages.
    ///
    /// # ZeroMQ Compatibility
    ///
    /// Corresponds to `ZMQ_RCVMORE` (13) option.
    pub fn has_more(&self) -> bool {
        self.decoder.has_more()
    }

    /// Get current socket events (read/write readiness).
    ///
    /// Returns a bitmask indicating which operations can proceed without blocking:
    /// - `POLLIN` (1): Socket has messages ready to read
    /// - `POLLOUT` (2): Socket can accept messages for sending
    ///
    /// # ZeroMQ Compatibility
    ///
    /// Corresponds to `ZMQ_EVENTS` (15) option.
    ///
    /// # Note
    ///
    /// This is a best-effort check based on current buffer state.
    /// For true async readiness, use the async recv/send operations.
    #[inline]
    pub fn events(&self) -> u32 {
        let mut events = 0u32;

        // POLLIN (1): Can receive if connected and buffers available
        if self.is_connected() && !self.is_poisoned {
            events |= 1; // POLLIN
        }

        // POLLOUT (2): Can send if connected and HWM not reached
        if self.is_connected() && !self.hwm_reached() && !self.is_poisoned {
            events |= 2; // POLLOUT
        }

        events
    }

    // ── Heartbeat helpers ────────────────────────────────────────────────────

    /// Record that a frame was received right now.
    ///
    /// Call this every time a complete ZMTP frame is read from the wire so that
    /// the heartbeat idle timer is reset.  When heartbeating is disabled
    /// (`options.heartbeat_ivl` is `None`) this is a no-op.
    #[inline]
    pub fn note_recv(&mut self) {
        if self.options.heartbeat_ivl.is_some() {
            self.last_recv_instant = Some(Instant::now());
        }
    }

    /// Record that a PONG was received, clearing the outstanding-PING flag.
    ///
    /// Call this when a command frame whose payload starts with `\x04PONG`
    /// is decoded in the active phase.
    #[inline]
    pub fn note_pong_received(&mut self) {
        if self.awaiting_pong {
            trace!("[SocketBase] PONG received  -  heartbeat round-trip complete");
            self.awaiting_pong = false;
            self.ping_sent_at = None;
        }
    }

    /// Check whether a PING should be sent and whether a pending PONG has
    /// timed out.
    ///
    /// This method should be called periodically in the socket's receive loop
    /// (e.g., after every successful `read_raw`).  It encodes the PING frame
    /// directly into `send_buffer` so the caller can flush it.
    ///
    /// # Return value
    ///
    /// - `Ok(true)`   -  a PING was appended to `send_buffer`; caller must flush.
    /// - `Ok(false)`  -  nothing to do or heartbeat is disabled.
    /// - `Err(e)`     -  a pending PONG timed out; connection should be closed.
    ///   Callers propagate this with `?` so the error surfaces to the application.
    pub fn check_heartbeat(&mut self) -> io::Result<bool> {
        let Some(ivl) = self.options.heartbeat_ivl else {
            return Ok(false);
        };

        let now = Instant::now();

        // ── Check PONG timeout ───────────────────────────────────────────
        if self.awaiting_pong {
            if let Some(ping_at) = self.ping_sent_at {
                let timeout = self.options.heartbeat_timeout.unwrap_or(ivl); // default to ivl when not set
                if now.duration_since(ping_at) > timeout {
                    warn!(
                        "[SocketBase] Heartbeat PONG not received within {:?}  -  peer considered dead",
                        timeout
                    );
                    // Mark disconnected
                    self.stream = None;
                    self.awaiting_pong = false;
                    self.ping_sent_at = None;
                    return Err(io::Error::new(
                        io::ErrorKind::TimedOut,
                        "ZMTP heartbeat: no PONG received within timeout",
                    ));
                }
            }
            // Still waiting for PONG  -  don't send another PING
            return Ok(false);
        }

        // ── Check whether we should send a PING ─────────────────────────
        let idle_since = self.last_recv_instant.unwrap_or_else(|| {
            // No frame received yet  -  treat as idle from the start
            now.checked_sub(ivl + ivl).unwrap_or(now)
        });

        if now.duration_since(idle_since) >= ivl {
            // Compute TTL to advertise: use heartbeat_ttl if set, else ivl
            let ttl_dur = self.options.heartbeat_ttl.unwrap_or(ivl);
            // ZMTP TTL is in tenths of a second, capped at u16::MAX
            let ttl_tenths = (ttl_dur.as_millis() / 100).min(u128::from(u16::MAX)) as u16;

            let ping = build_ping_frame(ttl_tenths);
            self.send_buffer.extend_from_slice(&ping);
            self.ping_sent_at = Some(now);
            self.awaiting_pong = true;

            debug!(
                "[SocketBase] Sending PING (ttl_tenths={}, idle={:?})",
                ttl_tenths,
                now.duration_since(idle_since)
            );
            return Ok(true);
        }

        Ok(false)
    }

    /// Read raw bytes from the stream into the recv buffer without decoding.
    ///
    /// This is the low-level read primitive used by socket implementations to
    /// accumulate multipart messages. Callers should manually decode frames
    /// from the recv buffer using `decoder.decode()`.
    ///
    /// Returns:
    /// - `Ok(n)` where n is the number of bytes read (n > 0)
    /// - `Ok(0)` if EOF was reached (connection closed)
    /// - `Err(e)` on I/O error
    ///
    /// On EOF, sets `stream = None` to mark disconnection.
    pub(crate) async fn read_raw(&mut self) -> io::Result<usize> {
        // Ensure we're connected
        if self.stream.is_none() {
            return Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "Socket not connected",
            ));
        }

        // Read from stream
        use compio_buf::BufResult;

        // Reject the non-blocking / zero-timeout case before taking a buffer.
        if matches!(self.options.recv_timeout, Some(dur) if dur.is_zero()) {
            return Err(io::Error::new(
                io::ErrorKind::WouldBlock,
                "Socket is in non-blocking mode and no data is available",
            ));
        }

        let read_size = self.options.read_buffer_size();

        let n = match self.options.recv_timeout {
            // No timeout: carve from the shared read slab and freeze in place.
            None => {
                // SAFETY: `buf` is passed straight to `read`; on the path that
                // exposes bytes it is first truncated to `n`, and the EOF/error
                // paths drop it without inspecting its contents.
                let buf = unsafe { take_read_buffer(&mut self.read_buf, read_size) };
                let stream = self
                    .stream
                    .as_mut()
                    .expect("BUG: stream must be Some  -  checked is_none() above");
                let BufResult(result, mut buf) = AsyncRead::read(stream, buf).await;
                let n = result?;
                if n == 0 {
                    trace!("[SocketBase] Connection closed (EOF)");
                    self.stream = None;
                    return Ok(0);
                }

                // Trim to what was actually read, then hand the unused tail of
                // the carve back to the slab so slab consumption tracks bytes
                // read, not the full carve. `reclaimed` (bytes past `n`) and
                // `self.read_buf` (the rest of the slab) are contiguous in the
                // same allocation, so `unsplit` rejoins them with no copy and the
                // next read reuses that space. This retires the old
                // per-small-frame copy_from_slice: freezing is zero-copy on every
                // size and the slab is not wasted on short reads.
                buf.truncate(n);
                let mut reclaimed = buf.split_off(n);
                // BytesMut::unsplit discards an empty `self` (`*self = other`),
                // and `reclaimed` has len 0, so give it a scratch len covering
                // its capacity first or the tail is dropped and the slab is not
                // reclaimed. SAFETY: uninitialized scratch that take_read_buffer
                // set_len's again before the next read overwrites it; never read
                // or exposed. `reclaimed` and `self.read_buf` are contiguous, so
                // unsplit merges without copying.
                unsafe {
                    reclaimed.set_len(reclaimed.capacity());
                }
                reclaimed.unsplit(std::mem::take(&mut self.read_buf));
                self.read_buf = reclaimed;
                self.recv.push(buf.freeze());
                n
            }

            // Timeout: read into a dedicated scratch buffer, never the shared
            // slab. A read that elapses loses its buffer inside the cancelled
            // future; keeping that off the slab means a timeout never wastes the
            // slab. On success the bytes are copied into recv and the scratch is
            // reused; on timeout the scratch is lost and the next read allocates.
            Some(dur) => {
                use monocoque_core::rt::timeout;

                // Arm the logical-recv deadline once; every read in this recv
                // then races the time that remains, so the whole recv is bounded
                // by `dur` however many reads a trickling peer forces.
                let now = Instant::now();
                let deadline = *self.recv_deadline.get_or_insert(now + dur);
                let remaining = deadline.saturating_duration_since(now);
                if remaining.is_zero() {
                    // Budget spent, including a deadline left stale by a dropped
                    // recv future: report the timeout and re-arm on the next call.
                    self.recv_deadline = None;
                    return Err(io::Error::new(
                        io::ErrorKind::TimedOut,
                        format!("Receive operation timed out after {:?}", dur),
                    ));
                }

                let mut scratch = std::mem::take(&mut self.recv_timeout_scratch);
                if scratch.capacity() < read_size {
                    scratch = BytesMut::with_capacity(read_size);
                }
                // SAFETY: uninitialized scratch; the read overwrites [0..n] and
                // only [0..n] is copied out below.
                unsafe {
                    scratch.set_len(read_size);
                }
                let stream = self
                    .stream
                    .as_mut()
                    .expect("BUG: stream must be Some  -  checked is_none() above");
                let Ok(BufResult(result, mut scratch)) =
                    timeout(remaining, AsyncRead::read(stream, scratch)).await
                else {
                    self.recv_deadline = None;
                    return Err(io::Error::new(
                        io::ErrorKind::TimedOut,
                        format!("Receive operation timed out after {:?}", dur),
                    ));
                };
                let n = result?;
                if n == 0 {
                    trace!("[SocketBase] Connection closed (EOF)");
                    self.stream = None;
                    return Ok(0);
                }
                self.recv.push(Bytes::copy_from_slice(&scratch[..n]));
                scratch.clear();
                self.recv_timeout_scratch = scratch;
                n
            }
        };

        // Update heartbeat idle timer: data was received so we are not idle
        self.note_recv();

        Ok(n)
    }

    /// Write buffered data from `send_buffer` to the stream.
    ///
    /// Uses PoisonGuard to ensure cancellation safety. If this method is
    /// cancelled during the write, the socket will be marked poisoned.
    ///
    /// Returns `Ok(())` on success, `Err(e)` on failure. On write failure,
    /// sets `stream = None` to mark disconnection.
    pub(crate) async fn flush_send_buffer(&mut self) -> io::Result<()> {
        if self.send_buffer.is_empty() {
            return Ok(());
        }

        // Check health before attempting I/O
        if self.is_poisoned {
            return Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "Socket poisoned by cancelled I/O - reconnect required",
            ));
        }

        // Ensure we have a connected stream
        let stream = self
            .stream
            .as_mut()
            .ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "Socket not connected"))?;

        if self.options.send_timeout.is_some_and(|dur| dur.is_zero()) {
            return Err(io::Error::new(
                io::ErrorKind::WouldBlock,
                "Socket is in non-blocking mode and cannot flush immediately",
            ));
        }

        trace!("[SocketBase] Flushing {} bytes", self.send_buffer.len());

        // Arm poison guard
        let guard = PoisonGuard::new(&mut self.is_poisoned);

        use compio_buf::BufResult;
        // Move send_buffer out by value: the owned-buffer write reuses this
        // allocation on completion instead of freezing a fresh Bytes per flush.
        let buf = std::mem::take(&mut self.send_buffer);

        // Apply send timeout
        let BufResult(result, buf) = match self.options.send_timeout {
            None => stream.write_all(buf).await,
            Some(dur) => {
                use monocoque_core::rt::timeout;
                match timeout(dur, stream.write_all(buf)).await {
                    Ok(result) => result,
                    Err(_) => {
                        return Err(io::Error::new(
                            io::ErrorKind::TimedOut,
                            format!("Flush operation timed out after {:?}", dur),
                        ));
                    }
                }
            }
        };

        // Recover the buffer's capacity for the next flush (see write_from_buf).
        let mut buf = buf;
        buf.clear();
        self.send_buffer = buf;

        let write_result = result;

        // If write failed, mark stream as disconnected
        if write_result.is_err() {
            self.stream = None;
        }

        write_result?;

        // Success - disarm guard and reset counter
        guard.disarm();
        self.buffered_messages = 0;

        trace!("[SocketBase] Flush completed");
        Ok(())
    }

    /// Write the contents of write_buf directly to the stream.
    ///
    /// This is used when the caller has already encoded data into write_buf
    /// and wants to send it without additional copying. Applies send_timeout
    /// from options and uses PoisonGuard for cancellation safety.
    pub(crate) async fn write_from_buf(&mut self) -> io::Result<()> {
        // Check health
        if self.is_poisoned {
            return Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "Socket poisoned by cancelled I/O",
            ));
        }

        // Ensure we have a connected stream
        let stream = self
            .stream
            .as_mut()
            .ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "Socket not connected"))?;

        if self.options.send_timeout.is_some_and(|dur| dur.is_zero()) {
            return Err(io::Error::new(
                io::ErrorKind::WouldBlock,
                "Socket is in non-blocking mode and cannot send immediately",
            ));
        }

        // Arm poison guard
        let guard = PoisonGuard::new(&mut self.is_poisoned);

        // Move write_buf out by value so the owned-buffer write reuses its
        // allocation instead of freezing a fresh refcounted Bytes each send:
        // no Arc alloc from freeze(), and the capacity comes straight back.
        let buf = std::mem::take(&mut self.write_buf);

        use compio_buf::BufResult;

        // Apply send timeout from options
        let BufResult(result, buf) = match self.options.send_timeout {
            None => {
                // Blocking mode - no timeout
                stream.write_all(buf).await
            }
            Some(dur) => {
                // Timed mode - apply timeout
                use monocoque_core::rt::timeout;
                match timeout(dur, stream.write_all(buf)).await {
                    Ok(result) => result,
                    Err(_) => {
                        return Err(io::Error::new(
                            io::ErrorKind::TimedOut,
                            format!("Send operation timed out after {:?}", dur),
                        ));
                    }
                }
            }
        };

        // Recover the buffer for the next send: clear() resets the length to
        // zero while keeping the capacity, so the hot path neither reallocates
        // nor churns a refcount. Reached only on completion; the timeout arm
        // returns early with the buffer still owned by the cancelled op.
        let mut buf = buf;
        buf.clear();
        self.write_buf = buf;

        // Mark disconnected on error
        if result.is_err() {
            self.stream = None;
        }

        result?;

        guard.disarm();
        Ok(())
    }

    /// Send one fully-assembled message over the best available write path.
    ///
    /// Mirrors the branch PUSH inlines (see `push::send`): opportunistic
    /// coalescing when enabled, an iovec (zero-copy) write for large frames,
    /// otherwise encode into the reused `write_buf` and write. Sharing it here
    /// lets DEALER/ROUTER/REQ/REP/PAIR reach the vectored and coalescing paths
    /// instead of always copying every body into `write_buf` for one `write`.
    pub(crate) async fn send_message(&mut self, msg: &[Bytes]) -> io::Result<()> {
        if self.options.write_coalescing {
            self.send_coalesced(msg).await
        } else if self.should_vectored_write(msg) {
            self.send_vectored(msg).await
        } else {
            self.encode_message_to_write_buf(msg)?;
            self.write_from_buf().await
        }
    }

    /// Return `true` if `msg` should be sent with a vectored write rather than
    /// the copy-into-`send_buffer` path.
    ///
    /// Vectored writes pay off once a frame body is large enough that copying it
    /// into the userspace send buffer dominates the per-message cost. Small
    /// frames stay on the copy path, where a single contiguous `write` beats the
    /// per-iovec bookkeeping. CURVE-encrypted connections never qualify: the
    /// cipher must transform each body into a fresh buffer regardless, so there
    /// is no copy to save.
    #[inline]
    pub(crate) fn should_vectored_write(&self, msg: &[Bytes]) -> bool {
        if self.curve_cipher.is_some() {
            return false;
        }
        let threshold = self.options.vectored_write_threshold;
        msg.iter().any(|frame| frame.len() >= threshold)
    }

    /// Send a multipart message using a vectored write, without copying any
    /// frame body into the userspace send buffer.
    ///
    /// Each frame contributes two iovec entries: a freshly built header (2 or 9
    /// bytes) and the frame body itself, an O(1) `Bytes::clone` with no data
    /// copy. The whole iovec list is handed to `write_vectored_all`, so the
    /// bodies travel straight to the kernel.
    ///
    /// Anything already pending in `send_buffer` is flushed first to preserve
    /// wire ordering. Uses `PoisonGuard` for cancellation safety and applies
    /// `send_timeout`, mirroring [`flush_send_buffer`](Self::flush_send_buffer).
    /// On write failure the stream is dropped (`stream = None`) to mark
    /// disconnection.
    pub(crate) async fn send_vectored(&mut self, msg: &[Bytes]) -> io::Result<()> {
        use crate::codec::write_frame_header;

        if msg.is_empty() {
            return Ok(());
        }

        // Preserve ordering: flush anything already buffered before this frame.
        if !self.send_buffer.is_empty() {
            self.flush_send_buffer().await?;
        }

        if self.is_poisoned {
            return Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "Socket poisoned by cancelled I/O - reconnect required",
            ));
        }

        if self.options.send_timeout.is_some_and(|dur| dur.is_zero()) {
            return Err(io::Error::new(
                io::ErrorKind::WouldBlock,
                "Socket is in non-blocking mode and cannot send immediately",
            ));
        }

        // Build all frame headers contiguously in the reused write_buf, then
        // slice each one back out (O(1), sharing write_buf's allocation). The
        // iovec list is reused across calls via `self.iov`, so the hot path
        // performs no per-message heap allocation.
        let last = msg.len() - 1;
        self.write_buf.clear();
        for (i, frame) in msg.iter().enumerate() {
            write_frame_header(&mut self.write_buf, frame.len(), i < last);
        }
        let mut headers = self.write_buf.split().freeze();

        let mut iovecs = std::mem::take(&mut self.iov);
        iovecs.clear();
        iovecs.reserve(msg.len() * 2);
        for frame in msg {
            let hlen = if frame.len() >= 256 { 9 } else { 2 };
            iovecs.push(headers.split_to(hlen));
            iovecs.push(frame.clone());
        }

        let Some(stream) = self.stream.as_mut() else {
            self.iov = iovecs; // keep the scratch capacity
            return Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "Socket not connected",
            ));
        };

        // Arm poison guard for cancellation safety.
        let guard = PoisonGuard::new(&mut self.is_poisoned);

        use compio_buf::BufResult;
        let BufResult(result, returned) = match self.options.send_timeout {
            None => stream.write_vectored_all(iovecs).await,
            Some(dur) => {
                use monocoque_core::rt::timeout;
                match timeout(dur, stream.write_vectored_all(iovecs)).await {
                    Ok(result) => result,
                    Err(_) => {
                        return Err(io::Error::new(
                            io::ErrorKind::TimedOut,
                            format!("Send operation timed out after {:?}", dur),
                        ));
                    }
                }
            }
        };

        // Reclaim the iovec allocation for the next call.
        self.iov = returned;

        if result.is_err() {
            self.stream = None;
        }
        result?;

        guard.disarm();
        Ok(())
    }

    /// Close the socket gracefully, honoring LINGER for any buffered send data.
    ///
    /// Coalesced-but-unflushed data in `send_buffer` is drained according to the
    /// `linger` option before the stream is shut down:
    /// - `Some(0)`: discard buffered data immediately.
    /// - `Some(dur)`: flush within `dur`, then close even if the flush did not
    ///   complete in time.
    /// - `None`: flush indefinitely until all buffered data is sent.
    ///
    /// Then sends a TCP FIN (or equivalent) and drops the connection. After this
    /// call `is_connected()` returns `false`.
    pub async fn close(&mut self) -> io::Result<()> {
        // Drain any coalesced-but-unflushed data per LINGER before shutdown, so
        // callers relying on close() to flush do not silently lose the tail of a
        // coalesced burst.
        if !self.send_buffer.is_empty() && self.stream.is_some() {
            match self.options.linger {
                Some(dur) if dur.is_zero() => {
                    // Linger 0: discard buffered data.
                    self.send_buffer.clear();
                    self.buffered_messages = 0;
                }
                Some(dur) => {
                    use monocoque_core::rt::timeout;
                    // Flush within the linger window; close anyway on timeout.
                    match timeout(dur, self.flush_send_buffer()).await {
                        Ok(Ok(())) | Err(_) => {}
                        Ok(Err(e)) => return Err(e),
                    }
                }
                None => {
                    // Linger indefinite: block until flushed.
                    self.flush_send_buffer().await?;
                }
            }
        }

        if let Some(ref mut stream) = self.stream {
            stream.shutdown().await?;
        }
        self.stream = None;
        Ok(())
    }

    /// Encode a multipart message into `write_buf`, encrypting if CURVE is active.
    pub fn encode_message_to_write_buf(&mut self, msg: &[Bytes]) -> io::Result<()> {
        use crate::codec::encode_multipart;
        self.write_buf.clear();
        if let Some(ref mut cipher) = self.curve_cipher {
            let last = msg.len().saturating_sub(1);
            for (i, frame) in msg.iter().enumerate() {
                let body = cipher
                    .encrypt_frame(frame, i < last)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
                append_zmtp_cmd_frame(&mut self.write_buf, &body);
            }
        } else {
            encode_multipart(msg, &mut self.write_buf);
        }
        Ok(())
    }

    /// Encode a multipart message into `send_buffer`, encrypting if CURVE is active.
    pub fn encode_message_to_send_buf(&mut self, msg: &[Bytes]) -> io::Result<()> {
        use crate::codec::encode_multipart;
        if let Some(ref mut cipher) = self.curve_cipher {
            let last = msg.len().saturating_sub(1);
            for (i, frame) in msg.iter().enumerate() {
                let body = cipher
                    .encrypt_frame(frame, i < last)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
                append_zmtp_cmd_frame(&mut self.send_buffer, &body);
            }
        } else {
            encode_multipart(msg, &mut self.send_buffer);
        }
        self.buffered_messages += 1;
        Ok(())
    }

    /// Encode `msg` into `send_buffer` and flush when the coalesce threshold is reached.
    ///
    /// This is the hot path used when `SocketOptions::write_coalescing` is enabled.
    /// It does **not** touch `buffered_messages` (which is reserved for the explicit
    /// `send_buffered` / `flush` batch API).  Callers must call `flush_send_buffer`
    /// after the last message in a burst.
    pub(crate) async fn send_coalesced(&mut self, msg: &[Bytes]) -> io::Result<()> {
        use crate::codec::encode_multipart;
        if let Some(ref mut cipher) = self.curve_cipher {
            let last = msg.len().saturating_sub(1);
            for (i, frame) in msg.iter().enumerate() {
                let body = cipher
                    .encrypt_frame(frame, i < last)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
                append_zmtp_cmd_frame(&mut self.send_buffer, &body);
            }
        } else {
            encode_multipart(msg, &mut self.send_buffer);
        }
        if self.send_buffer.len() >= self.options.write_coalesce_threshold {
            self.flush_send_buffer().await?;
        }
        Ok(())
    }

    /// Encode one data frame into `send_buffer`.
    ///
    /// Returns `true` when the coalescing threshold has been reached and the
    /// caller should flush.
    pub(crate) fn encode_one_coalesced(&mut self, frame: &Bytes) -> io::Result<bool> {
        if self.curve_cipher.is_none() {
            crate::codec::encode_single(frame, &mut self.send_buffer);
            return Ok(self.send_buffer.len() >= self.options.write_coalesce_threshold);
        }
        self.encode_one_curve_coalesced(frame)
    }

    fn encode_one_curve_coalesced(&mut self, frame: &Bytes) -> io::Result<bool> {
        let cipher = self
            .curve_cipher
            .as_mut()
            .expect("checked by encode_one_coalesced");
        let body = cipher
            .encrypt_frame(frame, false)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
        append_zmtp_cmd_frame(&mut self.send_buffer, &body);
        Ok(self.send_buffer.len() >= self.options.write_coalesce_threshold)
    }

    /// Account one data frame against the current logical message's caps.
    ///
    /// Enforces a bound on the frame count and aggregate size of a single
    /// multipart message. Without this a peer can stream MORE frames forever
    /// (`\x01\x01...`) and grow the socket's accumulator without bound; the
    /// per-frame size cap does not catch it. On exceed the connection is failed.
    /// The counters reset when a message completes (the MORE bit clears).
    fn account_multipart_frame(&mut self, more: bool, payload_len: usize) -> io::Result<()> {
        self.msg_frame_count += 1;
        self.msg_byte_count = self.msg_byte_count.saturating_add(payload_len);

        if self.msg_frame_count > MAX_FRAMES_PER_MESSAGE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "multipart message exceeds the per-message frame-count cap",
            ));
        }
        if self.msg_byte_count > MAX_MESSAGE_AGGREGATE_BYTES {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "multipart message exceeds the per-message aggregate-size cap",
            ));
        }

        if !more {
            // Logical message complete: reset for the next one.
            self.msg_frame_count = 0;
            self.msg_byte_count = 0;
            // The recv that produced this message is done; drop its deadline so
            // the next recv arms a fresh one (see `recv_deadline`).
            self.recv_deadline = None;
        }
        Ok(())
    }

    /// Decode the next frame from the receive buffer, handling CURVE decryption and PING/PONG.
    pub fn process_frame(&mut self) -> io::Result<FrameResult> {
        use crate::security::curve::CurveMessageCipher;
        match self
            .decoder
            .decode(&mut self.recv)
            .map_err(io::Error::from)?
        {
            None => Ok(FrameResult::NeedMore),
            Some(frame) => {
                if frame.is_command() {
                    // CURVE MESSAGE decryption
                    if let Some(ref mut cipher) = self.curve_cipher
                        && CurveMessageCipher::is_curve_message(&frame.payload)
                    {
                        let (more, payload) =
                            cipher.decrypt_frame(&frame.payload).map_err(|e| {
                                io::Error::new(io::ErrorKind::InvalidData, e.to_string())
                            })?;
                        self.account_multipart_frame(more, payload.len())?;
                        return Ok(FrameResult::Data(more, payload));
                    }
                    // PING/PONG
                    if is_ping_payload(&frame.payload) {
                        let pong = build_pong_frame();
                        self.send_buffer.extend_from_slice(&pong);
                    }
                    if is_pong_payload(&frame.payload) {
                        self.note_pong_received();
                    }
                    Ok(FrameResult::CommandHandled)
                } else {
                    // Reject raw data frames when CURVE is negotiated - all application
                    // messages must arrive as CURVE MESSAGE command frames.
                    if self.curve_cipher.is_some() {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "unexpected plaintext data frame in CURVE mode",
                        ));
                    }
                    let more = frame.more();
                    self.account_multipart_frame(more, frame.payload.len())?;
                    Ok(FrameResult::Data(more, frame.payload))
                }
            }
        }
    }
}

impl SocketBase<TcpStream> {
    /// Try to reconnect to the stored endpoint.
    ///
    /// This method:
    /// 1. Checks if endpoint is configured
    /// 2. Applies exponential backoff delay
    /// 3. Attempts new TCP connection
    /// 4. Performs ZMTP handshake
    /// 5. Resets socket state on success
    ///
    /// Returns `Ok(())` on successful reconnection, `Err(e)` otherwise.
    pub(crate) async fn try_reconnect(&mut self, socket_type: SocketType) -> io::Result<()> {
        // Can only reconnect if we have an endpoint
        let endpoint = self.endpoint.as_ref().ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::Unsupported,
                "Socket was not created with connect() - no endpoint stored for reconnection",
            )
        })?;

        // Apply the backoff delay if we have reconnection state. This is an
        // async sleep that yields the executor, so a reconnecting socket does
        // not stall other sockets colocated on the same single-threaded runtime.
        // (Earlier this had to be a blocking std::thread::sleep because compio
        // 0.10 left residual timer state after handshake timeouts that hung
        // rt::sleep; compio 0.19 fixed that, verified by
        // tests/sleep_after_timeout_probe.rs.)
        if let Some(reconnect) = &mut self.reconnect {
            let base_delay = reconnect.next_delay();
            let delay = jittered_backoff(base_delay);
            debug!(
                "[SocketBase] Reconnection attempt {} after {:?}",
                reconnect.attempt(),
                delay
            );
            monocoque_core::rt::sleep(delay).await;
        }

        // Attempt connection based on endpoint type
        let mut new_stream = match endpoint {
            Endpoint::Tcp(addr) => TcpStream::connect(addr).await?,
            #[cfg(unix)]
            Endpoint::Ipc(_) => {
                return Err(io::Error::new(
                    io::ErrorKind::Unsupported,
                    "IPC reconnection not supported for TcpStream base",
                ));
            }
            Endpoint::Inproc(_) => {
                return Err(io::Error::new(
                    io::ErrorKind::Unsupported,
                    "Inproc reconnection not supported for TcpStream base",
                ));
            }
        };

        // Re-apply TCP tuning to the fresh socket. The original connect set
        // TCP_NODELAY (and keepalive), but a reconnect is a brand-new fd that
        // starts with kernel defaults (Nagle on), so without this the socket
        // would silently run with Nagle enabled after any reconnect. This is a
        // one-time setsockopt at reconnect, off the send/recv hot path.
        crate::utils::configure_tcp_stream(&new_stream, &self.options, "RECONNECT")?;

        // Perform handshake  -  preserve routing identity from options
        let hr = perform_handshake_with_options(
            &mut new_stream,
            socket_type,
            self.options.routing_id.as_deref(),
            Some(self.options.handshake_timeout),
            &self.options,
        )
        .await
        .map_err(|e| io::Error::other(format!("Handshake failed during reconnect: {}", e)))?;

        // Success! Update socket state
        self.curve_cipher = hr.curve_cipher;
        self.stream = Some(new_stream);
        self.is_poisoned = false;
        self.recv = SegmentedBuffer::new();
        self.send_buffer.clear();
        self.buffered_messages = 0;

        // Reset heartbeat state for the fresh connection
        self.last_recv_instant = None;
        self.ping_sent_at = None;
        self.awaiting_pong = false;

        // Reset reconnection state
        if let Some(ref mut reconnect) = self.reconnect {
            reconnect.reset();
        }

        debug!("[SocketBase] Reconnection successful");
        Ok(())
    }
}

impl<S> fmt::Debug for SocketBase<S>
where
    S: AsyncRead + AsyncWrite + Unpin + fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SocketBase")
            .field("connected", &self.is_connected())
            .field("poisoned", &self.is_poisoned)
            .field("buffered_messages", &self.buffered_messages)
            .field("buffered_bytes", &self.buffered_bytes())
            .field("endpoint", &self.endpoint)
            .field("awaiting_pong", &self.awaiting_pong)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use compio_buf::{BufResult, IoBuf, IoBufMut};
    use std::collections::VecDeque;
    use std::io;
    use std::sync::{Arc, Mutex};

    #[test]
    fn jittered_backoff_stays_within_half_to_full_range() {
        use std::time::Duration;
        // Zero stays zero (no reconnect delay).
        assert_eq!(jittered_backoff(Duration::ZERO), Duration::ZERO);

        // Otherwise the jittered delay is always in [delay/2, delay]. Sample
        // enough draws to exercise the range without relying on a fixed seed.
        let delay = Duration::from_millis(100);
        for _ in 0..1000 {
            let j = jittered_backoff(delay);
            assert!(
                j >= delay / 2 && j <= delay,
                "jittered delay {j:?} out of [{:?}, {:?}]",
                delay / 2,
                delay
            );
        }
    }

    const PAYLOAD: &[u8] = b"abcdef";

    #[derive(Clone, Debug, Default)]
    struct WriteLog(Arc<Mutex<Vec<Vec<u8>>>>);

    impl WriteLog {
        fn push(&self, bytes: &[u8]) {
            self.0.lock().unwrap().push(bytes.to_vec());
        }

        fn bytes(&self) -> Vec<u8> {
            self.0.lock().unwrap().concat()
        }

        fn write_count(&self) -> usize {
            self.0.lock().unwrap().len()
        }

        fn is_empty(&self) -> bool {
            self.0.lock().unwrap().is_empty()
        }
    }

    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
    enum WriteStep {
        Bytes(usize),
        Error(io::ErrorKind),
    }

    #[derive(Debug)]
    struct ScriptedWriteStream {
        steps: VecDeque<WriteStep>,
        log: WriteLog,
    }

    impl ScriptedWriteStream {
        fn new(steps: impl IntoIterator<Item = WriteStep>) -> Self {
            Self {
                steps: steps.into_iter().collect(),
                log: WriteLog::default(),
            }
        }

        fn log(&self) -> WriteLog {
            self.log.clone()
        }
    }

    impl AsyncRead for ScriptedWriteStream {
        async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
            BufResult(Ok(0), buf)
        }
    }

    impl AsyncWrite for ScriptedWriteStream {
        async fn write<B: IoBuf>(&mut self, buf: B) -> BufResult<usize, B> {
            match self.steps.pop_front() {
                Some(WriteStep::Bytes(n)) => {
                    let n = n.min(buf.buf_len());
                    self.log.push(&buf.as_init()[..n]);
                    BufResult(Ok(n), buf)
                }
                Some(WriteStep::Error(kind)) => {
                    BufResult(Err(io::Error::new(kind, "scripted write error")), buf)
                }
                None => {
                    let n = buf.buf_len();
                    self.log.push(buf.as_init());
                    BufResult(Ok(n), buf)
                }
            }
        }

        async fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }

        async fn shutdown(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    #[derive(Clone, Copy, Debug)]
    enum WritePath {
        WriteFromBuf,
        FlushSendBuffer,
    }

    impl WritePath {
        fn buffer_payload(self, base: &mut SocketBase<ScriptedWriteStream>) {
            match self {
                Self::WriteFromBuf => base.write_buf.extend_from_slice(PAYLOAD),
                Self::FlushSendBuffer => {
                    base.send_buffer.extend_from_slice(PAYLOAD);
                    base.buffered_messages = 1;
                }
            }
        }

        async fn write(self, base: &mut SocketBase<ScriptedWriteStream>) -> io::Result<()> {
            match self {
                Self::WriteFromBuf => base.write_from_buf().await,
                Self::FlushSendBuffer => base.flush_send_buffer().await,
            }
        }

        fn assert_drained(self, base: &SocketBase<ScriptedWriteStream>) {
            match self {
                Self::WriteFromBuf => assert!(base.write_buf.is_empty()),
                Self::FlushSendBuffer => {
                    assert!(base.send_buffer.is_empty());
                    assert_eq!(base.buffered_messages, 0);
                }
            }
        }

        fn assert_payload_buffered(self, base: &SocketBase<ScriptedWriteStream>) {
            match self {
                Self::WriteFromBuf => assert_eq!(&base.write_buf[..], PAYLOAD),
                Self::FlushSendBuffer => {
                    assert_eq!(&base.send_buffer[..], PAYLOAD);
                    assert_eq!(base.buffered_messages, 1);
                }
            }
        }
    }

    #[derive(Clone, Copy, Debug)]
    enum StreamState {
        Connected,
        Disconnected,
    }

    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
    enum WriteStatus {
        Ok,
        Err(io::ErrorKind),
    }

    impl WriteStatus {
        fn from_result(result: &io::Result<()>) -> Self {
            match result {
                Ok(()) => Self::Ok,
                Err(err) => Self::Err(err.kind()),
            }
        }
    }

    #[derive(Debug)]
    struct ExpectedWriteOutcome {
        status: WriteStatus,
        written: Vec<u8>,
    }

    fn simulate_write_all(payload: &[u8], script: &[WriteStep]) -> ExpectedWriteOutcome {
        let mut offset = 0;
        let mut written = Vec::new();

        for step in script {
            if offset == payload.len() {
                return ExpectedWriteOutcome {
                    status: WriteStatus::Ok,
                    written,
                };
            }

            match *step {
                WriteStep::Bytes(0) => {
                    return ExpectedWriteOutcome {
                        status: WriteStatus::Err(io::ErrorKind::WriteZero),
                        written,
                    };
                }
                WriteStep::Bytes(n) => {
                    let n = n.min(payload.len() - offset);
                    written.extend_from_slice(&payload[offset..offset + n]);
                    offset += n;
                }
                WriteStep::Error(io::ErrorKind::Interrupted) => {}
                WriteStep::Error(kind) => {
                    return ExpectedWriteOutcome {
                        status: WriteStatus::Err(kind),
                        written,
                    };
                }
            }
        }

        if offset < payload.len() {
            written.extend_from_slice(&payload[offset..]);
        }

        ExpectedWriteOutcome {
            status: WriteStatus::Ok,
            written,
        }
    }

    fn byte_steps(chunks: impl IntoIterator<Item = usize>) -> Vec<WriteStep> {
        chunks.into_iter().map(WriteStep::Bytes).collect()
    }

    fn nonblocking_options() -> SocketOptions {
        SocketOptions {
            send_timeout: Some(std::time::Duration::ZERO),
            ..SocketOptions::default()
        }
    }

    fn socket_with_payload(
        path: WritePath,
        steps: impl IntoIterator<Item = WriteStep>,
        options: SocketOptions,
    ) -> (SocketBase<ScriptedWriteStream>, WriteLog) {
        let stream = ScriptedWriteStream::new(steps);
        let log = stream.log();
        let mut base = SocketBase::new(stream, SocketType::Dealer, options);
        path.buffer_payload(&mut base);
        (base, log)
    }

    async fn assert_short_writes_complete(
        path: WritePath,
        chunks: impl IntoIterator<Item = usize>,
    ) {
        let (mut base, log) =
            socket_with_payload(path, byte_steps(chunks), SocketOptions::default());

        path.write(&mut base).await.unwrap();

        assert_eq!(log.bytes(), PAYLOAD);
        assert!(log.write_count() > 1);
        assert!(base.stream.is_some());
        path.assert_drained(&base);
        assert!(!base.is_poisoned());
    }

    async fn assert_write_failure_after_progress(
        path: WritePath,
        script: impl IntoIterator<Item = WriteStep>,
        expected_kind: io::ErrorKind,
        expected_written: &[u8],
    ) {
        let (mut base, log) = socket_with_payload(path, script, SocketOptions::default());

        let err = path.write(&mut base).await.unwrap_err();

        assert_eq!(err.kind(), expected_kind);
        assert_eq!(log.bytes(), expected_written);
        assert!(base.stream.is_none());
        assert!(base.is_poisoned());
    }

    async fn assert_write_success(path: WritePath, script: impl IntoIterator<Item = WriteStep>) {
        let (mut base, log) = socket_with_payload(path, script, SocketOptions::default());

        path.write(&mut base).await.unwrap();

        assert_eq!(log.bytes(), PAYLOAD);
        assert!(base.stream.is_some());
        path.assert_drained(&base);
        assert!(!base.is_poisoned());
    }

    async fn assert_pre_io_error_preserves_buffer(
        path: WritePath,
        options: SocketOptions,
        stream_state: StreamState,
        expected_kind: io::ErrorKind,
    ) {
        let (mut base, log) = socket_with_payload(path, [], options);
        if matches!(stream_state, StreamState::Disconnected) {
            base.stream = None;
        }

        let err = path.write(&mut base).await.unwrap_err();

        assert_eq!(err.kind(), expected_kind);
        assert!(log.is_empty());
        match stream_state {
            StreamState::Connected => assert!(base.stream.is_some()),
            StreamState::Disconnected => assert!(base.stream.is_none()),
        }
        path.assert_payload_buffered(&base);
        assert!(!base.is_poisoned());
    }

    fn write_step_scripts(max_len: usize) -> Vec<Vec<WriteStep>> {
        const STEPS: &[WriteStep] = &[
            WriteStep::Bytes(0),
            WriteStep::Bytes(1),
            WriteStep::Bytes(2),
            WriteStep::Error(io::ErrorKind::Interrupted),
            WriteStep::Error(io::ErrorKind::BrokenPipe),
        ];

        fn extend_scripts(
            scripts: &mut Vec<Vec<WriteStep>>,
            current: &mut Vec<WriteStep>,
            steps: &[WriteStep],
            max_len: usize,
        ) {
            scripts.push(current.clone());
            if current.len() == max_len {
                return;
            }

            for step in steps {
                current.push(*step);
                extend_scripts(scripts, current, steps, max_len);
                current.pop();
            }
        }

        let mut scripts = Vec::new();
        extend_scripts(&mut scripts, &mut Vec::new(), STEPS, max_len);
        scripts
    }

    fn short_write_scripts(max_len: usize) -> Vec<Vec<WriteStep>> {
        write_step_scripts(max_len)
            .into_iter()
            .filter(|script| matches!(script.first(), Some(WriteStep::Bytes(1 | 2))))
            .collect()
    }

    async fn scripted_write_case_failures(path: WritePath, script: &[WriteStep]) -> Vec<String> {
        let stream = ScriptedWriteStream::new(script.iter().copied());
        let log = stream.log();
        let mut base = SocketBase::new(stream, SocketType::Dealer, SocketOptions::default());
        path.buffer_payload(&mut base);

        let expected = simulate_write_all(PAYLOAD, script);
        let result = path.write(&mut base).await;

        let mut failures = Vec::new();
        let actual_status = WriteStatus::from_result(&result);
        if actual_status != expected.status {
            failures.push(format!(
                "path={path:?} script={script:?}: result {actual_status:?} != {:?}",
                expected.status
            ));
        }

        let actual_written = log.bytes();
        if actual_written != expected.written {
            failures.push(format!(
                "path={path:?} script={script:?}: written {actual_written:?} != {:?}",
                expected.written
            ));
        }

        match expected.status {
            WriteStatus::Ok => {
                if base.stream.is_none() {
                    failures.push(format!(
                        "path={path:?} script={script:?}: stream disconnected after expected success"
                    ));
                }
                if base.is_poisoned() {
                    failures.push(format!(
                        "path={path:?} script={script:?}: socket poisoned after expected success"
                    ));
                }
                match path {
                    WritePath::WriteFromBuf if !base.write_buf.is_empty() => {
                        failures.push(format!(
                            "path={path:?} script={script:?}: write_buf not empty after expected success"
                        ));
                    }
                    WritePath::FlushSendBuffer if !base.send_buffer.is_empty() => {
                        failures.push(format!(
                            "path={path:?} script={script:?}: send_buffer not empty after expected success"
                        ));
                    }
                    WritePath::FlushSendBuffer if base.buffered_messages != 0 => {
                        failures.push(format!(
                            "path={path:?} script={script:?}: buffered_messages {} != 0 after expected success",
                            base.buffered_messages
                        ));
                    }
                    _ => {}
                }
            }
            WriteStatus::Err(_) => {
                if base.stream.is_some() {
                    failures.push(format!(
                        "path={path:?} script={script:?}: stream still connected after expected failure"
                    ));
                }
                if !base.is_poisoned() {
                    failures.push(format!(
                        "path={path:?} script={script:?}: socket not poisoned after expected failure"
                    ));
                }
            }
        }

        failures
    }

    async fn assert_scripted_write_cases(path: WritePath, scripts: Vec<Vec<WriteStep>>) {
        let mut failures = Vec::new();

        for script in scripts {
            failures.extend(scripted_write_case_failures(path, &script).await);
        }

        assert!(
            failures.is_empty(),
            "{} scripted write failures:\n{}",
            failures.len(),
            failures.join("\n")
        );
    }

    // ── PING / PONG frame builders ────────────────────────────────────────────

    #[test]
    fn test_build_ping_frame_structure() {
        // TTL = 100 tenths = 10 seconds
        let frame = build_ping_frame(100);
        // Byte 0: COMMAND flag (0x04), short frame
        assert_eq!(frame[0], 0x04, "PING frame must have COMMAND flag");
        // Byte 1: body length = 5 (\x04PING) + 2 (TTL) = 7
        assert_eq!(frame[1], 7, "PING body length must be 7 bytes");
        // Bytes 2-6: "\x04PING" (1-byte name-len + "PING")
        assert_eq!(
            &frame[2..7],
            b"\x04PING",
            "PING body must start with \\x04PING"
        );
        // Bytes 7-8: TTL big-endian
        let ttl = u16::from_be_bytes([frame[7], frame[8]]);
        assert_eq!(ttl, 100, "PING TTL field must match the argument");
    }

    #[test]
    fn test_build_pong_frame_structure() {
        let frame = build_pong_frame();
        // Byte 0: COMMAND flag (0x04)
        assert_eq!(frame[0], 0x04, "PONG frame must have COMMAND flag");
        // Byte 1: body length = 5 (\x04PONG = 1-byte length prefix + "PONG")
        assert_eq!(frame[1], 5, "PONG body length must be 5 bytes");
        // Bytes 2-6: "\x04PONG"
        assert_eq!(&frame[2..7], b"\x04PONG", "PONG body must be \\x04PONG");
    }

    #[test]
    fn test_is_ping_payload() {
        assert!(is_ping_payload(b"\x04PING\x00\x0A"));
        assert!(!is_ping_payload(b"\x04PONG"));
        assert!(!is_ping_payload(b"\x05READY"));
        assert!(!is_ping_payload(b""));
    }

    #[test]
    fn test_is_ping_payload_rejects_context_over_16_octets() {
        assert!(!is_ping_payload(b"\x04PING\x00\x0A12345678901234567"));
    }

    #[test]
    fn test_is_pong_payload() {
        assert!(is_pong_payload(b"\x04PONG"));
        assert!(!is_pong_payload(b"\x04PING\x00\x0A"));
        assert!(!is_pong_payload(b"\x05READY"));
        assert!(!is_pong_payload(b""));
    }

    #[test]
    fn test_is_pong_payload_rejects_context_over_16_octets() {
        assert!(!is_pong_payload(b"\x04PONG12345678901234567"));
    }

    // ── Heartbeat state helpers ───────────────────────────────────────────────

    /// `note_recv` must not panic when heartbeating is disabled.
    #[test]
    fn test_note_recv_no_op_when_disabled() {
        // We only test the public helpers  -  SocketBase itself requires a
        // concrete stream type which is difficult to instantiate in unit tests.
        // The logic here is purely tested through the helper functions.
        // Full integration is covered by the heartbeat field initialisation
        // verified in the constructor tests below.
        let _ = build_ping_frame(0); // smoke-test: no panic
        let _ = build_pong_frame(); // smoke-test: no panic
    }

    /// Verify that the PING frame's TTL is encoded as big-endian in tenths of
    /// a second.
    #[test]
    fn test_ping_ttl_encoding() {
        // 300 tenths = 30 seconds
        let frame = build_ping_frame(300);
        let ttl = u16::from_be_bytes([frame[7], frame[8]]);
        assert_eq!(ttl, 300);

        // Zero TTL is also valid (no peer-side timeout)
        let frame0 = build_ping_frame(0);
        let ttl0 = u16::from_be_bytes([frame0[7], frame0[8]]);
        assert_eq!(ttl0, 0);

        // Max u16
        let frame_max = build_ping_frame(u16::MAX);
        let ttl_max = u16::from_be_bytes([frame_max[7], frame_max[8]]);
        assert_eq!(ttl_max, u16::MAX);
    }

    #[test]
    fn test_write_from_buf_retries_short_writes_before_disarming() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_write_from_buf_retries_short_writes_before_disarming_impl());
    }

    async fn test_write_from_buf_retries_short_writes_before_disarming_impl() {
        assert_short_writes_complete(WritePath::WriteFromBuf, [2, 2]).await;
    }

    #[test]
    fn test_flush_send_buffer_retries_short_writes_before_disarming() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_flush_send_buffer_retries_short_writes_before_disarming_impl());
    }

    async fn test_flush_send_buffer_retries_short_writes_before_disarming_impl() {
        assert_short_writes_complete(WritePath::FlushSendBuffer, [2, 2]).await;
    }

    #[test]
    fn test_write_from_buf_write_zero_poisons_and_disconnects() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_write_from_buf_write_zero_poisons_and_disconnects_impl());
    }

    async fn test_write_from_buf_write_zero_poisons_and_disconnects_impl() {
        assert_write_failure_after_progress(
            WritePath::WriteFromBuf,
            [WriteStep::Bytes(2), WriteStep::Bytes(0)],
            io::ErrorKind::WriteZero,
            b"ab",
        )
        .await;
    }

    #[test]
    fn test_flush_send_buffer_write_zero_poisons_and_disconnects() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_flush_send_buffer_write_zero_poisons_and_disconnects_impl());
    }

    async fn test_flush_send_buffer_write_zero_poisons_and_disconnects_impl() {
        assert_write_failure_after_progress(
            WritePath::FlushSendBuffer,
            [WriteStep::Bytes(2), WriteStep::Bytes(0)],
            io::ErrorKind::WriteZero,
            b"ab",
        )
        .await;
    }

    #[test]
    fn test_write_from_buf_write_error_after_progress_poisons_and_disconnects() {
        monocoque_core::rt::LocalRuntime::new().unwrap().block_on(
            test_write_from_buf_write_error_after_progress_poisons_and_disconnects_impl(),
        );
    }

    async fn test_write_from_buf_write_error_after_progress_poisons_and_disconnects_impl() {
        assert_write_failure_after_progress(
            WritePath::WriteFromBuf,
            [
                WriteStep::Bytes(2),
                WriteStep::Error(io::ErrorKind::BrokenPipe),
            ],
            io::ErrorKind::BrokenPipe,
            b"ab",
        )
        .await;
    }

    #[test]
    fn test_flush_send_buffer_write_error_after_progress_poisons_and_disconnects() {
        monocoque_core::rt::LocalRuntime::new().unwrap().block_on(
            test_flush_send_buffer_write_error_after_progress_poisons_and_disconnects_impl(),
        );
    }

    async fn test_flush_send_buffer_write_error_after_progress_poisons_and_disconnects_impl() {
        assert_write_failure_after_progress(
            WritePath::FlushSendBuffer,
            [
                WriteStep::Bytes(2),
                WriteStep::Error(io::ErrorKind::BrokenPipe),
            ],
            io::ErrorKind::BrokenPipe,
            b"ab",
        )
        .await;
    }

    #[test]
    fn test_write_from_buf_interrupted_write_retries_without_poisoning() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_write_from_buf_interrupted_write_retries_without_poisoning_impl());
    }

    async fn test_write_from_buf_interrupted_write_retries_without_poisoning_impl() {
        assert_write_success(
            WritePath::WriteFromBuf,
            [
                WriteStep::Error(io::ErrorKind::Interrupted),
                WriteStep::Bytes(2),
                WriteStep::Bytes(4),
            ],
        )
        .await;
    }

    #[test]
    fn test_flush_send_buffer_interrupted_write_retries_without_poisoning() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_flush_send_buffer_interrupted_write_retries_without_poisoning_impl());
    }

    async fn test_flush_send_buffer_interrupted_write_retries_without_poisoning_impl() {
        assert_write_success(
            WritePath::FlushSendBuffer,
            [
                WriteStep::Error(io::ErrorKind::Interrupted),
                WriteStep::Bytes(2),
                WriteStep::Bytes(4),
            ],
        )
        .await;
    }

    #[test]
    fn test_scripted_write_from_buf_short_write_sequences_match_socket_state() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_scripted_write_from_buf_short_write_sequences_match_socket_state_impl());
    }

    async fn test_scripted_write_from_buf_short_write_sequences_match_socket_state_impl() {
        assert_scripted_write_cases(WritePath::WriteFromBuf, short_write_scripts(4)).await;
    }

    #[test]
    fn test_scripted_flush_send_buffer_short_write_sequences_match_socket_state() {
        monocoque_core::rt::LocalRuntime::new().unwrap().block_on(
            test_scripted_flush_send_buffer_short_write_sequences_match_socket_state_impl(),
        );
    }

    async fn test_scripted_flush_send_buffer_short_write_sequences_match_socket_state_impl() {
        assert_scripted_write_cases(WritePath::FlushSendBuffer, short_write_scripts(4)).await;
    }

    #[test]
    fn test_nonblocking_write_from_buf_keeps_buffer_and_health() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_nonblocking_write_from_buf_keeps_buffer_and_health_impl());
    }

    async fn test_nonblocking_write_from_buf_keeps_buffer_and_health_impl() {
        assert_pre_io_error_preserves_buffer(
            WritePath::WriteFromBuf,
            nonblocking_options(),
            StreamState::Connected,
            io::ErrorKind::WouldBlock,
        )
        .await;
    }

    #[test]
    fn test_write_from_buf_not_connected_does_not_poison() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_write_from_buf_not_connected_does_not_poison_impl());
    }

    async fn test_write_from_buf_not_connected_does_not_poison_impl() {
        assert_pre_io_error_preserves_buffer(
            WritePath::WriteFromBuf,
            SocketOptions::default(),
            StreamState::Disconnected,
            io::ErrorKind::NotConnected,
        )
        .await;
    }

    #[test]
    fn test_write_from_buf_not_connected_takes_precedence_over_nonblocking() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_write_from_buf_not_connected_takes_precedence_over_nonblocking_impl());
    }

    async fn test_write_from_buf_not_connected_takes_precedence_over_nonblocking_impl() {
        assert_pre_io_error_preserves_buffer(
            WritePath::WriteFromBuf,
            nonblocking_options(),
            StreamState::Disconnected,
            io::ErrorKind::NotConnected,
        )
        .await;
    }

    #[test]
    fn test_nonblocking_flush_keeps_buffer_and_health() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_nonblocking_flush_keeps_buffer_and_health_impl());
    }

    async fn test_nonblocking_flush_keeps_buffer_and_health_impl() {
        assert_pre_io_error_preserves_buffer(
            WritePath::FlushSendBuffer,
            nonblocking_options(),
            StreamState::Connected,
            io::ErrorKind::WouldBlock,
        )
        .await;
    }

    #[test]
    fn test_flush_send_buffer_not_connected_keeps_buffer_and_health() {
        monocoque_core::rt::LocalRuntime::new()
            .unwrap()
            .block_on(test_flush_send_buffer_not_connected_keeps_buffer_and_health_impl());
    }

    async fn test_flush_send_buffer_not_connected_keeps_buffer_and_health_impl() {
        assert_pre_io_error_preserves_buffer(
            WritePath::FlushSendBuffer,
            SocketOptions::default(),
            StreamState::Disconnected,
            io::ErrorKind::NotConnected,
        )
        .await;
    }

    #[test]
    fn test_flush_send_buffer_not_connected_takes_precedence_over_nonblocking() {
        monocoque_core::rt::LocalRuntime::new().unwrap().block_on(
            test_flush_send_buffer_not_connected_takes_precedence_over_nonblocking_impl(),
        );
    }

    async fn test_flush_send_buffer_not_connected_takes_precedence_over_nonblocking_impl() {
        assert_pre_io_error_preserves_buffer(
            WritePath::FlushSendBuffer,
            nonblocking_options(),
            StreamState::Disconnected,
            io::ErrorKind::NotConnected,
        )
        .await;
    }

    #[test]
    fn process_frame_caps_multipart_frame_count() {
        let stream = ScriptedWriteStream::new(Vec::<WriteStep>::new());
        let mut base = SocketBase::new(stream, SocketType::Dealer, SocketOptions::default());

        // Feed more MORE-frames than the cap allows. Each is a short data frame
        // with the MORE bit set: flags=0x01, length=1, payload 'x'.
        let more_frame = Bytes::from_static(b"\x01\x01x");
        for _ in 0..=MAX_FRAMES_PER_MESSAGE {
            base.recv.push(more_frame.clone());
        }

        let mut accepted = 0usize;
        let err = loop {
            match base.process_frame() {
                Ok(FrameResult::Data(_, _)) => accepted += 1,
                Ok(FrameResult::NeedMore) => panic!("ran out of data before the cap triggered"),
                Ok(_) => {}
                Err(e) => break e,
            }
        };

        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
        assert_eq!(
            accepted, MAX_FRAMES_PER_MESSAGE,
            "the decoder should accept exactly {MAX_FRAMES_PER_MESSAGE} frames of one \
             logical message, then reject the next"
        );
    }

    #[test]
    fn recv_timeout_reads_via_scratch_without_touching_the_slab() {
        use monocoque_core::rt::{LocalRuntime, TcpListener, TcpStream};
        use std::time::Duration;

        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            let peer = monocoque_core::rt::spawn(async move {
                let (mut s, _) = listener.accept().await.unwrap();
                use compio_io::AsyncWriteExt;
                let _ = s.write_all(b"hello").await;
                // Keep the connection open so the reader sees the bytes.
                monocoque_core::rt::sleep(Duration::from_millis(300)).await;
            });

            let stream = TcpStream::connect(addr).await.unwrap();
            let opts = SocketOptions::default().with_recv_timeout(Duration::from_secs(2));
            let mut base = SocketBase::new(stream, SocketType::Dealer, opts);

            let slab_cap_before = base.read_buf.capacity();
            let n = base.read_raw().await.unwrap();
            assert_eq!(n, 5, "should read the 5 bytes the peer sent");

            // The recv-timeout path used the dedicated scratch, never the slab.
            assert_eq!(
                base.read_buf.capacity(),
                slab_cap_before,
                "a recv-timeout read must not consume the shared read slab"
            );
            assert!(
                base.recv_timeout_scratch.capacity() > 0,
                "the scratch buffer should be retained for reuse after a successful timed read"
            );

            monocoque_core::rt::join(peer).await;
        });
    }

    #[test]
    fn recv_timeout_bounds_the_whole_logical_recv_not_each_read() {
        use monocoque_core::rt::{LocalRuntime, TcpListener, TcpStream};
        use std::time::{Duration, Instant};

        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            // The peer trickles one byte every 100ms and never sends a full
            // frame. Each individual read lands well inside the 300ms timeout,
            // so a per-read timer (the old bug) would never fire and the recv
            // would stall forever. The per-logical-recv deadline must fire ~300ms
            // in regardless of how many reads the trickle forces.
            let peer = monocoque_core::rt::spawn(async move {
                let (mut s, _) = listener.accept().await.unwrap();
                use compio_buf::BufResult;
                use compio_io::AsyncWriteExt;
                for _ in 0..40 {
                    let BufResult(res, _) = s.write_all(vec![0xffu8]).await;
                    // Stop once the reader is gone and the write fails.
                    if res.is_err() {
                        break;
                    }
                    monocoque_core::rt::sleep(Duration::from_millis(100)).await;
                }
            });

            let stream = TcpStream::connect(addr).await.unwrap();
            let opts = SocketOptions::default().with_recv_timeout(Duration::from_millis(300));
            let mut base = SocketBase::new(stream, SocketType::Dealer, opts);

            let start = Instant::now();
            let err = loop {
                match base.read_raw().await {
                    Ok(0) => panic!("unexpected EOF while trickling"),
                    // Got a trickle byte; keep reading, same as a socket recv
                    // loop that has not yet decoded a whole message.
                    Ok(_) => continue,
                    Err(e) => break e,
                }
            };
            let elapsed = start.elapsed();

            assert_eq!(err.kind(), io::ErrorKind::TimedOut);
            // Bounded by the 300ms deadline plus scheduling slack, NOT by
            // 40 * 100ms of trickling.
            assert!(
                elapsed < Duration::from_millis(1500),
                "the logical recv should time out ~300ms in, but took {elapsed:?}"
            );
            assert!(
                base.recv_deadline.is_none(),
                "the deadline must be cleared after a timeout so the next recv re-arms"
            );

            // Drop the reader so the peer's next write fails and it stops.
            drop(base);
            monocoque_core::rt::join(peer).await;
        });
    }

    #[test]
    fn completing_a_message_clears_the_recv_deadline() {
        use monocoque_core::rt::{LocalRuntime, TcpListener, TcpStream};
        use std::time::{Duration, Instant};

        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();
            let peer = monocoque_core::rt::spawn(async move {
                let _ = listener.accept().await.unwrap();
                monocoque_core::rt::sleep(Duration::from_millis(50)).await;
            });

            let stream = TcpStream::connect(addr).await.unwrap();
            let opts = SocketOptions::default().with_recv_timeout(Duration::from_millis(100));
            let mut base = SocketBase::new(stream, SocketType::Dealer, opts);

            base.recv_deadline = Some(Instant::now() + Duration::from_millis(100));
            // A non-final frame leaves the logical message in progress, so the
            // deadline stays armed across the reads that assemble it.
            base.account_multipart_frame(true, 4).unwrap();
            assert!(
                base.recv_deadline.is_some(),
                "mid-message: the recv deadline must stay armed"
            );
            // The final frame completes the message; the deadline is dropped so
            // the next recv starts a fresh budget.
            base.account_multipart_frame(false, 4).unwrap();
            assert!(
                base.recv_deadline.is_none(),
                "message complete: the recv deadline must be cleared"
            );

            monocoque_core::rt::join(peer).await;
        });
    }
}