moonpool_sim/network/sim/

stream.rs

use super::types::ConnectionId;
use crate::TcpListenerTrait;
use crate::sim::state::CloseReason;
use crate::{Event, WeakSimWorld};
use futures::io::{AsyncRead, AsyncWrite};
use std::{
    future::Future,
    io::{self, IoSlice},
    pin::Pin,
    task::{Context, Poll},
};
use tracing::instrument;

/// Create an `io::Error` for simulation shutdown.
///
/// Used when the simulation world has been dropped but stream operations are still attempted.
fn sim_shutdown_error() -> io::Error {
    io::Error::new(io::ErrorKind::BrokenPipe, "simulation shutdown")
}

/// Create an `io::Error` for random connection failure (chaos injection).
fn random_connection_failure_error() -> io::Error {
    io::Error::new(
        io::ErrorKind::ConnectionReset,
        "Random connection failure (explicit)",
    )
}

/// Create an `io::Error` for half-open connection timeout.
fn half_open_timeout_error() -> io::Error {
    io::Error::new(
        io::ErrorKind::ConnectionReset,
        "Connection reset (half-open timeout)",
    )
}

/// Create an `io::Error` for aborted connection (RST).
fn connection_aborted_error() -> io::Error {
    io::Error::new(
        io::ErrorKind::ConnectionReset,
        "Connection was aborted (RST)",
    )
}

/// Simulated TCP stream that implements async read/write operations.
///
/// `SimTcpStream` provides a realistic simulation of TCP socket behavior by implementing
/// the `AsyncRead` and `AsyncWrite` traits. It interfaces with the simulation event system
/// to provide ordered, reliable data delivery with configurable network delays.
///
/// ## Architecture Overview
///
/// Each `SimTcpStream` represents one endpoint of a TCP connection:
///
/// ```text
/// Application Layer          SimTcpStream Layer          Simulation Layer
/// ─────────────────          ──────────────────          ─────────────────
///                                                        
/// stream.write_all(data) ──► poll_write(data) ────────► buffer_send(data)
///                                                        └─► ProcessSendBuffer event
///                                                            └─► DataDelivery event
///                                                                └─► paired connection
///                                                        
/// stream.read(buf) ◄────── poll_read(buf) ◄──────────── receive_buffer
///                          │                           └─► waker registration
///                          └─► Poll::Pending/Ready     
/// ```
///
/// ## TCP Semantics Implemented
///
/// This implementation provides the core TCP guarantees required for realistic simulation:
///
/// ### 1. Reliable Delivery
/// - All written data will eventually be delivered to the paired connection
/// - No data loss (unless explicitly simulated via fault injection)
/// - Delivery confirmation through the event system
///
/// ### 2. Ordered Delivery (FIFO)
/// - Messages written first will arrive first at the destination
/// - Achieved through per-connection send buffering
/// - Critical for protocols that depend on message ordering
///
/// ### 3. Flow Control Simulation
/// - Read operations block (`Poll::Pending`) when no data is available
/// - Write operations complete immediately (buffering model)
/// - Backpressure handled at the application layer
///
/// ## Usage Examples
///
/// Provides async read/write operations for client and server connections.
///
/// ## Performance Characteristics
///
/// - **Write Latency**: O(1) - writes are buffered immediately
/// - **Read Latency**: `O(network_delay)` - depends on simulation configuration
/// - **Memory Usage**: `O(buffered_data)` - proportional to unread data
/// - **CPU Overhead**: Minimal - leverages efficient event system
///
/// ## Connection Lifecycle
///
/// 1. **Creation**: Stream created with reference to simulation and connection ID
/// 2. **Active Phase**: Read/write operations interact with simulation buffers
/// 3. **Data Transfer**: Asynchronous event processing handles network simulation
/// 4. **Termination**: Stream dropped when connection ends (automatic cleanup)
///
/// ## Thread Safety
///
/// `SimTcpStream` is `Send + Sync + Unpin + 'static` via its `Arc<RwLock<…>>`
/// backed `WeakSimWorld` handle. The simulation runtime itself runs on a
/// single OS thread (`new_current_thread().build()`), but the stream type
/// satisfies Send-bounded traits so it composes naturally with
/// `tokio::spawn`, hyper/reqwest connectors, and customer code that uses
/// `Arc<RwLock<…>>` / `DashMap` / `Arc<AtomicBool>`.
pub struct SimTcpStream {
    /// Weak reference to the simulation world.
    ///
    /// Uses `WeakSimWorld` to avoid circular references while allowing the stream
    /// to detect if the simulation has been dropped. Operations return errors
    /// gracefully if the simulation is no longer available.
    sim: WeakSimWorld,

    /// Unique identifier for this connection within the simulation.
    ///
    /// This ID corresponds to a `ConnectionState` entry in the simulation's
    /// connection table. Used to route read/write operations to the correct
    /// connection buffers and waker registrations.
    connection_id: ConnectionId,
}

impl SimTcpStream {
    /// Create a new simulated TCP stream
    pub(crate) fn new(sim: WeakSimWorld, connection_id: ConnectionId) -> Self {
        Self { sim, connection_id }
    }

    /// Get the connection ID (for test introspection and chaos injection)
    #[must_use]
    pub fn connection_id(&self) -> ConnectionId {
        self.connection_id
    }

    /// Returns `true`: `SimTcpStream` implements an efficient vectored write that
    /// records each `IoSlice` as its own ordered delivery event, so the chaos pack
    /// can act on individual segments.
    #[must_use]
    pub fn is_write_vectored(&self) -> bool {
        true
    }

    /// Run the chaos/closure checks that precede backpressure for a write.
    ///
    /// Returns `Some(poll)` to short-circuit, `None` to proceed.
    fn write_guard_pre_backpressure(
        &self,
        sim: &crate::sim::SimWorld,
        cx: &mut Context<'_>,
    ) -> Option<Poll<Result<usize, io::Error>>> {
        // Random close chaos injection (FDB rollRandomClose pattern)
        // Check at start of every write operation - sim2.actor.cpp:423
        // Returns Some(true) for explicit error, Some(false) for silent (connection marked closed)
        if let Some(true) = sim.roll_random_close(self.connection_id) {
            // 30% explicit exception - throw connection_failed immediately
            return Some(Poll::Ready(Err(random_connection_failure_error())));
            // 70% silent case: connection already marked as closed, will fail below
        }

        // Check if send side is closed (asymmetric closure)
        if sim.is_send_closed(self.connection_id) {
            return Some(Poll::Ready(Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "Connection send side closed",
            ))));
        }

        // Check if connection is closed or cut
        if sim.is_connection_closed(self.connection_id) {
            // Check how the connection was closed
            return Some(match sim.close_reason(self.connection_id) {
                CloseReason::Aborted => Poll::Ready(Err(connection_aborted_error())),
                _ => Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::BrokenPipe,
                    "Connection was closed (FIN)",
                ))),
            });
        }

        if sim.is_connection_cut(self.connection_id) {
            // Connection is temporarily cut - register waker and wait for restoration
            tracing::debug!(
                "SimTcpStream::poll_write connection_id={} is cut, registering cut waker",
                self.connection_id.0
            );
            sim.register_cut_waker(self.connection_id, cx.waker().clone());
            tracing::debug!(
                "SimTcpStream::poll_write connection_id={} registered waker for cut connection",
                self.connection_id.0
            );
            return Some(Poll::Pending);
        }

        // Check for half-open connection (peer crashed)
        if sim.is_half_open(self.connection_id) && sim.should_half_open_error(self.connection_id) {
            // Error time reached - return ECONNRESET
            tracing::debug!(
                "SimTcpStream::poll_write connection_id={} half-open error time reached, returning ECONNRESET",
                self.connection_id.0
            );
            return Some(Poll::Ready(Err(half_open_timeout_error())));
        }
        // Half-open but not yet error time - writes succeed but data goes nowhere
        // (paired_connection is already None, so buffer_send will silently succeed)

        None
    }

    /// Run the write-clog checks after the backpressure decision.
    ///
    /// Returns `Some(poll)` to short-circuit, `None` to proceed.
    fn write_guard_clog(
        &self,
        sim: &crate::sim::SimWorld,
        cx: &mut Context<'_>,
    ) -> Option<Poll<Result<usize, io::Error>>> {
        // Phase 7: Check for write clogging
        if sim.is_write_clogged(self.connection_id) {
            // Already clogged, register waker and return Pending
            sim.register_clog_waker(self.connection_id, cx.waker().clone());
            return Some(Poll::Pending);
        }

        // Check if this write should be clogged
        if sim.should_clog_write(self.connection_id) {
            sim.clog_write(self.connection_id);
            sim.register_clog_waker(self.connection_id, cx.waker().clone());
            return Some(Poll::Pending);
        }

        None
    }
}

impl Drop for SimTcpStream {
    fn drop(&mut self) {
        // Close the connection in the simulation when the stream is dropped
        // This matches real TCP behavior where dropping a socket always closes it
        if let Ok(sim) = self.sim.upgrade() {
            tracing::debug!(
                "SimTcpStream dropping, closing connection {}",
                self.connection_id.0
            );
            sim.close_connection(self.connection_id);
        }
    }
}

impl AsyncRead for SimTcpStream {
    #[instrument(skip(self, cx, buf))]
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        tracing::trace!(
            "SimTcpStream::poll_read called on connection_id={}",
            self.connection_id.0
        );
        let sim = self.sim.upgrade().map_err(|_| sim_shutdown_error())?;

        // Random close chaos injection (FDB rollRandomClose pattern)
        // Check at start of every read operation - sim2.actor.cpp:408
        // Returns Some(true) for explicit error, Some(false) for silent (connection marked closed)
        if let Some(true) = sim.roll_random_close(self.connection_id) {
            // 30% explicit exception - throw connection_failed immediately
            return Poll::Ready(Err(random_connection_failure_error()));
            // 70% silent case: connection already marked as closed, will return EOF below
        }

        // Check if receive side is closed (asymmetric closure)
        if sim.is_recv_closed(self.connection_id) {
            tracing::debug!(
                "SimTcpStream::poll_read connection_id={} recv side closed, returning EOF",
                self.connection_id.0
            );
            return Poll::Ready(Ok(0)); // EOF
        }

        // Check for half-open connection (peer crashed)
        if sim.is_half_open(self.connection_id) && sim.should_half_open_error(self.connection_id) {
            // Error time reached - return ECONNRESET
            tracing::debug!(
                "SimTcpStream::poll_read connection_id={} half-open error time reached, returning ECONNRESET",
                self.connection_id.0
            );
            return Poll::Ready(Err(half_open_timeout_error()));
        }
        // Half-open but not yet error time - will block (Pending) below waiting for data
        // that will never come, which is the correct half-open behavior

        // Check for read clogging (symmetric with write clogging)
        if sim.is_read_clogged(self.connection_id) {
            // Already clogged, register waker and return Pending
            sim.register_read_clog_waker(self.connection_id, cx.waker().clone());
            return Poll::Pending;
        }

        // Check if this read should be clogged
        if sim.should_clog_read(self.connection_id) {
            sim.clog_read(self.connection_id);
            sim.register_read_clog_waker(self.connection_id, cx.waker().clone());
            return Poll::Pending;
        }

        // Try to read from connection's receive buffer first
        // We should be able to read buffered data even if connection is currently cut
        let mut temp_buf = vec![0u8; buf.len()];
        let bytes_read = sim
            .read_from_connection(self.connection_id, &mut temp_buf)
            .map_err(|e| io::Error::other(format!("read error: {e}")))?;

        tracing::trace!(
            "SimTcpStream::poll_read connection_id={} read {} bytes",
            self.connection_id.0,
            bytes_read
        );

        if bytes_read > 0 {
            let data_preview = String::from_utf8_lossy(&temp_buf[..std::cmp::min(bytes_read, 20)]);
            tracing::trace!(
                "SimTcpStream::poll_read connection_id={} returning data: '{}'",
                self.connection_id.0,
                data_preview
            );
            buf[..bytes_read].copy_from_slice(&temp_buf[..bytes_read]);
            return Poll::Ready(Ok(bytes_read));
        }
        Self::poll_read_no_data(&sim, self.connection_id, cx, buf)
    }
}

impl SimTcpStream {
    /// Handle the `poll_read` branch where no buffered data is currently
    /// available. Checks for graceful FIN, abort, cut, then registers a
    /// read waker and rechecks the buffer to avoid races.
    fn poll_read_no_data(
        sim: &crate::sim::SimWorld,
        connection_id: crate::network::sim::ConnectionId,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        // No data available - check if connection has received FIN, is closed, or cut.
        if sim.is_remote_fin_received(connection_id) {
            tracing::info!(
                "SimTcpStream::poll_read connection_id={} remote FIN received, returning EOF",
                connection_id.0
            );
            return Poll::Ready(Ok(0));
        }
        if sim.is_connection_closed(connection_id) {
            if sim.close_reason(connection_id) == CloseReason::Aborted {
                tracing::info!(
                    "SimTcpStream::poll_read connection_id={} was aborted (RST)",
                    connection_id.0
                );
                return Poll::Ready(Err(connection_aborted_error()));
            }
            tracing::info!(
                "SimTcpStream::poll_read connection_id={} closed gracefully (FIN)",
                connection_id.0
            );
            return Poll::Ready(Ok(0));
        }
        if sim.is_connection_cut(connection_id) {
            tracing::debug!(
                "SimTcpStream::poll_read connection_id={} is cut, registering cut waker",
                connection_id.0
            );
            sim.register_cut_waker(connection_id, cx.waker().clone());
            return Poll::Pending;
        }

        // Register for notification when data arrives, then re-check for race.
        tracing::trace!(
            "SimTcpStream::poll_read connection_id={} no data, registering waker",
            connection_id.0
        );
        sim.register_read_waker(connection_id, cx.waker().clone());

        let mut temp_buf_recheck = vec![0u8; buf.len()];
        let bytes_read_recheck = sim
            .read_from_connection(connection_id, &mut temp_buf_recheck)
            .map_err(|e| io::Error::other(format!("recheck read error: {e}")))?;

        if bytes_read_recheck > 0 {
            buf[..bytes_read_recheck].copy_from_slice(&temp_buf_recheck[..bytes_read_recheck]);
            return Poll::Ready(Ok(bytes_read_recheck));
        }

        // Final check after waker registration.
        if sim.is_remote_fin_received(connection_id) {
            return Poll::Ready(Ok(0));
        }
        if sim.is_connection_closed(connection_id) {
            if sim.close_reason(connection_id) == CloseReason::Aborted {
                return Poll::Ready(Err(connection_aborted_error()));
            }
            return Poll::Ready(Ok(0));
        }
        Poll::Pending
    }
}

impl AsyncWrite for SimTcpStream {
    #[instrument(skip(self, cx, buf))]
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        let sim = self.sim.upgrade().map_err(|_| sim_shutdown_error())?;

        if let Some(poll) = self.write_guard_pre_backpressure(&sim, cx) {
            return poll;
        }

        // Check for send buffer space (backpressure)
        let available_buffer = sim.available_send_buffer(self.connection_id);
        if available_buffer < buf.len() {
            // Not enough buffer space, register waker and return Pending
            tracing::debug!(
                "SimTcpStream::poll_write connection_id={} buffer full (available={}, needed={}), waiting",
                self.connection_id.0,
                available_buffer,
                buf.len()
            );
            sim.register_send_buffer_waker(self.connection_id, cx.waker().clone());
            return Poll::Pending;
        }

        if let Some(poll) = self.write_guard_clog(&sim, cx) {
            return poll;
        }

        // Use buffered send to maintain TCP ordering
        let data_preview = String::from_utf8_lossy(&buf[..std::cmp::min(buf.len(), 20)]);
        tracing::trace!(
            "SimTcpStream::poll_write buffering {} bytes: '{}' for ordered delivery",
            buf.len(),
            data_preview
        );

        // Buffer the data for ordered processing instead of direct event scheduling
        sim.buffer_send(self.connection_id, buf.to_vec())
            .map_err(|e| io::Error::other(format!("buffer send error: {e}")))?;

        Poll::Ready(Ok(buf.len()))
    }

    #[instrument(skip(self, cx, bufs))]
    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[IoSlice<'_>],
    ) -> Poll<Result<usize, io::Error>> {
        let sim = self.sim.upgrade().map_err(|_| sim_shutdown_error())?;

        if let Some(poll) = self.write_guard_pre_backpressure(&sim, cx) {
            return poll;
        }

        let total: usize = bufs.iter().map(|slice| slice.len()).sum();
        if total == 0 {
            return Poll::Ready(Ok(0));
        }

        // writev(2) partial-accept semantics: if there's SOME room, accept what
        // fits and report a short count; only block when there is NO room at all.
        let available = sim.available_send_buffer(self.connection_id);
        if available == 0 {
            sim.register_send_buffer_waker(self.connection_id, cx.waker().clone());
            return Poll::Pending;
        }

        if let Some(poll) = self.write_guard_clog(&sim, cx) {
            return poll;
        }

        let accepted = total.min(available);

        // Buffer each IoSlice as its own ordered delivery event, truncating the
        // boundary slice when `accepted < total`. Skip empty slices so they do not
        // create empty delivery events.
        let mut remaining = accepted;
        for slice in bufs {
            if remaining == 0 {
                break;
            }
            if slice.is_empty() {
                continue;
            }
            let take = remaining.min(slice.len());
            sim.buffer_send(self.connection_id, slice[..take].to_vec())
                .map_err(|e| io::Error::other(format!("buffer send error: {e}")))?;
            remaining -= take;
        }

        Poll::Ready(Ok(accepted))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        let sim = self.sim.upgrade().map_err(|_| sim_shutdown_error())?;

        // Close the connection in the simulation when close is called
        tracing::debug!(
            "SimTcpStream::poll_close closing connection {}",
            self.connection_id.0
        );
        sim.close_connection(self.connection_id);

        Poll::Ready(Ok(()))
    }
}

/// Future representing an accept operation
pub struct AcceptFuture {
    sim: WeakSimWorld,
    local_addr: String,
}

impl Future for AcceptFuture {
    type Output = io::Result<(SimTcpStream, String)>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let Ok(sim) = self.sim.upgrade() else {
            return Poll::Ready(Err(sim_shutdown_error()));
        };

        let Some(connection_id) = sim.pending_connection(&self.local_addr) else {
            // No connection available yet - register waker for when connection becomes available
            sim.register_accept_waker(&self.local_addr, cx.waker().clone());
            return Poll::Pending;
        };

        // Get accept delay from network configuration
        let delay = sim
            .with_network_config(|config| crate::network::sample_latency(&config.accept_latency));

        // Schedule accept completion event to advance simulation time
        sim.schedule_event(
            Event::Connection {
                id: connection_id.0,
                state: crate::ConnectionStateChange::ConnectionReady,
            },
            delay,
        );

        // FDB Pattern (sim2.actor.cpp:1149-1175):
        // Return the synthesized ephemeral peer address, not the client's real address.
        // This simulates real TCP where servers see client ephemeral ports.
        let peer_addr = sim
            .connection_peer_address(connection_id)
            .unwrap_or_else(|| "unknown:0".to_string());

        let stream = SimTcpStream::new(self.sim.clone(), connection_id);
        Poll::Ready(Ok((stream, peer_addr)))
    }
}

/// Simulated TCP listener
pub struct SimTcpListener {
    sim: WeakSimWorld,
    local_addr: String,
}

impl SimTcpListener {
    /// Create a new simulated TCP listener
    pub(crate) fn new(sim: WeakSimWorld, local_addr: String) -> Self {
        Self { sim, local_addr }
    }
}

impl TcpListenerTrait for SimTcpListener {
    type TcpStream = SimTcpStream;

    #[instrument(skip(self))]
    async fn accept(&self) -> io::Result<(Self::TcpStream, String)> {
        AcceptFuture {
            sim: self.sim.clone(),
            local_addr: self.local_addr.clone(),
        }
        .await
    }

    fn local_addr(&self) -> io::Result<String> {
        Ok(self.local_addr.clone())
    }
}