supermachine 0.1.1

// Single mio::Poll thread that owns ALL accepted TCP streams.
// Replaces a per-conn reader+writer-thread design (which cost
// ~50–100µs of pthread_create overhead per inbound connection).
// Architecture: one epoll thread, commands via mpsc + mio::Waker,
// no per-conn threads.

use std::collections::HashMap;
use std::fmt;
use std::io::{Read, Write};
use std::net::{TcpStream as StdTcpStream, UdpSocket as StdUdpSocket};
use std::os::unix::net::UnixStream as StdUnixStream;
use std::sync::mpsc;
use std::sync::Arc;
use std::thread;
use std::time::Instant;

use super::mux_profile::{self, Stage};
use mio::net::{TcpStream as MioTcpStream, UdpSocket as MioUdpSocket, UnixStream as MioUnixStream};
use mio::{Events, Interest, Poll, Registry, Token};

pub enum MuxerStream {
    Tcp(StdTcpStream),
    Unix(StdUnixStream),
    /// TCP with a pre-buffered prefix that must be pushed into the
    /// guest as if it had been read from the socket. Used by the
    /// SCM_RIGHTS handoff path: the router reads enough of the HTTP
    /// request to route, then passes the TCP fd plus those already-
    /// consumed bytes here. The muxer-io thread injects the prefix
    /// via `on_data` before adding READABLE interest, so the guest
    /// sees a contiguous byte stream.
    TcpWithPrefix(StdTcpStream, Vec<u8>),
}

/// Commands the dispatcher / muxer sends to the I/O thread.
pub enum MuxerCmd {
    Register {
        host_src_port: u32,
        stream: MuxerStream,
        on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync>,
    },
    Write {
        host_src_port: u32,
        bytes: Vec<u8>,
    },
    Close {
        host_src_port: u32,
    },
    /// Register a UDP socket for inbound datagrams. The dispatcher
    /// owns the SEND side via try_clone; we own the RECV side here
    /// and call on_data per datagram received.
    RegisterUdp {
        key: u32, // (cid, peer_port) packed → unique per UDP "conn"
        udp: StdUdpSocket,
        on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync>,
    },
    /// Drop every registered TCP/UDP. Used by `VsockMuxer::reset()`
    /// between pool-worker RESTOREs so the next dispatch starts
    /// with a clean io-thread state. The thread itself stays alive.
    Reset {
        done: mpsc::Sender<()>,
    },
}

struct Conn {
    stream: MioStream,
    host_src_port: u32,
    pending_out: Vec<u8>,
    on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync>,
    eof_pushed: bool,
}

enum MioStream {
    Tcp(MioTcpStream),
    Unix(MioUnixStream),
}

impl MioStream {
    fn register(
        &mut self,
        registry: &Registry,
        tk: Token,
        interest: Interest,
    ) -> std::io::Result<()> {
        match self {
            Self::Tcp(s) => registry.register(s, tk, interest),
            Self::Unix(s) => registry.register(s, tk, interest),
        }
    }

    fn reregister(
        &mut self,
        registry: &Registry,
        tk: Token,
        interest: Interest,
    ) -> std::io::Result<()> {
        match self {
            Self::Tcp(s) => registry.reregister(s, tk, interest),
            Self::Unix(s) => registry.reregister(s, tk, interest),
        }
    }

    fn deregister(&mut self, registry: &Registry) -> std::io::Result<()> {
        match self {
            Self::Tcp(s) => registry.deregister(s),
            Self::Unix(s) => registry.deregister(s),
        }
    }
}

impl Read for MioStream {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        match self {
            Self::Tcp(s) => s.read(buf),
            Self::Unix(s) => s.read(buf),
        }
    }
}

impl Write for MioStream {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        match self {
            Self::Tcp(s) => s.write(buf),
            Self::Unix(s) => s.write(buf),
        }
    }

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

struct UdpConn {
    udp: MioUdpSocket,
    on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync>,
}

#[derive(Debug)]
pub enum StartError {
    Poll(std::io::Error),
    Waker(std::io::Error),
    ThreadSpawn {
        name: String,
        source: std::io::Error,
    },
}

impl fmt::Display for StartError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StartError::Poll(e) => write!(f, "muxer I/O Poll::new: {e}"),
            StartError::Waker(e) => write!(f, "muxer I/O Waker::new: {e}"),
            StartError::ThreadSpawn { name, source } => {
                write!(f, "spawn thread {name}: {source}")
            }
        }
    }
}

impl std::error::Error for StartError {}

/// Spawn the I/O thread.
pub fn spawn() -> Result<(mpsc::Sender<MuxerCmd>, Arc<mio::Waker>), StartError> {
    let (tx, rx) = mpsc::channel::<MuxerCmd>();
    let poll = Poll::new().map_err(StartError::Poll)?;
    let waker = Arc::new(mio::Waker::new(poll.registry(), Token(0)).map_err(StartError::Waker)?);
    let waker_for_thread = waker.clone();
    let name = "vsock-muxer-io".to_string();
    thread::Builder::new()
        .name(name.clone())
        .spawn(move || run(poll, waker_for_thread, rx))
        .map_err(|source| StartError::ThreadSpawn { name, source })?;
    Ok((tx, waker))
}

fn run(mut poll: Poll, _waker: Arc<mio::Waker>, rx: mpsc::Receiver<MuxerCmd>) {
    let mut conns: HashMap<Token, Conn> = HashMap::new();
    let mut udps: HashMap<Token, UdpConn> = HashMap::new();
    let mut by_port: HashMap<u32, Token> = HashMap::new();
    let mut next_token: usize = 1;
    let mut events = Events::with_capacity(128);
    // Per-vsock-packet cap. Linux's virtio-vsock allocates skbs
    // sized after the inbound packet header's `len` field; if we
    // send more than its skb buffer can hold, the kernel panics
    // in skb_over_panic. Linux's vsock_stream typically allocates
    // 4 KiB skbs. Cap at 4 KiB to be safe. A future improvement
    // is to track peer_buf_alloc credit per-conn and use that as
    // the cap instead.
    const MAX_PKT_PAYLOAD: usize = 4096;
    let mut buf = [0u8; MAX_PKT_PAYLOAD];

    loop {
        // 1. Drain commands.
        loop {
            let cmd = match rx.try_recv() {
                Ok(c) => c,
                Err(mpsc::TryRecvError::Empty) => break,
                Err(mpsc::TryRecvError::Disconnected) => return,
            };
            match cmd {
                MuxerCmd::Register {
                    host_src_port,
                    stream,
                    on_data,
                } => {
                    let (mut stream, prefix) = match stream {
                        MuxerStream::Tcp(tcp) => {
                            tcp.set_nonblocking(true).ok();
                            (MioStream::Tcp(MioTcpStream::from_std(tcp)), Vec::new())
                        }
                        MuxerStream::Unix(unix) => {
                            unix.set_nonblocking(true).ok();
                            (MioStream::Unix(MioUnixStream::from_std(unix)), Vec::new())
                        }
                        MuxerStream::TcpWithPrefix(tcp, prefix) => {
                            tcp.set_nonblocking(true).ok();
                            (MioStream::Tcp(MioTcpStream::from_std(tcp)), prefix)
                        }
                    };
                    let tk = Token(next_token);
                    next_token += 1;
                    if let Err(e) = stream.register(poll.registry(), tk, Interest::READABLE) {
                        eprintln!("[muxer-io] register err: {e}");
                        continue;
                    }
                    // Inject any already-read bytes into the guest first,
                    // chunked to the per-packet payload cap that the
                    // poll loop uses for normal reads.
                    if !prefix.is_empty() {
                        for chunk in prefix.chunks(MAX_PKT_PAYLOAD) {
                            on_data(chunk.to_vec());
                        }
                    }
                    by_port.insert(host_src_port, tk);
                    conns.insert(
                        tk,
                        Conn {
                            stream,
                            host_src_port,
                            pending_out: Vec::new(),
                            on_data,
                            eof_pushed: false,
                        },
                    );
                }
                MuxerCmd::Write {
                    host_src_port,
                    bytes,
                } => {
                    let tk = match by_port.get(&host_src_port) {
                        Some(t) => *t,
                        None => continue,
                    };
                    if let Some(c) = conns.get_mut(&tk) {
                        // Inline write first (cheap when buf has room).
                        if c.pending_out.is_empty() {
                            let t0 = Instant::now();
                            match c.stream.write(&bytes) {
                                Ok(n) if n == bytes.len() => {
                                    mux_profile::record(
                                        Stage::IoTcpWrite,
                                        n,
                                        t0.elapsed().as_micros() as u64,
                                    );
                                    continue;
                                }
                                Ok(n) => {
                                    mux_profile::record(
                                        Stage::IoTcpWrite,
                                        n,
                                        t0.elapsed().as_micros() as u64,
                                    );
                                    c.pending_out.extend_from_slice(&bytes[n..]);
                                }
                                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                                    mux_profile::record(
                                        Stage::IoTcpWrite,
                                        0,
                                        t0.elapsed().as_micros() as u64,
                                    );
                                    c.pending_out.extend_from_slice(&bytes);
                                }
                                Err(_) => continue,
                            }
                        } else {
                            c.pending_out.extend_from_slice(&bytes);
                        }
                        let _ = c.stream.reregister(
                            poll.registry(),
                            tk,
                            Interest::READABLE | Interest::WRITABLE,
                        );
                    }
                }
                MuxerCmd::Close { host_src_port } => {
                    if let Some(tk) = by_port.remove(&host_src_port) {
                        if let Some(mut c) = conns.remove(&tk) {
                            let _ = c.stream.deregister(poll.registry());
                        }
                    }
                }
                MuxerCmd::RegisterUdp {
                    key: _,
                    udp,
                    on_data,
                } => {
                    udp.set_nonblocking(true).ok();
                    let mut mu = MioUdpSocket::from_std(udp);
                    let tk = Token(next_token);
                    next_token += 1;
                    if let Err(e) = poll.registry().register(&mut mu, tk, Interest::READABLE) {
                        eprintln!("[muxer-io] register udp err: {e}");
                        continue;
                    }
                    udps.insert(tk, UdpConn { udp: mu, on_data });
                }
                MuxerCmd::Reset { done } => {
                    for (_, mut c) in conns.drain() {
                        let _ = c.stream.deregister(poll.registry());
                    }
                    for (_, mut u) in udps.drain() {
                        let _ = poll.registry().deregister(&mut u.udp);
                    }
                    by_port.clear();
                    let _ = done.send(());
                }
            }
        }

        if let Err(e) = poll.poll(&mut events, Some(std::time::Duration::from_millis(50))) {
            if e.kind() == std::io::ErrorKind::Interrupted {
                continue;
            }
            eprintln!("[muxer-io] poll err: {e}");
            continue;
        }

        let mut to_close: Vec<Token> = Vec::new();
        for ev in events.iter() {
            let tk = ev.token();
            if tk == Token(0) {
                continue;
            }
            // UDP recv side?
            if let Some(uc) = udps.get_mut(&tk) {
                if ev.is_readable() {
                    let mut dbuf = [0u8; 64 * 1024];
                    loop {
                        match uc.udp.recv_from(&mut dbuf) {
                            Ok((n, _peer)) => (uc.on_data)(dbuf[..n].to_vec()),
                            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
                            Err(_) => break,
                        }
                    }
                }
                continue;
            }
            let Some(c) = conns.get_mut(&tk) else {
                continue;
            };
            if ev.is_readable() {
                loop {
                    let t0 = Instant::now();
                    match c.stream.read(&mut buf) {
                        Ok(0) => {
                            mux_profile::record(
                                Stage::IoTcpRead,
                                0,
                                t0.elapsed().as_micros() as u64,
                            );
                            if !c.eof_pushed {
                                let cb_t0 = Instant::now();
                                (c.on_data)(Vec::new());
                                mux_profile::record(
                                    Stage::IoCallback,
                                    0,
                                    cb_t0.elapsed().as_micros() as u64,
                                );
                                c.eof_pushed = true;
                            }
                            to_close.push(tk);
                            break;
                        }
                        Ok(n) => {
                            mux_profile::record(
                                Stage::IoTcpRead,
                                n,
                                t0.elapsed().as_micros() as u64,
                            );
                            let cb_t0 = Instant::now();
                            (c.on_data)(buf[..n].to_vec());
                            mux_profile::record(
                                Stage::IoCallback,
                                n,
                                cb_t0.elapsed().as_micros() as u64,
                            );
                        }
                        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
                        Err(_) => {
                            mux_profile::record(
                                Stage::IoTcpRead,
                                0,
                                t0.elapsed().as_micros() as u64,
                            );
                            if !c.eof_pushed {
                                let cb_t0 = Instant::now();
                                (c.on_data)(Vec::new());
                                mux_profile::record(
                                    Stage::IoCallback,
                                    0,
                                    cb_t0.elapsed().as_micros() as u64,
                                );
                                c.eof_pushed = true;
                            }
                            to_close.push(tk);
                            break;
                        }
                    }
                }
            }
            if ev.is_writable() && !c.pending_out.is_empty() {
                loop {
                    if c.pending_out.is_empty() {
                        let _ = c.stream.reregister(poll.registry(), tk, Interest::READABLE);
                        break;
                    }
                    let t0 = Instant::now();
                    match c.stream.write(&c.pending_out) {
                        Ok(0) => {
                            mux_profile::record(
                                Stage::IoTcpWrite,
                                0,
                                t0.elapsed().as_micros() as u64,
                            );
                            to_close.push(tk);
                            break;
                        }
                        Ok(n) => {
                            mux_profile::record(
                                Stage::IoTcpWrite,
                                n,
                                t0.elapsed().as_micros() as u64,
                            );
                            c.pending_out.drain(..n);
                        }
                        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
                        Err(_) => {
                            to_close.push(tk);
                            break;
                        }
                    }
                }
            }
        }
        for tk in to_close {
            if let Some(mut c) = conns.remove(&tk) {
                by_port.remove(&c.host_src_port);
                let _ = c.stream.deregister(poll.registry());
            }
        }
    }
}