Skip to main content

simoom_sys/driver/
mod.rs

1//! A completion-style I/O driver, one per event-loop thread.
2//!
3//! The model is the one io_uring and IOCP provide natively, and that epoll and kqueue emulate
4//! inside [`Driver::turn`] with readiness events plus non-blocking calls:
5//!
6//! - **Receive** into buffers the driver owns. The caller hands buffers to the driver with
7//!   [`Driver::provide`]; the driver takes one only when data has actually arrived (io_uring's
8//!   provided-buffer ring, a read after readiness on epoll and kqueue, a read after a zero-byte
9//!   `WSARecv` on IOCP). So **an idle connection holds no buffer**. Data comes back as
10//!   [`Event::Data`] carrying the filled `Vec<u8>`; the caller owns it from then on.
11//! - **Send** owned buffers with [`Driver::send`]. Sends on one connection go out in order and
12//!   are batched: everything queued during a turn is written at the start of the next
13//!   [`Driver::turn`], with one vectored call (or one submission) per connection. Each buffer
14//!   comes back as [`Event::Sent`] once written, so it can be reused.
15//! - **Accept** on listeners registered with [`Driver::listen`]: [`Event::Accepted`] carries the
16//!   new non-blocking socket, which the caller registers with [`Driver::stream`] (possibly on
17//!   another loop's driver).
18//! - **Close** with [`Driver::close`]: the driver cancels what is in flight, returns every
19//!   buffer, closes the socket and reports [`Event::Closed`]; only then may the caller reuse the
20//!   key.
21//! - **Wake** a driver blocked in [`Driver::turn`] from any thread with a [`Waker`].
22//!
23//! The driver never calls back into user code: [`Driver::turn`] appends [`Event`]s to a vector.
24
25use crate::net::Socket;
26use crate::udp::{RecvMeta, Transmit};
27use std::fmt;
28use std::io;
29use std::net::SocketAddr;
30use std::sync::atomic::{AtomicBool, Ordering};
31use std::sync::Arc;
32use std::time::Duration;
33
34#[cfg(any(target_os = "linux", target_os = "macos", target_os = "ios", target_os = "freebsd"))]
35mod readiness;
36
37#[cfg(target_os = "linux")]
38mod epoll;
39#[cfg(all(target_os = "linux", feature = "io-uring"))]
40pub mod uring;
41
42#[cfg(windows)]
43mod iocp;
44
45#[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
46mod kqueue;
47
48/// Which I/O back end a driver uses.
49#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
50pub enum Backend {
51    /// Pick for this OS: epoll on Linux (io_uring when `SIMOOM_BACKEND=io_uring` is set), IOCP
52    /// on Windows, kqueue on macOS and BSD.
53    Auto,
54    /// Linux epoll, edge-triggered.
55    Epoll,
56    /// Linux io_uring (6.1+). Falls back to epoll when the kernel refuses it.
57    IoUring,
58    /// Windows I/O completion ports.
59    Iocp,
60    /// macOS and BSD kqueue.
61    Kqueue,
62}
63
64impl Backend {
65    /// The back end `Auto` resolves to on this OS, honouring `SIMOOM_BACKEND`
66    /// (`epoll`, `io_uring`/`uring`, `iocp`, `kqueue`).
67    pub fn resolve(self) -> Backend {
68        if self != Backend::Auto {
69            return self;
70        }
71        match env_bytes("SIMOOM_BACKEND").as_deref() {
72            Some(b"epoll") => return Backend::Epoll,
73            Some(b"io_uring") | Some(b"uring") => return Backend::IoUring,
74            Some(b"iocp") => return Backend::Iocp,
75            Some(b"kqueue") => return Backend::Kqueue,
76            _ => {}
77        }
78        if cfg!(windows) {
79            Backend::Iocp
80        } else if cfg!(target_os = "linux") {
81            Backend::Epoll
82        } else {
83            Backend::Kqueue
84        }
85    }
86
87    /// A short lower-case name.
88    pub fn name(self) -> &'static str {
89        match self {
90            Backend::Auto => "auto",
91            Backend::Epoll => "epoll",
92            Backend::IoUring => "io_uring",
93            Backend::Iocp => "iocp",
94            Backend::Kqueue => "kqueue",
95        }
96    }
97}
98
99/// Driver settings. The defaults suit a server loop.
100#[derive(Clone, Debug)]
101pub struct DriverConfig {
102    /// Which back end to use.
103    pub backend: Backend,
104    /// Most OS events taken per wait (epoll and kqueue events, IOCP entries). The array starts
105    /// at 64 entries and doubles, up to this, whenever a wait fills it, so a busy loop drains
106    /// every ready event in one call and an idle one keeps a small array.
107    pub events: usize,
108    /// Most receives per connection per turn before others get a go (epoll, kqueue, IOCP).
109    pub recv_budget: usize,
110    /// io_uring: submission queue entries (the completion queue is four times larger).
111    pub uring_entries: u32,
112    /// io_uring: entries of the provided-buffer ring (a power of two, at most 32768).
113    pub uring_buffers: u16,
114    /// io_uring: use registered (fixed) file descriptors for connections.
115    pub uring_fixed_files: bool,
116    /// io_uring: size of the fixed file table (connections per loop that can use it; others
117    /// use plain descriptors).
118    pub uring_max_files: u32,
119    /// io_uring: pretend `io_uring_setup` failed with this errno, to test the fallback. The
120    /// environment variable `SIMOOM_IO_URING_FAIL=<errno>` does the same.
121    pub uring_simulate_failure: Option<i32>,
122    /// IOCP: post receives with a buffer instead of the zero-byte idiom: one system call per
123    /// receive instead of two, but every connection waiting for data holds a receive buffer.
124    /// Off by default.
125    pub iocp_eager_recv: bool,
126    /// IOCP: `AcceptEx` calls kept outstanding per listener.
127    pub accepts_outstanding: usize,
128    /// UDP: receives kept outstanding per socket on IOCP.
129    pub udp_recvs_outstanding: usize,
130    /// Adaptive busy polling: after a turn that delivered events, the driver may poll without
131    /// blocking for up to this long before it sleeps, saving a sleep and a wake-up when the next
132    /// event comes within a few microseconds. It spins only while events have been arriving
133    /// sooner than a quarter of this window after busy turns (measured by an occasional probe
134    /// spin), so at moderate load it sleeps exactly as without busy polling and costs no CPU per
135    /// request; an idle loop never spins. `Duration::ZERO` turns it off. The environment variable
136    /// `SIMOOM_BUSY_POLL_US` overrides it. Default 30 µs.
137    pub busy_poll: Duration,
138}
139
140impl Default for DriverConfig {
141    fn default() -> Self {
142        DriverConfig {
143            backend: Backend::Auto,
144            events: 1024,
145            recv_budget: 16,
146            uring_entries: 1024,
147            uring_buffers: 1024,
148            uring_fixed_files: true,
149            uring_max_files: 16384,
150            uring_simulate_failure: None,
151            iocp_eager_recv: false,
152            accepts_outstanding: 32,
153            udp_recvs_outstanding: 4,
154            busy_poll: Duration::from_micros(30),
155        }
156    }
157}
158
159/// A handle to a socket registered with a [`Driver`]. Stale handles (after the socket was
160/// closed) are detected by a generation count and ignored.
161#[derive(Clone, Copy, PartialEq, Eq, Hash)]
162pub struct Id(pub(crate) u64);
163
164impl Id {
165    pub(crate) fn new(idx: u32, gen: u32) -> Id {
166        Id(idx as u64 | (gen as u64) << 32)
167    }
168    pub(crate) fn idx(self) -> u32 {
169        self.0 as u32
170    }
171    pub(crate) fn gen(self) -> u32 {
172        (self.0 >> 32) as u32
173    }
174    /// The raw value (for logging).
175    pub fn raw(self) -> u64 {
176        self.0
177    }
178}
179
180impl fmt::Debug for Id {
181    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
182        write!(f, "Id({}#{})", self.idx(), self.gen())
183    }
184}
185
186/// A buffer to send: owned (returned by [`Event::Sent`] for reuse) or static.
187#[derive(Debug)]
188pub enum OutBuf {
189    /// An owned buffer; all of `len()` bytes are sent.
190    Owned(Vec<u8>),
191    /// Static bytes, such as a fixed reply.
192    Static(&'static [u8]),
193}
194
195impl OutBuf {
196    /// The bytes to send.
197    #[inline]
198    pub fn as_slice(&self) -> &[u8] {
199        match self {
200            OutBuf::Owned(v) => v,
201            OutBuf::Static(s) => s,
202        }
203    }
204    /// Length in bytes.
205    #[inline]
206    pub fn len(&self) -> usize {
207        self.as_slice().len()
208    }
209    /// Whether it is empty.
210    #[inline]
211    pub fn is_empty(&self) -> bool {
212        self.len() == 0
213    }
214}
215
216/// Something that happened, reported by [`Driver::turn`].
217#[derive(Debug)]
218pub enum Event {
219    /// A listener accepted a connection. The socket is non-blocking and not yet registered.
220    Accepted { key: u64, socket: Socket, peer: Option<SocketAddr> },
221    /// Bytes arrived on a stream: `buf[..]` holds them (`buf.len()` > 0).
222    Data { key: u64, buf: Vec<u8> },
223    /// The peer closed its sending side. Receiving stops; sending may continue.
224    Eof { key: u64 },
225    /// A receive, send or accept failed. The connection should be closed.
226    Error { key: u64, error: io::Error },
227    /// A buffer was written completely (or discarded because the connection failed or closed).
228    Sent { key: u64, buf: OutBuf },
229    /// A datagram (or a GRO batch) arrived on a UDP socket.
230    Datagram { key: u64, buf: Vec<u8>, meta: RecvMeta },
231    /// A socket given to [`Driver::close`] is closed and the driver holds nothing of it.
232    Closed { key: u64 },
233    /// A [`Waker`] was triggered.
234    Wake,
235}
236
237/// System calls and events counted by a driver (plain counters; read them on the loop thread).
238#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
239pub struct DriverStats {
240    /// Waits: `epoll_wait`, `kevent`, `io_uring_enter`, `GetQueuedCompletionStatusEx`.
241    pub waits: u64,
242    /// Receive calls: `recv`, `recvmmsg`, `WSARecv`/`WSARecvMsg` (io_uring: none, completions only).
243    pub recv_calls: u64,
244    /// Send calls: `sendmsg`/`send`, `WSASend` (io_uring: none, submissions only).
245    pub send_calls: u64,
246    /// Accept calls: `accept4`, `AcceptEx`.
247    pub accept_calls: u64,
248    /// Other system calls: registration, `epoll_ctl`, eventfd reads, `io_uring_register`, ...
249    pub other_calls: u64,
250    /// io_uring submission queue entries pushed (buffer hand-overs with `PROVIDE_BUFFERS` not
251    /// counted).
252    pub sqes: u64,
253    /// Completions or readiness events taken from the OS.
254    pub os_events: u64,
255    /// Waits that blocked (put the thread to sleep, or could have: a non-zero timeout).
256    pub sleeps: u64,
257    /// Non-blocking polls made while busy polling (also counted in `waits` when they were system
258    /// calls; io_uring polls its rings in user space and enters only when work is flagged).
259    pub spin_polls: u64,
260}
261
262impl DriverStats {
263    /// All counted system calls.
264    pub fn syscalls(&self) -> u64 {
265        self.waits + self.recv_calls + self.send_calls + self.accept_calls + self.other_calls
266    }
267}
268
269// ---------------------------------------------------------------------------------------------
270// Adaptive busy polling
271// ---------------------------------------------------------------------------------------------
272
273/// Entries of an event array before it grows (see [`DriverConfig::events`]).
274pub(crate) const EVENTS_START: usize = 64;
275
276/// Most turns in a row that report results without asking the OS for new events.
277pub(crate) const MAX_SKIPS: u8 = 3;
278
279/// Busy waits between probe spins while the loop sleeps (see [`Spin`]).
280const PROBE_EVERY: u8 = 64;
281
282/// Busy polling of one driver (see [`DriverConfig::busy_poll`]).
283///
284/// Spinning costs the CPU time until the next event; sleeping costs a sleep and a wake-up (a
285/// few microseconds of kernel and scheduler work, more in a VM). So the driver spins only while
286/// the next event comes sooner than about one sleep and wake-up: it keeps `gap`, a smoothed time
287/// from the end of a busy turn to the next event, and spins while it is under a quarter of the
288/// window. The gap is measured only by spinning (a sleep's length includes the wake-up delay, so
289/// it would overstate it): while the loop sleeps, every [`PROBE_EVERY`]th busy wait spins once
290/// for the whole window to sample it. A spin that finds nothing counts as a long gap, so when
291/// load drops the loop goes back to sleeping at once. At moderate load (tens of microseconds
292/// between requests per loop) it sleeps, as without busy polling; on chatty connections with
293/// short round trips and near saturation it spins and saves the wake-ups. An idle loop never
294/// spins.
295pub(crate) struct Spin {
296    max: u64,
297    /// Smoothed time to the next event after a busy turn, in nanoseconds (spins only).
298    gap: u64,
299    /// Busy waits since the last probe spin.
300    probe: u8,
301    /// The last turn delivered events: the loop is busy, so spinning may be worth it.
302    pub(crate) worked: bool,
303}
304
305/// An environment variable's bytes (`std::env::var` and its UTF-8 check cost more code).
306#[inline(never)]
307pub(crate) fn env_bytes(name: &str) -> Option<Vec<u8>> {
308    std::env::var_os(name).map(|v| v.into_encoded_bytes())
309}
310
311/// An environment variable holding a decimal number; `None` if unset or not a number.
312pub(crate) fn env_u64(name: &str) -> Option<u64> {
313    let v = env_bytes(name)?;
314    if v.is_empty() || v.len() > 19 || !v.iter().all(u8::is_ascii_digit) {
315        return None;
316    }
317    Some(v.iter().fold(0, |n, &d| n * 10 + u64::from(d - b'0')))
318}
319
320impl Spin {
321    pub(crate) fn new(cfg: &DriverConfig) -> Spin {
322        let max = env_u64("SIMOOM_BUSY_POLL_US").map_or(cfg.busy_poll.as_nanos().min(u64::MAX as u128) as u64, |us| us.saturating_mul(1000));
323        Spin { max, gap: max, probe: 0, worked: false }
324    }
325
326    /// Whether events come soon enough after busy turns to spin for them.
327    fn spins(&self) -> bool {
328        self.worked && self.max > 0 && self.gap <= self.max / 4
329    }
330
331    /// Folds one observed time-to-event into the smoothed gap (weight 1/4).
332    fn record(&mut self, ns: u64) {
333        let ns = ns.min(4 * self.max);
334        self.gap = self.gap - self.gap / 4 + ns / 4;
335    }
336
337    /// Waits for events with `poll`, which waits up to the given timeout (zero: does not block)
338    /// and returns whether it found anything. Spins first when [`Spin::spins`] or a probe is
339    /// due; `syscall_polls` paces polls that are system calls to one a microsecond.
340    ///
341    /// `w` is the driver's waker: work [`Waker::post`]ed meanwhile ends a spin, and the driver
342    /// announces that it parks in the OS before it blocks, so that a post wakes it only then.
343    // One copy for every back end (K9): `poll` is a trait object, its cost is nothing next to
344    // the system call it makes.
345    #[inline(never)]
346    pub(crate) fn wait(&mut self, timeout: Option<Duration>, syscall_polls: bool, stats: &mut DriverStats, w: &WakerInner, poll: &mut dyn FnMut(Option<Duration>, &mut DriverStats) -> io::Result<bool>) -> io::Result<()> {
347        if timeout == Some(Duration::ZERO) {
348            return poll(timeout, stats).map(drop);
349        }
350        let mut spin = self.spins();
351        if !spin && self.worked && self.max > 0 {
352            self.probe += 1;
353            if self.probe >= PROBE_EVERY {
354                self.probe = 0;
355                spin = true;
356            }
357        }
358        let mut spun = 0;
359        if spin {
360            let limit = timeout.map_or(u64::MAX, |t| t.as_nanos().min(u64::MAX as u128) as u64);
361            let window = self.max.min(limit);
362            let t0 = crate::time::now_ns();
363            let mut last = 0;
364            loop {
365                let now = crate::time::now_ns() - t0;
366                if w.posted.load(Ordering::Relaxed) {
367                    // A message for the loop (no system call on either side). Not counted as
368                    // a short gap: replies from worker threads must not keep the loop
369                    // spinning, or it takes the CPU those threads need (K3).
370                    return Ok(());
371                }
372                if !syscall_polls || now - last >= 1_000 || last == 0 {
373                    last = now.max(1);
374                    stats.spin_polls += 1;
375                    if poll(Some(Duration::ZERO), stats)? {
376                        self.record(crate::time::now_ns() - t0);
377                        return Ok(());
378                    }
379                }
380                if now >= window {
381                    spun = now;
382                    break;
383                }
384                core::hint::spin_loop();
385            }
386            if window == self.max {
387                self.record(u64::MAX); // nothing within the whole window
388            }
389        }
390        let rest = timeout.map(|t| t.saturating_sub(Duration::from_nanos(spun)));
391        // Announce the park, then look for posted work once more: with both sides sequentially
392        // consistent, either this sees the post or the poster sees the park and signals.
393        w.parked.store(true, Ordering::SeqCst);
394        if w.posted.load(Ordering::SeqCst) {
395            w.parked.store(false, Ordering::Relaxed);
396            return poll(Some(Duration::ZERO), stats).map(drop);
397        }
398        stats.sleeps += 1;
399        let r = poll(rest, stats);
400        w.parked.store(false, Ordering::Relaxed);
401        r.map(drop)
402    }
403}
404
405// ---------------------------------------------------------------------------------------------
406// Waker
407// ---------------------------------------------------------------------------------------------
408
409/// Wakes a driver blocked in [`Driver::turn`] from any thread. Cheap to clone; repeated wakes
410/// before the driver runs collapse into one OS signal.
411#[derive(Clone)]
412pub struct Waker {
413    pub(crate) inner: Arc<WakerInner>,
414}
415
416pub(crate) struct WakerInner {
417    pub(crate) pending: AtomicBool,
418    pub(crate) signal: Box<dyn Fn() + Send + Sync>,
419    /// Work was [`Waker::post`]ed since the driver last reported it.
420    pub(crate) posted: AtomicBool,
421    /// The driver is blocked in the OS, or about to be (set by [`Spin::wait`]).
422    pub(crate) parked: AtomicBool,
423}
424
425impl Waker {
426    pub(crate) fn new(signal: Box<dyn Fn() + Send + Sync>) -> Waker {
427        Waker { inner: Arc::new(WakerInner { pending: AtomicBool::new(false), signal, posted: AtomicBool::new(false), parked: AtomicBool::new(false) }) }
428    }
429
430    /// Makes the driver's current or next [`Driver::turn`] return [`Event::Wake`].
431    pub fn wake(&self) {
432        if !self.inner.pending.swap(true, Ordering::AcqRel) {
433            (self.inner.signal)();
434        }
435    }
436
437    /// Like [`Waker::wake`], for work the caller has just queued for the driver's thread (with
438    /// a sequentially consistent push, such as [`crate::mpsc::Mpsc::push`]): the current or next
439    /// [`Driver::turn`] returns [`Event::Wake`], but the OS is signalled only when the driver is
440    /// parked in it. A driver that is running, or spinning in its busy poll, picks the work up
441    /// without a system call on either side, and posts that arrive together cost one signal.
442    pub fn post(&self) {
443        let i = &*self.inner;
444        // Sequentially consistent against the driver's park announcement and re-check in
445        // `Spin::wait`: either the driver sees `posted`, or this sees `parked` and signals
446        // (`wake` signals once however many posters see it).
447        i.posted.store(true, Ordering::SeqCst);
448        if i.parked.load(Ordering::SeqCst) {
449            self.wake();
450        }
451    }
452
453    /// Whether work was posted since the driver last reported it (and clears that).
454    pub(crate) fn take_posted(&self) -> bool {
455        self.inner.posted.load(Ordering::Relaxed) && self.inner.posted.swap(false, Ordering::AcqRel)
456    }
457
458    /// Called by the driver when it reports the wake, before the caller looks at its queues.
459    pub(crate) fn clear(&self) {
460        self.inner.pending.store(false, Ordering::Release);
461    }
462}
463
464impl fmt::Debug for Waker {
465    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
466        f.debug_struct("Waker").field("pending", &self.inner.pending.load(Ordering::Relaxed)).finish()
467    }
468}
469
470// ---------------------------------------------------------------------------------------------
471// Slab with generations
472// ---------------------------------------------------------------------------------------------
473
474pub(crate) struct Slab<T> {
475    slots: Vec<(u32, Option<T>)>,
476    free: Vec<u32>,
477    len: usize,
478}
479
480impl<T> Slab<T> {
481    pub(crate) fn new() -> Slab<T> {
482        Slab { slots: Vec::new(), free: Vec::new(), len: 0 }
483    }
484    pub(crate) fn insert(&mut self, v: T) -> Id {
485        self.len += 1;
486        if let Some(i) = self.free.pop() {
487            let s = &mut self.slots[i as usize];
488            s.1 = Some(v);
489            Id::new(i, s.0)
490        } else {
491            self.slots.push((0, Some(v)));
492            Id::new(self.slots.len() as u32 - 1, 0)
493        }
494    }
495    pub(crate) fn get(&self, id: Id) -> Option<&T> {
496        match self.slots.get(id.idx() as usize) {
497            Some((g, Some(v))) if *g == id.gen() => Some(v),
498            _ => None,
499        }
500    }
501    pub(crate) fn get_mut(&mut self, id: Id) -> Option<&mut T> {
502        match self.slots.get_mut(id.idx() as usize) {
503            Some((g, Some(v))) if *g == id.gen() => Some(v),
504            _ => None,
505        }
506    }
507    #[cfg_attr(not(all(target_os = "linux", feature = "io-uring")), allow(dead_code))]
508    pub(crate) fn id_at(&self, idx: u32) -> Option<Id> {
509        match self.slots.get(idx as usize) {
510            Some((g, Some(_))) => Some(Id::new(idx, *g)),
511            _ => None,
512        }
513    }
514    pub(crate) fn remove(&mut self, id: Id) -> Option<T> {
515        let s = self.slots.get_mut(id.idx() as usize)?;
516        if s.0 != id.gen() || s.1.is_none() {
517            return None;
518        }
519        s.0 = s.0.wrapping_add(1);
520        self.free.push(id.idx());
521        self.len -= 1;
522        s.1.take()
523    }
524    pub(crate) fn len(&self) -> usize {
525        self.len
526    }
527    #[cfg_attr(not(windows), allow(dead_code))]
528    pub(crate) fn ids(&self) -> Vec<Id> {
529        self.slots.iter().enumerate().filter(|(_, s)| s.1.is_some()).map(|(i, s)| Id::new(i as u32, s.0)).collect()
530    }
531}
532
533// ---------------------------------------------------------------------------------------------
534// Large UDP receive buffers (GRO)
535// ---------------------------------------------------------------------------------------------
536
537/// Receive buffers per batched receive call on a GRO socket. Each holds up to 64 datagrams of
538/// one peer, so a batch takes up to 512 datagrams (and 512 KiB of buffers).
539pub(crate) const LARGE_BATCH: usize = 8;
540
541/// The driver's own stock of large buffers for UDP sockets registered with [`Driver::udp_large`]:
542/// allocated on first use, given back through [`Driver::provide`], at most two batches kept.
543#[derive(Default)]
544pub(crate) struct LargeBufs {
545    pub(crate) size: usize,
546    free: Vec<Vec<u8>>,
547}
548
549impl LargeBufs {
550    /// A cleared buffer of at least `size` bytes.
551    #[inline(never)]
552    pub(crate) fn take(&mut self) -> Vec<u8> {
553        self.free.pop().unwrap_or_else(|| Vec::with_capacity(self.size))
554    }
555
556    /// Takes `buf` back if it is a large buffer (`None`), else hands it back. The check is
557    /// inlined; keeping the buffer (GRO sockets only) is not.
558    #[inline]
559    pub(crate) fn keep(&mut self, buf: Vec<u8>) -> Option<Vec<u8>> {
560        if !cfg!(feature = "udp") || self.size == 0 || buf.capacity() < self.size {
561            return Some(buf);
562        }
563        self.store(buf);
564        None
565    }
566
567    #[inline(never)]
568    fn store(&mut self, mut buf: Vec<u8>) {
569        if self.free.len() < 2 * LARGE_BATCH {
570            buf.clear();
571            self.free.push(buf);
572        }
573    }
574}
575
576// ---------------------------------------------------------------------------------------------
577// The driver
578// ---------------------------------------------------------------------------------------------
579
580enum Imp {
581    #[cfg(target_os = "linux")]
582    Epoll(readiness::Driver<epoll::Poller>),
583    #[cfg(all(target_os = "linux", feature = "io-uring"))]
584    Uring(Box<uring::Driver>),
585    #[cfg(windows)]
586    Iocp(iocp::Driver),
587    #[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
588    Kqueue(readiness::Driver<kqueue::Poller>),
589}
590
591macro_rules! dispatch {
592    ($self:expr, $d:ident => $e:expr) => {
593        match &mut $self.imp {
594            #[cfg(target_os = "linux")]
595            Imp::Epoll($d) => $e,
596            #[cfg(all(target_os = "linux", feature = "io-uring"))]
597            Imp::Uring($d) => $e,
598            #[cfg(windows)]
599            Imp::Iocp($d) => $e,
600            #[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
601            Imp::Kqueue($d) => $e,
602        }
603    };
604}
605
606macro_rules! dispatch_ref {
607    ($self:expr, $d:ident => $e:expr) => {
608        match &$self.imp {
609            #[cfg(target_os = "linux")]
610            Imp::Epoll($d) => $e,
611            #[cfg(all(target_os = "linux", feature = "io-uring"))]
612            Imp::Uring($d) => $e,
613            #[cfg(windows)]
614            Imp::Iocp($d) => $e,
615            #[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
616            Imp::Kqueue($d) => $e,
617        }
618    };
619}
620
621/// A completion-style I/O driver for one thread. See the [module docs](self).
622///
623/// A driver is not `Send`: create it on the thread that runs it (io_uring's single-issuer mode
624/// requires that). Its [`Waker`] is `Send + Sync`.
625pub struct Driver {
626    imp: Imp,
627    kind: Backend,
628    fallback: Option<String>,
629    _not_send: core::marker::PhantomData<*const ()>,
630}
631
632impl fmt::Debug for Driver {
633    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
634        f.debug_struct("Driver").field("backend", &self.kind).field("fallback", &self.fallback).finish()
635    }
636}
637
638impl Driver {
639    /// Creates a driver. With [`Backend::IoUring`] it probes the kernel and falls back to epoll
640    /// when io_uring is missing, disabled or too old ([`Driver::fallback_reason`] says why).
641    /// Asking for a back end this OS doesn't have is an `Unsupported` error.
642    pub fn new(cfg: &DriverConfig) -> io::Result<Driver> {
643        let want = cfg.backend.resolve();
644        #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
645        let mut fallback = None;
646        // io_uring, or why it fell back to epoll: the epoll driver is built in one place below
647        // (it is inlined into the loop thread's start, so a second copy would cost its code twice).
648        #[cfg(target_os = "linux")]
649        if want == Backend::IoUring {
650            #[cfg(feature = "io-uring")]
651            match uring::Driver::new(cfg) {
652                Ok(d) => {
653                    let fallback = d.degraded();
654                    return Ok(Driver { imp: Imp::Uring(d), kind: Backend::IoUring, fallback, _not_send: core::marker::PhantomData });
655                }
656                Err(e) => fallback = Some(uring_unavailable(&e)),
657            }
658            #[cfg(not(feature = "io-uring"))]
659            {
660                fallback = Some("io_uring unavailable (built without the `io-uring` feature); using epoll".into());
661            }
662        }
663        let (imp, kind) = match want {
664            #[cfg(target_os = "linux")]
665            Backend::Epoll | Backend::IoUring => (Imp::Epoll(readiness::Driver::new(cfg)?), Backend::Epoll),
666            #[cfg(windows)]
667            Backend::Iocp => (Imp::Iocp(iocp::Driver::new(cfg)?), Backend::Iocp),
668            #[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
669            Backend::Kqueue => (Imp::Kqueue(readiness::Driver::new(cfg)?), Backend::Kqueue),
670            other => return Err(unsupported_backend(other)),
671        };
672        Ok(Driver { imp, kind, fallback, _not_send: core::marker::PhantomData })
673    }
674
675    /// The back end in use.
676    pub fn backend(&self) -> Backend {
677        self.kind
678    }
679
680    /// Why the requested back end was replaced, if it was, or which of its optional features
681    /// the kernel lacks (io_uring set-up flags or registered files, each dropped on its own).
682    pub fn fallback_reason(&self) -> Option<&str> {
683        self.fallback.as_deref()
684    }
685
686    /// A handle that wakes this driver from any thread.
687    pub fn waker(&self) -> Waker {
688        dispatch_ref!(self, d => d.waker())
689    }
690
691    /// Registers a listening socket and starts accepting on it. Accepted connections arrive as
692    /// [`Event::Accepted`] with `key`.
693    pub fn listen(&mut self, socket: Socket, key: u64) -> io::Result<Id> {
694        dispatch!(self, d => d.listen(socket, key))
695    }
696
697    /// Registers a connected, non-blocking stream socket. Nothing is received until
698    /// [`Driver::recv`] turns receiving on.
699    pub fn stream(&mut self, socket: Socket, key: u64) -> io::Result<Id> {
700        dispatch!(self, d => d.stream(socket, key))
701    }
702
703    /// Registers a bound UDP socket. Datagrams arrive as [`Event::Datagram`] once
704    /// [`Driver::recv`] turns receiving on.
705    pub fn udp(&mut self, socket: Socket, key: u64) -> io::Result<Id> {
706        dispatch!(self, d => d.udp(socket, key))
707    }
708
709    /// Registers a bound UDP socket that receives into `buf_size`-byte buffers of the driver's own
710    /// (instead of the stock given with [`Driver::provide`]), for a socket with GRO on: a buffer
711    /// then takes a whole batch of datagrams (up to 64 KiB, cut by
712    /// [`RecvMeta::segment_size`]). Give the buffers of its [`Event::Datagram`]s back with
713    /// [`Driver::provide`]; the driver keeps a few for reuse.
714    pub fn udp_large(&mut self, socket: Socket, key: u64, buf_size: usize) -> io::Result<Id> {
715        dispatch!(self, d => d.udp_large(socket, key, buf_size))
716    }
717
718    /// Turns receiving on or off for a stream or UDP socket. Turning it off is backpressure: the
719    /// driver stops taking data, so the kernel's receive window closes. A few `Data` events
720    /// already under way may still arrive after it is turned off.
721    pub fn recv(&mut self, id: Id, on: bool) {
722        dispatch!(self, d => d.recv(id, on))
723    }
724
725    /// Queues a buffer to send on a stream. It is written at the start of the next turn, after
726    /// everything queued before it.
727    pub fn send(&mut self, id: Id, buf: OutBuf) {
728        dispatch!(self, d => d.send(id, buf))
729    }
730
731    /// Sends datagrams on a UDP socket at once, without queueing. Returns how many were handed
732    /// to the kernel; `WouldBlock` means the socket buffer is full and nothing was sent.
733    pub fn send_to(&mut self, id: Id, msgs: &[Transmit<'_>]) -> io::Result<usize> {
734        dispatch!(self, d => d.send_to(id, msgs))
735    }
736
737    /// Closes a socket. Queued and in-flight sends are discarded (their buffers come back as
738    /// [`Event::Sent`]), then [`Event::Closed`] is reported. Flush first for a graceful close.
739    pub fn close(&mut self, id: Id) {
740        dispatch!(self, d => d.close(id))
741    }
742
743    /// The socket behind `id`, for addresses and options.
744    pub fn socket(&self, id: Id) -> Option<&Socket> {
745        dispatch_ref!(self, d => d.socket(id))
746    }
747
748    /// Gives the driver a receive buffer (an empty `Vec` with spare capacity). The driver keeps
749    /// it until data arrives and then hands it back in an event. A buffer of the size given to
750    /// [`Driver::udp_large`] goes back to the large stock.
751    pub fn provide(&mut self, buf: Vec<u8>) {
752        dispatch!(self, d => d.provide(buf))
753    }
754
755    /// How many more receive buffers the driver would like (to reach its target stock).
756    pub fn buffers_wanted(&self) -> usize {
757        dispatch_ref!(self, d => d.buffers_wanted())
758    }
759
760    /// Sets how many receive buffers the driver keeps in stock (io_uring: at most the ring size).
761    pub fn set_buffer_target(&mut self, n: usize) {
762        dispatch!(self, d => d.set_buffer_target(n))
763    }
764
765    /// Writes a stream's queued output now, where that costs no extra system call (epoll and
766    /// kqueue write each socket with its own call anyway; io_uring batches its submissions and
767    /// IOCP posts them, so there this does nothing). The results come with the next turn (K4).
768    pub fn flush_now(&mut self, id: Id) {
769        match &mut self.imp {
770            #[cfg(target_os = "linux")]
771            Imp::Epoll(d) => d.flush_now(id),
772            #[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
773            Imp::Kqueue(d) => d.flush_now(id),
774            #[allow(unreachable_patterns)]
775            _ => {
776                let _ = id;
777            }
778        }
779    }
780
781    /// Frees spare receive buffers down to `keep` and lets the stock grow again from small
782    /// (an idle loop after a burst; buffers lent to an io_uring ring stay).
783    pub fn trim_buffers(&mut self, keep: usize) {
784        dispatch!(self, d => d.trim_buffers(keep))
785    }
786
787    /// Flushes queued sends, waits for I/O up to `timeout` (`None` waits until something happens
788    /// or a [`Waker`] fires; zero only polls) and appends what happened to `out`. Returns at once
789    /// when events are already pending.
790    pub fn turn(&mut self, timeout: Option<Duration>, out: &mut Vec<Event>) -> io::Result<()> {
791        dispatch!(self, d => d.turn(timeout, out))?;
792        if dispatch_ref!(self, d => d.waker.take_posted()) {
793            out.push(Event::Wake);
794        }
795        Ok(())
796    }
797
798    /// System call and event counters.
799    pub fn stats(&self) -> DriverStats {
800        dispatch_ref!(self, d => d.stats())
801    }
802
803    /// Registered sockets (listeners, streams and UDP sockets, including ones closing).
804    pub fn len(&self) -> usize {
805        dispatch_ref!(self, d => d.len())
806    }
807
808    /// Whether no socket is registered.
809    pub fn is_empty(&self) -> bool {
810        self.len() == 0
811    }
812}
813
814#[cfg(all(target_os = "linux", feature = "io-uring"))]
815#[cold]
816#[inline(never)]
817fn uring_unavailable(e: &io::Error) -> String {
818    format!("io_uring unavailable ({e}); using epoll")
819}
820
821#[cold]
822#[inline(never)]
823fn unsupported_backend(b: Backend) -> io::Error {
824    io::Error::new(io::ErrorKind::Unsupported, format!("the {} back end is not available on this OS", b.name()))
825}
826
827/// Converts an optional timeout to milliseconds for epoll and IOCP, rounding up so a timer
828/// never fires early (`-1` / `INFINITE` for none).
829#[allow(dead_code)]
830pub(crate) fn timeout_ms(t: Option<Duration>) -> i32 {
831    match t {
832        None => -1,
833        Some(d) => {
834            let ms = d.as_nanos().div_ceil(1_000_000);
835            ms.min(i32::MAX as u128) as i32
836        }
837    }
838}
839
840#[cfg(test)]
841mod tests;