smarm/src/cluster/transport/loopback.rs

//! In-memory loopback transport — a shipped **test** transport.
//!
//! Lets Phases 2–4 exercise protocol logic (connector, membership,
//! monitors) through the real transport trait and the real framed codec
//! without sockets or timing flake.
//!
//! Blocking model: calls block the **OS thread** on a condvar. That is the
//! right shape for plain `#[test]`s driving protocol state machines; it is
//! the wrong shape for scheduler threads. Do not drive a loopback conn from
//! inside an actor — use the TCP impl there.
//!
//! Semantics mirror TCP shutdown where it matters for the codec: bytes
//! written before `close` remain readable at the peer, which then sees EOF;
//! writes toward a closed peer fail with `BrokenPipe`. Write buffers are
//! unbounded, so writes never block — backpressure is not simulated.

use std::collections::{HashMap, VecDeque};
use std::io;
use std::sync::{Arc, Condvar, Mutex, MutexGuard};

use super::{Conn, Listener, Transport};

/// Poison-tolerant lock: a panicked holder in a *test* transport must not
/// cascade; the byte-queue state stays consistent under every early return.
fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
    match m.lock() {
        Ok(g) => g,
        Err(poisoned) => poisoned.into_inner(),
    }
}

// ---------------------------------------------------------------------------
// One direction of a duplex: a byte queue with close flags for both ends
// ---------------------------------------------------------------------------

#[derive(Default)]
struct PipeState {
    bytes: VecDeque<u8>,
    /// The writing end closed: readers drain remaining bytes, then EOF.
    write_closed: bool,
    /// The reading end closed: writers fail with `BrokenPipe`.
    read_closed: bool,
}

#[derive(Default)]
struct Pipe {
    state: Mutex<PipeState>,
    cv: Condvar,
}

impl Pipe {
    fn write_all(&self, buf: &[u8]) -> io::Result<()> {
        let mut st = lock(&self.state);
        if st.write_closed {
            return Err(io::Error::new(
                io::ErrorKind::NotConnected,
                "loopback conn closed locally",
            ));
        }
        if st.read_closed {
            return Err(io::Error::new(
                io::ErrorKind::BrokenPipe,
                "loopback peer closed",
            ));
        }
        st.bytes.extend(buf);
        self.cv.notify_all();
        Ok(())
    }

    fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        let mut st = lock(&self.state);
        loop {
            if !st.bytes.is_empty() {
                let n = st.bytes.len().min(buf.len());
                for (slot, byte) in buf.iter_mut().zip(st.bytes.drain(..n)) {
                    *slot = byte;
                }
                return Ok(n);
            }
            if st.write_closed || st.read_closed {
                return Ok(0); // EOF: peer closed, or our own end closed.
            }
            st = match self.cv.wait(st) {
                Ok(g) => g,
                Err(poisoned) => poisoned.into_inner(),
            };
        }
    }

    /// Close from the writer side: remaining bytes stay readable, then EOF.
    fn close_write(&self) {
        lock(&self.state).write_closed = true;
        self.cv.notify_all();
    }

    /// Close from the reader side: peer writes fail from now on.
    fn close_read(&self) {
        lock(&self.state).read_closed = true;
        self.cv.notify_all();
    }
}

// ---------------------------------------------------------------------------
// Conn: two pipes, one per direction
// ---------------------------------------------------------------------------

/// One end of an established loopback connection.
pub struct LoopbackConn {
    tx: Arc<Pipe>,
    rx: Arc<Pipe>,
    peer: String,
}

impl Conn for LoopbackConn {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.rx.read(buf)
    }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.tx.write_all(buf)
    }

    fn close(&mut self) {
        self.tx.close_write();
        self.rx.close_read();
    }

    fn peer_addr(&self) -> String {
        self.peer.clone()
    }
}

impl Drop for LoopbackConn {
    fn drop(&mut self) {
        self.close();
    }
}

fn conn_pair(listen_addr: &str, conn_no: u64) -> (LoopbackConn, LoopbackConn) {
    let a_to_b = Arc::new(Pipe::default());
    let b_to_a = Arc::new(Pipe::default());
    let dialer = LoopbackConn {
        tx: a_to_b.clone(),
        rx: b_to_a.clone(),
        peer: listen_addr.to_string(),
    };
    let accepted = LoopbackConn {
        tx: b_to_a,
        rx: a_to_b,
        peer: format!("{listen_addr}#dialer-{conn_no}"),
    };
    (dialer, accepted)
}

// ---------------------------------------------------------------------------
// Listener + registry
// ---------------------------------------------------------------------------

#[derive(Default)]
struct AcceptState {
    pending: VecDeque<LoopbackConn>,
    closed: bool,
}

#[derive(Default)]
struct AcceptQueue {
    state: Mutex<AcceptState>,
    cv: Condvar,
}

/// A bound loopback listen point.
pub struct LoopbackListener {
    addr: String,
    queue: Arc<AcceptQueue>,
    registry: Arc<Mutex<Registry>>,
}

impl Listener for LoopbackListener {
    fn accept(&mut self) -> io::Result<Box<dyn Conn>> {
        let mut st = lock(&self.queue.state);
        loop {
            if let Some(conn) = st.pending.pop_front() {
                return Ok(Box::new(conn));
            }
            if st.closed {
                return Err(io::Error::new(
                    io::ErrorKind::NotConnected,
                    "loopback listener closed",
                ));
            }
            st = match self.queue.cv.wait(st) {
                Ok(g) => g,
                Err(poisoned) => poisoned.into_inner(),
            };
        }
    }

    fn local_addr(&self) -> String {
        self.addr.clone()
    }
}

impl Drop for LoopbackListener {
    fn drop(&mut self) {
        lock(&self.registry).listeners.remove(&self.addr);
        let mut st = lock(&self.queue.state);
        st.closed = true;
        self.queue.cv.notify_all();
    }
}

#[derive(Default)]
struct Registry {
    listeners: HashMap<String, Arc<AcceptQueue>>,
    dial_count: u64,
}

/// The loopback transport. Addresses are arbitrary strings scoped to one
/// transport instance; distinct instances never see each other's listeners.
#[derive(Default)]
pub struct LoopbackTransport {
    registry: Arc<Mutex<Registry>>,
}

impl Transport for LoopbackTransport {
    fn dial(&self, addr: &str) -> io::Result<Box<dyn Conn>> {
        let (queue, conn_no) = {
            let mut reg = lock(&self.registry);
            reg.dial_count += 1;
            let no = reg.dial_count;
            match reg.listeners.get(addr) {
                Some(q) => (q.clone(), no),
                None => {
                    return Err(io::Error::new(
                        io::ErrorKind::ConnectionRefused,
                        format!("no loopback listener at {addr:?}"),
                    ));
                }
            }
        };
        let (dialer, accepted) = conn_pair(addr, conn_no);
        let mut st = lock(&queue.state);
        if st.closed {
            return Err(io::Error::new(
                io::ErrorKind::ConnectionRefused,
                format!("loopback listener at {addr:?} closed"),
            ));
        }
        st.pending.push_back(accepted);
        queue.cv.notify_all();
        Ok(Box::new(dialer))
    }

    fn listen(&self, addr: &str) -> io::Result<Box<dyn Listener>> {
        let queue = Arc::new(AcceptQueue::default());
        let mut reg = lock(&self.registry);
        if reg.listeners.contains_key(addr) {
            return Err(io::Error::new(
                io::ErrorKind::AddrInUse,
                format!("loopback listener already bound at {addr:?}"),
            ));
        }
        reg.listeners.insert(addr.to_string(), queue.clone());
        Ok(Box::new(LoopbackListener {
            addr: addr.to_string(),
            queue,
            registry: self.registry.clone(),
        }))
    }
}