tokio/runtime/io/
driver.rs

1// Signal handling
2cfg_signal_internal_and_unix! {
3    mod signal;
4}
5cfg_io_uring! {
6    mod uring;
7    use uring::UringContext;
8    pub(crate) use uring::RegisterResult;
9    use crate::sync::OnceCell;
10}
11
12use crate::io::interest::Interest;
13use crate::io::ready::Ready;
14use crate::loom::sync::Mutex;
15use crate::runtime::driver;
16use crate::runtime::io::registration_set;
17use crate::runtime::io::{IoDriverMetrics, RegistrationSet, ScheduledIo};
18
19use crate::loom::sync::Arc;
20use mio::event::Source;
21use std::fmt;
22use std::io;
23use std::time::Duration;
24
25/// I/O driver, backed by Mio.
26pub(crate) struct Driver {
27    /// True when an event with the signal token is received
28    signal_ready: bool,
29
30    /// Reuse the `mio::Events` value across calls to poll.
31    events: mio::Events,
32
33    /// Buffer for polls that do not wait, if
34    /// `Builder::max_io_events_per_busy_tick` is set.
35    events_busy: Option<mio::Events>,
36
37    /// The system event queue.
38    poll: mio::Poll,
39}
40
41/// A reference to an I/O driver.
42pub(crate) struct Handle {
43    /// Registers I/O resources.
44    registry: mio::Registry,
45
46    /// Tracks all registrations
47    registrations: RegistrationSet,
48
49    /// State that should be synchronized
50    synced: Mutex<registration_set::Synced>,
51
52    /// Used to wake up the reactor from a call to `turn`.
53    /// Not supported on `Wasi` due to lack of threading support.
54    #[cfg(not(target_os = "wasi"))]
55    waker: mio::Waker,
56
57    pub(crate) metrics: IoDriverMetrics,
58
59    #[cfg(all(
60        tokio_unstable,
61        feature = "io-uring",
62        feature = "rt",
63        feature = "fs",
64        target_os = "linux",
65    ))]
66    pub(crate) uring_context: Mutex<UringContext>,
67
68    #[cfg(all(
69        tokio_unstable,
70        feature = "io-uring",
71        feature = "rt",
72        feature = "fs",
73        target_os = "linux",
74    ))]
75    pub(crate) uring_probe: OnceCell<Option<io_uring::Probe>>,
76}
77
78#[derive(Debug)]
79pub(crate) struct ReadyEvent {
80    pub(super) tick: u16,
81    pub(crate) ready: Ready,
82    pub(super) is_shutdown: bool,
83}
84
85cfg_net_unix!(
86    impl ReadyEvent {
87        pub(crate) fn with_ready(&self, ready: Ready) -> Self {
88            Self {
89                ready,
90                tick: self.tick,
91                is_shutdown: self.is_shutdown,
92            }
93        }
94    }
95);
96
97#[derive(Debug, Eq, PartialEq, Clone, Copy)]
98pub(super) enum Direction {
99    Read,
100    Write,
101}
102
103pub(super) enum Tick {
104    Set,
105    Clear(u16),
106}
107
108const TOKEN_WAKEUP: mio::Token = mio::Token(0);
109const TOKEN_SIGNAL: mio::Token = mio::Token(1);
110
111fn _assert_kinds() {
112    fn _assert<T: Send + Sync>() {}
113
114    _assert::<Handle>();
115}
116
117// ===== impl Driver =====
118
119impl Driver {
120    /// Creates a new event loop, returning any error that happened during the
121    /// creation.
122    pub(crate) fn new(nevents: usize, nevents_busy: Option<usize>) -> io::Result<(Driver, Handle)> {
123        let poll = mio::Poll::new()?;
124        #[cfg(not(target_os = "wasi"))]
125        let waker = mio::Waker::new(poll.registry(), TOKEN_WAKEUP)?;
126        let registry = poll.registry().try_clone()?;
127
128        let driver = Driver {
129            signal_ready: false,
130            events: mio::Events::with_capacity(nevents),
131            events_busy: nevents_busy.map(mio::Events::with_capacity),
132            poll,
133        };
134
135        let (registrations, synced) = RegistrationSet::new();
136
137        let handle = Handle {
138            registry,
139            registrations,
140            synced: Mutex::new(synced),
141            #[cfg(not(target_os = "wasi"))]
142            waker,
143            metrics: IoDriverMetrics::default(),
144            #[cfg(all(
145                tokio_unstable,
146                feature = "io-uring",
147                feature = "rt",
148                feature = "fs",
149                target_os = "linux",
150            ))]
151            uring_context: Mutex::new(UringContext::new()),
152            #[cfg(all(
153                tokio_unstable,
154                feature = "io-uring",
155                feature = "rt",
156                feature = "fs",
157                target_os = "linux",
158            ))]
159            uring_probe: OnceCell::new(),
160        };
161
162        Ok((driver, handle))
163    }
164
165    pub(crate) fn park(&mut self, rt_handle: &driver::Handle) {
166        let handle = rt_handle.io();
167        self.turn(handle, None);
168    }
169
170    pub(crate) fn park_timeout(&mut self, rt_handle: &driver::Handle, duration: Duration) {
171        let handle = rt_handle.io();
172        self.turn(handle, Some(duration));
173    }
174
175    pub(crate) fn shutdown(&mut self, rt_handle: &driver::Handle) {
176        let handle = rt_handle.io();
177        let ios = handle.registrations.shutdown(&mut handle.synced.lock());
178
179        // `shutdown()` must be called without holding the lock.
180        for io in ios {
181            io.shutdown();
182        }
183    }
184
185    fn turn(&mut self, handle: &Handle, max_wait: Option<Duration>) {
186        debug_assert!(!handle.registrations.is_shutdown(&handle.synced.lock()));
187
188        handle.release_pending_registrations();
189
190        // A poll that does not wait takes the busy batch. Events it leaves
191        // behind stay queued in the kernel, so the next poll returns them.
192        let events = match (&mut self.events_busy, max_wait) {
193            (Some(busy), Some(wait)) if wait.is_zero() => busy,
194            _ => &mut self.events,
195        };
196
197        // Block waiting for an event to happen, peeling out how many events
198        // happened.
199        match self.poll.poll(events, max_wait) {
200            Ok(()) => {}
201            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
202            #[cfg(target_os = "wasi")]
203            Err(e) if e.kind() == io::ErrorKind::InvalidInput => {
204                // In case of wasm32_wasi this error happens, when trying to poll without subscriptions
205                // just return from the park, as there would be nothing, which wakes us up.
206            }
207            Err(e) => panic!("unexpected error when polling the I/O driver: {e:?}"),
208        }
209
210        // Process all the events that came in, dispatching appropriately
211        let mut ready_count = 0;
212        for event in events.iter() {
213            let token = event.token();
214
215            if token == TOKEN_WAKEUP {
216                // Nothing to do, the event is used to unblock the I/O driver
217            } else if token == TOKEN_SIGNAL {
218                self.signal_ready = true;
219            } else {
220                let ready = Ready::from_mio(event);
221
222                // Safety: we ensure that the pointers used as tokens are not freed
223                // until they are both deregistered from mio **and** we know the I/O
224                // driver is not concurrently polling. The I/O driver holds ownership of
225                // an `Arc<ScheduledIo>` so we can safely cast this to a ref.
226                unsafe { dispatch_event(token, ready) };
227
228                ready_count += 1;
229            }
230        }
231
232        #[cfg(all(
233            tokio_unstable,
234            feature = "io-uring",
235            feature = "rt",
236            feature = "fs",
237            target_os = "linux",
238        ))]
239        {
240            let mut guard = handle.get_uring().lock();
241            let ctx = &mut *guard;
242            ctx.dispatch_completions();
243
244            // There might be some cases where the CQ overflows, so we need to flush
245            // the remaining buffered CQEs.
246            while ctx
247                .uring
248                .as_mut()
249                .is_some_and(|uring| uring.submission().cq_overflow())
250            {
251                ctx.submit()
252                    .expect("failed to flush io_uring completion queue overflow");
253                ctx.dispatch_completions();
254            }
255        }
256
257        handle.metrics.incr_ready_count_by(ready_count);
258    }
259}
260
261impl fmt::Debug for Driver {
262    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
263        write!(f, "Driver")
264    }
265}
266
267impl Handle {
268    /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise
269    /// makes the next call to `turn` return immediately.
270    ///
271    /// This method is intended to be used in situations where a notification
272    /// needs to otherwise be sent to the main reactor. If the reactor is
273    /// currently blocked inside of `turn` then it will wake up and soon return
274    /// after this method has been called. If the reactor is not currently
275    /// blocked in `turn`, then the next call to `turn` will not block and
276    /// return immediately.
277    pub(crate) fn unpark(&self) {
278        #[cfg(not(target_os = "wasi"))]
279        self.waker.wake().expect("failed to wake I/O driver");
280    }
281
282    /// Registers an I/O resource with the reactor for a given `mio::Ready` state.
283    ///
284    /// The registration token is returned.
285    pub(super) fn add_source(
286        &self,
287        source: &mut impl mio::event::Source,
288        interest: Interest,
289    ) -> io::Result<Arc<ScheduledIo>> {
290        let scheduled_io = self.registrations.allocate(&mut self.synced.lock())?;
291        let token = scheduled_io.token();
292
293        // we should remove the `scheduled_io` from the `registrations` set if registering
294        // the `source` with the OS fails. Otherwise it will leak the `scheduled_io`.
295        if let Err(e) = self.registry.register(source, token, interest.to_mio()) {
296            // safety: `scheduled_io` is part of the `registrations` set.
297            unsafe {
298                self.registrations
299                    .remove(&mut self.synced.lock(), &scheduled_io)
300            };
301
302            return Err(e);
303        }
304
305        // TODO: move this logic to `RegistrationSet` and use a `CountedLinkedList`
306        self.metrics.incr_fd_count();
307
308        Ok(scheduled_io)
309    }
310
311    /// Deregisters an I/O resource from the reactor.
312    pub(super) fn deregister_source(
313        &self,
314        registration: &Arc<ScheduledIo>,
315        source: &mut impl Source,
316    ) -> io::Result<()> {
317        // Deregister the source with the OS poller **first**.
318        //
319        // If `deregister` fails (for example, because the caller closed the raw
320        // file descriptor before dropping `AsyncFd`), we must NOT release the
321        // `ScheduledIo` registration. On Linux, if the file descriptor was
322        // duplicated before the original fd number was closed, the open file
323        // description remains registered in epoll with the `ScheduledIo` token
324        // pointer even though `EPOLL_CTL_DEL` fails with `EBADF`.
325        self.registry.deregister(source)?;
326
327        if self
328            .registrations
329            .deregister(&mut self.synced.lock(), registration)
330        {
331            self.unpark();
332        }
333
334        self.metrics.dec_fd_count();
335
336        Ok(())
337    }
338
339    fn release_pending_registrations(&self) {
340        if self.registrations.needs_release() {
341            self.registrations.release(&mut self.synced.lock());
342        }
343    }
344}
345
346/// Dispatches an event received from the poller to the corresponding `ScheduledIo`.
347///
348/// # Safety
349///
350/// `token` must have been obtained from `ScheduledIo::token()`, and that
351/// `ScheduledIo` must remain alive throughout this call. The I/O driver ensures
352/// that pending registrations are not freed until after event dispatch completes.
353pub(super) unsafe fn dispatch_event(token: mio::Token, ready: Ready) {
354    let ptr = super::EXPOSE_IO.from_exposed_addr(token.0);
355
356    // SAFETY: the caller guarantees that the token refers to a live `ScheduledIo`.
357    let io: &ScheduledIo = unsafe { &*ptr };
358
359    io.set_readiness(Tick::Set, |curr| curr | ready);
360    io.wake(ready);
361}
362
363impl fmt::Debug for Handle {
364    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
365        write!(f, "Handle")
366    }
367}
368
369impl Direction {
370    pub(super) fn mask(self) -> Ready {
371        match self {
372            Direction::Read => Ready::READABLE | Ready::READ_CLOSED,
373            Direction::Write => Ready::WRITABLE | Ready::WRITE_CLOSED,
374        }
375    }
376}
377
378#[cfg(all(test, unix, feature = "net", not(loom), not(miri)))]
379mod tests {
380    use super::*;
381    use std::io::Write;
382
383    #[test]
384    fn busy_turn_takes_busy_batch() {
385        let (mut driver, handle) = Driver::new(16, Some(2)).unwrap();
386        let mut sources = Vec::new();
387        for _ in 0..5 {
388            let (mut rx, mut tx) = mio::net::UnixStream::pair().unwrap();
389            tx.write_all(b"x").unwrap();
390            let reg = handle.add_source(&mut rx, Interest::READABLE).unwrap();
391            sources.push((rx, tx, reg));
392        }
393
394        // A poll that does not wait takes the busy batch.
395        driver.turn(&handle, Some(Duration::ZERO));
396        assert_eq!(driver.events_busy.as_ref().unwrap().iter().count(), 2);
397
398        // The rest stays queued for the next poll, which takes the main batch.
399        driver.turn(&handle, Some(Duration::from_millis(100)));
400        assert_eq!(driver.events.iter().count(), 3);
401
402        for (mut rx, _tx, reg) in sources {
403            handle.deregister_source(&reg, &mut rx).unwrap();
404        }
405        handle.release_pending_registrations();
406    }
407}
tokio/runtime/io/driver.rs

tokio/runtime/io/
driver.rs
