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;