How an Async Runtime Works in Rust — Building an Executor, Timer, and epoll from Scratch

//! The minimal executor: a run queue of tasks and a loop that polls them.
//!
//! Waker strategy: we implement `std::task::Wake` for `Arc<Task>` instead of
//! hand-rolling a `RawWakerVTable`. The trait generates exactly that vtable
//! for us (clone = `Arc::clone`, wake = `Wake::wake`, drop = `Arc::drop`),
//! so we get the same machinery without any unsafe pointer bookkeeping.
//! Rolling the vtable by hand buys nothing here but bugs.

use anyhow::{Context as _, Result, anyhow};
use std::{
    future::Future,
    pin::Pin,
    sync::{
        Arc, Mutex, PoisonError,
        atomic::{AtomicUsize, Ordering},
        mpsc::{Receiver, Sender, channel},
    },
    task::{Context, Poll, Wake, Waker},
};

type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send>>;

/// One spawned future plus the sender that puts it back on the run queue.
pub struct Task {
    /// The future lives behind a `Mutex` so a wake arriving from another
    /// thread can never observe a poll that is still in progress.
    future: Mutex<Option<BoxFuture>>,
    ready: Sender<Arc<Task>>,
    live: Arc<AtomicUsize>,
}

impl Wake for Task {
    /// The entire waker contract in one line: to wake a task is to put it
    /// back on the queue. Nothing is polled here; the executor thread will
    /// pick the task up on its next loop iteration.
    fn wake(self: Arc<Self>) {
        if self.ready.send(self.clone()).is_err() {
            // The executor is gone, so there is nothing left to schedule
            // onto. A wake aimed at a dead executor is a no-op by contract.
        }
    }
}

impl Task {
    /// Poll the stored future once, with a fresh waker pointing back at us.
    fn poll(self: &Arc<Self>) {
        let mut slot = self.future.lock().unwrap_or_else(PoisonError::into_inner);
        // An empty slot means the future already completed; a stale wake
        // that was queued twice lands here and does nothing.
        let Some(mut future) = slot.take() else {
            return;
        };
        let waker = Waker::from(self.clone());
        let mut context = Context::from_waker(&waker);
        match future.as_mut().poll(&mut context) {
            Poll::Ready(()) => {
                self.live.fetch_sub(1, Ordering::SeqCst);
            }
            Poll::Pending => {
                *slot = Some(future);
            }
        }
    }
}

/// A cloneable handle for putting new tasks on the queue.
#[derive(Clone)]
pub struct Spawner {
    ready: Sender<Arc<Task>>,
    live: Arc<AtomicUsize>,
}

impl Spawner {
    pub fn spawn(&self, future: impl Future<Output = ()> + Send + 'static) -> Result<()> {
        self.live.fetch_add(1, Ordering::SeqCst);
        let task = Arc::new(Task {
            future: Mutex::new(Some(Box::pin(future))),
            ready: self.ready.clone(),
            live: self.live.clone(),
        });
        self.ready
            .send(task)
            .map_err(|_| anyhow!("executor is gone, cannot spawn"))
    }
}

/// The executor: pops tasks off the queue and polls them, one at a time.
pub struct Executor {
    queue: Receiver<Arc<Task>>,
    spawner: Spawner,
    polls: u64,
}

impl Default for Executor {
    fn default() -> Self {
        Self::new()
    }
}

impl Executor {
    pub fn new() -> Self {
        let (sender, receiver) = channel();
        let spawner = Spawner {
            ready: sender,
            live: Arc::new(AtomicUsize::new(0)),
        };
        Self {
            queue: receiver,
            spawner,
            polls: 0,
        }
    }

    pub fn spawner(&self) -> Spawner {
        self.spawner.clone()
    }

    /// Total number of `poll` calls this executor has made.
    pub fn polls(&self) -> u64 {
        self.polls
    }

    /// Run until every spawned task has completed. `recv()` blocks when the
    /// queue is empty, so a runtime with sleeping tasks parks the thread
    /// instead of spinning; an external wake (a timer, the reactor) is what
    /// makes `recv()` return.
    pub fn run_all(&mut self) -> Result<()> {
        while self.spawner.live.load(Ordering::SeqCst) > 0 {
            let task = self
                .queue
                .recv()
                .context("run queue closed with live tasks remaining")?;
            self.polls += 1;
            task.poll();
        }
        Ok(())
    }

    /// Drive the queue until this one future completes, then return its
    /// output. Other spawned tasks make progress while we wait.
    pub fn block_on<F>(&mut self, future: F) -> Result<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let slot = Arc::new(Mutex::new(None));
        let out = slot.clone();
        self.spawner.spawn(async move {
            let value = future.await;
            *out.lock().unwrap_or_else(PoisonError::into_inner) = Some(value);
        })?;
        loop {
            let task = self
                .queue
                .recv()
                .context("run queue closed before the main future finished")?;
            self.polls += 1;
            task.poll();
            let mut guard = slot.lock().unwrap_or_else(PoisonError::into_inner);
            if let Some(value) = guard.take() {
                return Ok(value);
            }
        }
    }
}

//! The timer: one thread, one deadline heap, and `Waker::wake` called from
//! outside the executor. This is the smallest possible proof that wakes
//! cross threads: the executor sleeps inside `recv()` until the timer
//! thread pushes a task back onto the run queue.

use anyhow::{Context as _, Result};
use std::{
    cmp::Ordering as CmpOrdering,
    collections::BinaryHeap,
    future::Future,
    pin::Pin,
    sync::{Arc, Condvar, Mutex, PoisonError},
    task::{Context, Poll, Waker},
    thread::{Builder, JoinHandle},
    time::{Duration, Instant},
};

/// State shared between one `Sleep` future and the timer thread. The single
/// mutex is what makes the handoff race-free: either the poll stores its
/// waker before the deadline fires (the thread will find and wake it), or
/// the deadline fired first (the poll sees `fired` and returns Ready).
struct SleepState {
    fired: bool,
    waker: Option<Waker>,
}

/// A heap entry: min-ordered by deadline (via reversed `Ord`).
struct Entry {
    deadline: Instant,
    state: Arc<Mutex<SleepState>>,
}

impl Eq for Entry {}

impl PartialEq for Entry {
    fn eq(&self, other: &Self) -> bool {
        self.deadline == other.deadline
    }
}

impl Ord for Entry {
    fn cmp(&self, other: &Self) -> CmpOrdering {
        // Reversed: BinaryHeap is a max-heap, we want the nearest deadline
        // on top.
        other.deadline.cmp(&self.deadline)
    }
}

impl PartialOrd for Entry {
    fn partial_cmp(&self, other: &Self) -> Option<CmpOrdering> {
        Some(self.cmp(other))
    }
}

struct Shared {
    heap: Mutex<BinaryHeap<Entry>>,
    cvar: Condvar,
}

/// A cloneable handle to the timer thread.
#[derive(Clone)]
pub struct Timer {
    shared: Arc<Shared>,
    /// Held only so the thread's lifetime is visible in the type; the
    /// thread itself runs for the life of the process.
    _thread: Arc<JoinHandle<()>>,
}

impl Timer {
    /// Start the timer thread. It sleeps on a condvar until the nearest
    /// deadline (or forever when the heap is empty) and never busy-waits.
    pub fn spawn() -> Result<Timer> {
        let shared = Arc::new(Shared {
            heap: Mutex::new(BinaryHeap::new()),
            cvar: Condvar::new(),
        });
        let worker = shared.clone();
        let thread = Builder::new()
            .name("minirt-timer".into())
            .spawn(move || run(&worker))
            .context("failed to spawn the timer thread")?;
        Ok(Timer {
            shared,
            _thread: Arc::new(thread),
        })
    }

    /// A future that completes `duration` from now.
    pub fn sleep(&self, duration: Duration) -> Sleep {
        Sleep {
            timer: self.clone(),
            deadline: Instant::now() + duration,
            state: Arc::new(Mutex::new(SleepState {
                fired: false,
                waker: None,
            })),
            registered: false,
        }
    }

    fn register(&self, deadline: Instant, state: Arc<Mutex<SleepState>>) {
        let mut heap = self
            .shared
            .heap
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        heap.push(Entry { deadline, state });
        // Wake the thread so it can re-check the nearest deadline.
        self.shared.cvar.notify_one();
    }
}

/// The timer thread body: pop due entries, wake their tasks, sleep until
/// the next deadline.
fn run(shared: &Shared) {
    let mut heap = shared.heap.lock().unwrap_or_else(PoisonError::into_inner);
    loop {
        let now = Instant::now();
        while let Some(top) = heap.peek() {
            if top.deadline > now {
                break;
            }
            let Some(entry) = heap.pop() else {
                break;
            };
            let mut state = entry.state.lock().unwrap_or_else(PoisonError::into_inner);
            state.fired = true;
            if let Some(waker) = state.waker.take() {
                // The cross-thread moment: this re-enqueues the task on the
                // executor's run queue from the timer thread.
                waker.wake();
            }
        }
        let wait = heap.peek().map(|top| top.deadline - now);
        heap = match wait {
            Some(timeout) => {
                let (guard, _timed_out) = shared
                    .cvar
                    .wait_timeout(heap, timeout)
                    .unwrap_or_else(PoisonError::into_inner);
                guard
            }
            None => shared
                .cvar
                .wait(heap)
                .unwrap_or_else(PoisonError::into_inner),
        };
    }
}

/// The sleeping future. Two polls in the common case: the first stores the
/// waker and registers the deadline, the second (after the timer thread
/// fired) observes `fired` and returns Ready.
pub struct Sleep {
    timer: Timer,
    deadline: Instant,
    state: Arc<Mutex<SleepState>>,
    registered: bool,
}

impl Future for Sleep {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let this = self.get_mut();
        {
            let mut state = this.state.lock().unwrap_or_else(PoisonError::into_inner);
            if state.fired {
                return Poll::Ready(());
            }
            // Store a fresh waker on every poll: the task may have moved
            // and the old waker could point at a stale generation.
            state.waker = Some(cx.waker().clone());
        }
        if !this.registered {
            this.registered = true;
            this.timer.register(this.deadline, this.state.clone());
        }
        Poll::Pending
    }
}

//! The reactor: one epoll instance and one thread blocked in `epoll_wait`.
//!
//! A future that hits `WouldBlock` parks its `Waker` here, keyed by fd and
//! direction; when the kernel reports readiness, the thread wakes exactly
//! the parked task. We use level-triggered epoll with `EPOLLONESHOT`:
//! level-triggered closes the check-then-park race (arming an interest on
//! an already-ready fd fires immediately), and oneshot means a fired fd
//! stays silent until somebody parks a waker again — no busy loops.
//!
//! Why epoll via `libc` at all: std has no readiness API. The honest
//! std-only alternatives are a thread per connection (which makes the
//! executor pointless) or a polling loop over non-blocking sockets (which
//! burns a core doing nothing). The epoll surface we need is three calls.

use anyhow::{Context as _, Result, ensure};
use libc::{
    EPOLL_CLOEXEC, EPOLL_CTL_ADD, EPOLL_CTL_DEL, EPOLL_CTL_MOD, EPOLLERR, EPOLLHUP, EPOLLIN,
    EPOLLONESHOT, EPOLLOUT, c_int, close, epoll_create1, epoll_ctl, epoll_event, epoll_wait,
};
use std::{
    collections::HashMap,
    io,
    os::fd::RawFd,
    ptr,
    sync::{Arc, Mutex, PoisonError},
    task::Waker,
    thread::{Builder, JoinHandle},
};

/// Which readiness a task is waiting for.
#[derive(Clone, Copy)]
pub enum Direction {
    Read,
    Write,
}

/// Parked wakers for one registered fd.
#[derive(Default)]
struct Slot {
    read: Option<Waker>,
    write: Option<Waker>,
}

fn interest_mask(slot: &Slot) -> u32 {
    let mut mask = 0;
    if slot.read.is_some() {
        mask |= EPOLLIN as u32;
    }
    if slot.write.is_some() {
        mask |= EPOLLOUT as u32;
    }
    mask
}

/// Re-arm the fd with the given interest; oneshot keeps it silent after
/// the next event until somebody arms it again.
fn arm(epoll: RawFd, fd: RawFd, mask: u32) -> Result<()> {
    let mut event = epoll_event {
        events: mask | EPOLLONESHOT as u32,
        u64: fd as u64,
    };
    let code = unsafe { epoll_ctl(epoll, EPOLL_CTL_MOD, fd, &mut event) };
    ensure!(
        code == 0,
        "epoll_ctl(MOD) failed: {}",
        io::Error::last_os_error()
    );
    Ok(())
}

struct Shared {
    epoll: RawFd,
    slots: Mutex<HashMap<RawFd, Slot>>,
}

impl Drop for Shared {
    fn drop(&mut self) {
        // Best effort: the reactor lives for the whole process anyway.
        if unsafe { close(self.epoll) } != 0 {
            eprintln!("reactor: close failed: {}", io::Error::last_os_error());
        }
    }
}

/// A cloneable handle to the reactor thread.
#[derive(Clone)]
pub struct Reactor {
    shared: Arc<Shared>,
    /// Held for lifetime clarity; the thread runs as long as the process.
    _thread: Arc<JoinHandle<()>>,
}

impl Reactor {
    pub fn spawn() -> Result<Reactor> {
        let epoll = unsafe { epoll_create1(EPOLL_CLOEXEC) };
        ensure!(
            epoll >= 0,
            "epoll_create1 failed: {}",
            io::Error::last_os_error()
        );
        let shared = Arc::new(Shared {
            epoll,
            slots: Mutex::new(HashMap::new()),
        });
        let worker = shared.clone();
        let thread = Builder::new()
            .name("minirt-reactor".into())
            .spawn(move || run(&worker))
            .context("failed to spawn the reactor thread")?;
        Ok(Reactor {
            shared,
            _thread: Arc::new(thread),
        })
    }

    /// Add the fd to the epoll set, disarmed: nothing fires until a task
    /// parks a waker via `wait`.
    pub fn register(&self, fd: RawFd) -> Result<()> {
        let mut slots = self
            .shared
            .slots
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        ensure!(!slots.contains_key(&fd), "fd {fd} is already registered");
        let mut event = epoll_event {
            events: EPOLLONESHOT as u32,
            u64: fd as u64,
        };
        let code = unsafe { epoll_ctl(self.shared.epoll, EPOLL_CTL_ADD, fd, &mut event) };
        ensure!(
            code == 0,
            "epoll_ctl(ADD) failed: {}",
            io::Error::last_os_error()
        );
        slots.insert(fd, Slot::default());
        Ok(())
    }

    /// Park a waker until the fd is ready in the given direction. Called
    /// by a leaf future right after an op returned `WouldBlock`.
    pub fn wait(&self, fd: RawFd, direction: Direction, waker: Waker) -> Result<()> {
        let mut slots = self
            .shared
            .slots
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        let slot = slots
            .get_mut(&fd)
            .with_context(|| format!("fd {fd} is not registered with the reactor"))?;
        match direction {
            Direction::Read => slot.read = Some(waker),
            Direction::Write => slot.write = Some(waker),
        }
        arm(self.shared.epoll, fd, interest_mask(slot))
    }

    /// Forget the fd; any still-parked task is woken so it can observe the
    /// closed socket instead of hanging forever.
    pub fn deregister(&self, fd: RawFd) -> Result<()> {
        let mut slots = self
            .shared
            .slots
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        let Some(mut slot) = slots.remove(&fd) else {
            return Ok(());
        };
        let code = unsafe { epoll_ctl(self.shared.epoll, EPOLL_CTL_DEL, fd, ptr::null_mut()) };
        ensure!(
            code == 0,
            "epoll_ctl(DEL) failed: {}",
            io::Error::last_os_error()
        );
        drop(slots);
        if let Some(waker) = slot.read.take() {
            waker.wake();
        }
        if let Some(waker) = slot.write.take() {
            waker.wake();
        }
        Ok(())
    }
}

/// The reactor thread body: wait for kernel events, hand the parked wakers
/// back to the executor, re-arm whatever interest is still parked.
fn run(shared: &Shared) {
    let mut events = [epoll_event { events: 0, u64: 0 }; 64];
    loop {
        let count =
            unsafe { epoll_wait(shared.epoll, events.as_mut_ptr(), events.len() as c_int, -1) };
        if count < 0 {
            let error = io::Error::last_os_error();
            if error.kind() == io::ErrorKind::Interrupted {
                continue;
            }
            eprintln!("reactor: epoll_wait failed: {error}");
            return;
        }
        let mut ready = Vec::new();
        {
            let mut slots = shared.slots.lock().unwrap_or_else(PoisonError::into_inner);
            for event in &events[..count as usize] {
                let bits = event.events;
                let fd = event.u64 as RawFd;
                let Some(slot) = slots.get_mut(&fd) else {
                    continue;
                };
                // Errors and hangups wake both directions: the parked task
                // must retry its op and see the real error itself.
                let fail = bits & (EPOLLERR | EPOLLHUP) as u32 != 0;
                if fail || bits & EPOLLIN as u32 != 0 {
                    if let Some(waker) = slot.read.take() {
                        ready.push(waker);
                    }
                }
                if fail || bits & EPOLLOUT as u32 != 0 {
                    if let Some(waker) = slot.write.take() {
                        ready.push(waker);
                    }
                }
                // Oneshot disarmed the fd; keep it armed for the direction
                // that is still parked (if any).
                let mask = interest_mask(slot);
                if mask != 0
                    && let Err(error) = arm(shared.epoll, fd, mask)
                {
                    eprintln!("reactor: re-arm of fd {fd} failed: {error:#}");
                }
            }
        }
        // Wake outside the lock: wakes go straight to the run queue and
        // must never contend with a poll that is parking a new waker.
        for waker in ready {
            waker.wake();
        }
    }
}

//! Non-blocking TCP on top of the reactor. Every async operation here is
//! the same two-step dance: try the syscall; on `WouldBlock`, park the
//! waker in the reactor and return `Pending`.

use crate::reactor::{Direction, Reactor};
use anyhow::{Context as _, Result, bail, ensure};
use libc::{
    AF_INET, EINPROGRESS, ENOTCONN, SO_ERROR, SOCK_CLOEXEC, SOCK_NONBLOCK, SOCK_STREAM, SOL_SOCKET,
    c_int, c_void, connect, getsockopt, in_addr, sockaddr, sockaddr_in, socket, socklen_t,
};
use std::{
    future::Future,
    io::{self, Read as _, Write as _},
    net::{SocketAddr, TcpListener, TcpStream},
    os::fd::{AsRawFd, FromRawFd, RawFd},
    pin::Pin,
    ptr,
    task::{Context, Poll},
};

/// The universal leaf future: run `op` until it stops saying `WouldBlock`.
/// Parking happens after the failed try; level-triggered epoll makes that
/// order safe (arming an already-ready fd fires immediately).
struct ReadyOp<'a, F> {
    reactor: &'a Reactor,
    fd: RawFd,
    direction: Direction,
    op: F,
}

impl<F, T> Future for ReadyOp<'_, F>
where
    F: FnMut() -> io::Result<T> + Unpin,
{
    type Output = Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<T>> {
        let this = self.get_mut();
        loop {
            match (this.op)() {
                Ok(value) => return Poll::Ready(Ok(value)),
                Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
                    let parked = this
                        .reactor
                        .wait(this.fd, this.direction, cx.waker().clone());
                    return match parked {
                        Ok(()) => Poll::Pending,
                        Err(error) => Poll::Ready(Err(error)),
                    };
                }
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(error) => return Poll::Ready(Err(error.into())),
            }
        }
    }
}

/// An async TCP listener: std's listener, switched to non-blocking and
/// registered with the reactor.
pub struct Listener {
    reactor: Reactor,
    inner: TcpListener,
}

impl Listener {
    pub fn bind(reactor: &Reactor, addr: SocketAddr) -> Result<Listener> {
        let inner = TcpListener::bind(addr).with_context(|| format!("failed to bind {addr}"))?;
        inner.set_nonblocking(true)?;
        reactor.register(inner.as_raw_fd())?;
        Ok(Listener {
            reactor: reactor.clone(),
            inner,
        })
    }

    pub fn local_addr(&self) -> Result<SocketAddr> {
        Ok(self.inner.local_addr()?)
    }

    pub async fn accept(&self) -> Result<(Stream, SocketAddr)> {
        let fd = self.inner.as_raw_fd();
        let inner = &self.inner;
        let (stream, peer) = ReadyOp {
            reactor: &self.reactor,
            fd,
            direction: Direction::Read,
            op: move || inner.accept(),
        }
        .await?;
        Ok((Stream::from_std(&self.reactor, stream)?, peer))
    }
}

impl Drop for Listener {
    fn drop(&mut self) {
        if let Err(error) = self.reactor.deregister(self.inner.as_raw_fd()) {
            eprintln!("listener deregister failed: {error:#}");
        }
    }
}

/// An async TCP stream over the reactor.
pub struct Stream {
    reactor: Reactor,
    inner: TcpStream,
}

impl Stream {
    fn from_std(reactor: &Reactor, inner: TcpStream) -> Result<Stream> {
        inner.set_nonblocking(true)?;
        reactor.register(inner.as_raw_fd())?;
        Ok(Stream {
            reactor: reactor.clone(),
            inner,
        })
    }

    /// Non-blocking connect, the classic sequence: `connect(2)` returns
    /// `EINPROGRESS`, we park for writability, and the final verdict is
    /// read back with `SO_ERROR`.
    pub async fn connect(reactor: &Reactor, addr: SocketAddr) -> Result<Stream> {
        let SocketAddr::V4(v4) = addr else {
            bail!("only IPv4 addresses are supported in this demo");
        };
        let fd = unsafe { socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0) };
        ensure!(fd >= 0, "socket(2) failed: {}", io::Error::last_os_error());
        // From here the fd is owned by TcpStream: every error path below
        // closes it via Drop.
        let inner = unsafe { TcpStream::from_raw_fd(fd) };
        let sin = sockaddr_in {
            sin_family: AF_INET as u16,
            sin_port: v4.port().to_be(),
            sin_addr: in_addr {
                s_addr: u32::from_ne_bytes(v4.ip().octets()),
            },
            sin_zero: [0; 8],
        };
        let code = unsafe {
            connect(
                fd,
                ptr::from_ref(&sin).cast::<sockaddr>(),
                size_of::<sockaddr_in>() as socklen_t,
            )
        };
        if code != 0 {
            let error = io::Error::last_os_error();
            ensure!(
                error.raw_os_error() == Some(EINPROGRESS),
                "connect(2) to {addr} failed: {error}"
            );
        }
        let stream = Stream::from_std(reactor, inner)?;
        let probe = &stream.inner;
        ReadyOp {
            reactor: &stream.reactor,
            fd,
            direction: Direction::Write,
            op: move || {
                // A failed connect surfaces via SO_ERROR, not via connect(2).
                let pending = so_error(fd)?;
                if pending != 0 {
                    return Err(io::Error::from_raw_os_error(pending));
                }
                match probe.peer_addr() {
                    Ok(_) => Ok(()),
                    Err(error) if error.raw_os_error() == Some(ENOTCONN) => {
                        Err(io::ErrorKind::WouldBlock.into())
                    }
                    Err(error) => Err(error),
                }
            },
        }
        .await
        .with_context(|| format!("connect to {addr} did not complete"))?;
        Ok(stream)
    }

    pub async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let fd = self.inner.as_raw_fd();
        let mut sock = &self.inner;
        ReadyOp {
            reactor: &self.reactor,
            fd,
            direction: Direction::Read,
            op: move || sock.read(buf),
        }
        .await
    }

    /// Read until `buf` is full; a peer that closes mid-message is an error.
    pub async fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        let mut filled = 0;
        while filled < buf.len() {
            let count = self.read(&mut buf[filled..]).await?;
            ensure!(count > 0, "connection closed mid-message");
            filled += count;
        }
        Ok(())
    }

    /// Write the whole buffer, parking on a full send queue.
    pub async fn write_all(&mut self, data: &[u8]) -> Result<()> {
        let fd = self.inner.as_raw_fd();
        let mut offset = 0;
        while offset < data.len() {
            let rest = &data[offset..];
            let mut sock = &self.inner;
            let count = ReadyOp {
                reactor: &self.reactor,
                fd,
                direction: Direction::Write,
                op: move || sock.write(rest),
            }
            .await?;
            ensure!(count > 0, "socket accepted zero bytes");
            offset += count;
        }
        Ok(())
    }
}

impl Drop for Stream {
    fn drop(&mut self) {
        if let Err(error) = self.reactor.deregister(self.inner.as_raw_fd()) {
            eprintln!("stream deregister failed: {error:#}");
        }
    }
}

fn so_error(fd: RawFd) -> io::Result<i32> {
    let mut value: c_int = 0;
    let mut len = size_of::<c_int>() as socklen_t;
    let code = unsafe {
        getsockopt(
            fd,
            SOL_SOCKET,
            SO_ERROR,
            ptr::from_mut(&mut value).cast::<c_void>(),
            &mut len,
        )
    };
    if code != 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(value)
}

//! TCP echo server on minirt: an accept loop, one task per connection and
//! a stats heartbeat on the timer — all interleaved on a single executor
//! thread, with wakes arriving from the epoll and timer threads.

use anyhow::Result;
use clap::Parser;
use minirt::{
    executor::{Executor, Spawner},
    net::{Listener, Stream},
    reactor::Reactor,
    timer::Timer,
};
use std::{
    net::SocketAddr,
    sync::{
        Arc,
        atomic::{AtomicU64, AtomicUsize, Ordering},
    },
    time::Duration,
};

/// TCP echo server running on the minirt executor.
#[derive(Parser)]
struct Args {
    /// Address to listen on.
    #[arg(long, default_value = "127.0.0.1:7000")]
    addr: SocketAddr,
}

#[derive(Default)]
struct Stats {
    active: AtomicUsize,
    peak: AtomicUsize,
    served: AtomicUsize,
    bytes: AtomicU64,
}

/// One connection: read whatever arrives, write it straight back, stop at
/// EOF. The 4 KiB buffer lives inside the task's future, not on a thread
/// stack — a thousand connections cost a thousand futures, not a thousand
/// stacks.
async fn handle(mut stream: Stream, stats: Arc<Stats>) -> Result<()> {
    let mut buf = [0u8; 4096];
    loop {
        let count = stream.read(&mut buf).await?;
        if count == 0 {
            return Ok(());
        }
        stream.write_all(&buf[..count]).await?;
        stats.bytes.fetch_add(count as u64, Ordering::Relaxed);
    }
}

/// The accept loop: every new connection becomes one more task on the
/// same executor.
async fn serve(listener: Listener, spawner: Spawner, stats: Arc<Stats>) -> Result<()> {
    loop {
        let (stream, peer) = listener.accept().await?;
        let stats = stats.clone();
        stats.active.fetch_add(1, Ordering::SeqCst);
        stats
            .peak
            .fetch_max(stats.active.load(Ordering::SeqCst), Ordering::SeqCst);
        spawner.spawn(async move {
            let outcome = handle(stream, stats.clone()).await;
            stats.active.fetch_sub(1, Ordering::SeqCst);
            stats.served.fetch_add(1, Ordering::SeqCst);
            if let Err(error) = outcome {
                eprintln!("[{peer}] error: {error:#}");
            }
        })?;
    }
}

fn main() -> Result<()> {
    let args = Args::parse();
    let mut executor = Executor::new();
    let spawner = executor.spawner();
    let reactor = Reactor::spawn()?;
    let timer = Timer::spawn()?;
    let listener = Listener::bind(&reactor, args.addr)?;
    println!("echo listening on {}", listener.local_addr()?);

    // Heartbeat: once a second print the counters, but only when they
    // changed — an idle server stays silent.
    let stats = Arc::new(Stats::default());
    let watched = stats.clone();
    spawner.spawn(async move {
        let mut last = (0, 0, 0, 0);
        loop {
            timer.sleep(Duration::from_secs(1)).await;
            let snapshot = (
                watched.active.load(Ordering::SeqCst),
                watched.peak.load(Ordering::SeqCst),
                watched.served.load(Ordering::SeqCst),
                watched.bytes.load(Ordering::SeqCst),
            );
            if snapshot != last {
                last = snapshot;
                let (active, peak, served, bytes) = snapshot;
                println!("conns {active} | peak {peak} | served {served} | echoed {bytes} B");
            }
        }
    })?;

    executor.block_on(serve(listener, spawner.clone(), stats))??;
    Ok(())
}

//! Concurrency probe for the echo server, itself running on minirt: open
//! all connections first (so the server really holds them at once), then
//! run echo rounds on every connection and verify each byte.

use anyhow::{Result, ensure};
use clap::Parser;
use minirt::{executor::Executor, net::Stream, reactor::Reactor, timer::Timer};
use std::{
    net::SocketAddr,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::{Duration, Instant},
};

/// Echo probe: N concurrent connections, M verified rounds each.
#[derive(Parser)]
struct Args {
    /// Echo server address.
    #[arg(long, default_value = "127.0.0.1:7000")]
    addr: SocketAddr,
    /// Number of concurrent connections.
    #[arg(long, default_value_t = 1000)]
    conns: usize,
    /// Echo rounds per connection.
    #[arg(long, default_value_t = 10)]
    rounds: usize,
    /// Payload size per round, bytes.
    #[arg(long, default_value_t = 64)]
    payload: usize,
}

/// Everything one worker task needs.
struct Job {
    id: usize,
    addr: SocketAddr,
    rounds: usize,
    payload: usize,
    total: usize,
    reactor: Reactor,
    timer: Timer,
    connected: Arc<AtomicUsize>,
    verified: Arc<AtomicUsize>,
}

async fn worker(job: Job) -> Result<()> {
    let mut stream = Stream::connect(&job.reactor, job.addr).await?;
    job.connected.fetch_add(1, Ordering::SeqCst);
    // Poor man's barrier: keep the socket open until every worker is
    // connected, so peak concurrency on the server really reaches N.
    while job.connected.load(Ordering::SeqCst) < job.total {
        job.timer.sleep(Duration::from_millis(5)).await;
    }
    let message: Vec<u8> = (0..job.payload)
        .map(|i| ((job.id + i) % 256) as u8)
        .collect();
    let mut echo = vec![0u8; message.len()];
    for round in 0..job.rounds {
        stream.write_all(&message).await?;
        stream.read_exact(&mut echo).await?;
        ensure!(echo == message, "echo mismatch on round {round}");
        job.verified.fetch_add(1, Ordering::SeqCst);
    }
    Ok(())
}

fn main() -> Result<()> {
    let args = Args::parse();
    let mut executor = Executor::new();
    let spawner = executor.spawner();
    let reactor = Reactor::spawn()?;
    let timer = Timer::spawn()?;
    let connected = Arc::new(AtomicUsize::new(0));
    let verified = Arc::new(AtomicUsize::new(0));
    let failed = Arc::new(AtomicUsize::new(0));
    let start = Instant::now();
    for id in 0..args.conns {
        let job = Job {
            id,
            addr: args.addr,
            rounds: args.rounds,
            payload: args.payload,
            total: args.conns,
            reactor: reactor.clone(),
            timer: timer.clone(),
            connected: connected.clone(),
            verified: verified.clone(),
        };
        let failures = failed.clone();
        spawner.spawn(async move {
            let id = job.id;
            if let Err(error) = worker(job).await {
                failures.fetch_add(1, Ordering::SeqCst);
                eprintln!("[conn {id}] error: {error:#}");
            }
        })?;
    }
    executor.run_all()?;
    println!(
        "{} conns, {} rounds x {} B: {} echoes verified, {} failed, {} ms wall, {} polls",
        args.conns,
        args.rounds,
        args.payload,
        verified.load(Ordering::SeqCst),
        failed.load(Ordering::SeqCst),
        start.elapsed().as_millis(),
        executor.polls()
    );
    ensure!(
        failed.load(Ordering::SeqCst) == 0,
        "{} connections failed",
        failed.load(Ordering::SeqCst)
    );
    Ok(())
}