nuclei/
proactor.rs

1//! The proactor facade of earlier nuclei versions, over the runtime's
2//! shards: every shard drives its own IO, and these calls reach the shard
3//! of the calling thread.
4
5use core::future::Future;
6use std::io;
7use std::sync::{Arc, OnceLock};
8use std::time::Duration;
9
10use crate::config::NucleiConfig;
11use crate::rt::context;
12use crate::rt::shard::Shared;
13
14/// Which kernel interface nuclei's IO runs on.
15#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
16pub enum IoBackend {
17    /// kqueue (macOS and the BSDs).
18    Kqueue,
19    /// epoll (Linux).
20    Epoll,
21    /// io_uring (Linux).
22    IoUring,
23    /// I/O completion ports (Windows).
24    Iocp,
25}
26
27/// The process's proactor: its configuration, and the IO of the calling
28/// thread's shard.
29#[derive(Debug)]
30pub struct Proactor {
31    config: NucleiConfig,
32}
33
34static PROACTOR: OnceLock<Proactor> = OnceLock::new();
35
36/// The shards of threads inside [`Proactor::wait`], by address.
37type Waiting = kovan_map::HashMap<usize, Arc<Shared>>;
38
39static WAITING: OnceLock<Waiting> = OnceLock::new();
40
41fn waiting() -> &'static Waiting {
42    // 64 is the smallest table the map accepts.
43    WAITING.get_or_init(|| Waiting::with_capacity(64))
44}
45
46/// Ends the waits of every shard in `waiting`: a wake is left on each
47/// one's poller, so a wait that has not begun yet (its thread is inside
48/// `Proactor::wait`, before the poller) ends as soon as it does.
49fn wake_all(waiting: &Waiting) {
50    for (_, shard) in waiting.iter() {
51        shard.notify();
52    }
53}
54
55/// A thread's shard in [`WAITING`] while its `Proactor::wait` lasts.
56struct InWait(usize);
57
58impl InWait {
59    fn new(shard: &Arc<Shared>) -> Self {
60        let key = Arc::as_ptr(shard) as usize;
61        waiting().insert(key, Arc::clone(shard));
62        Self(key)
63    }
64}
65
66impl Drop for InWait {
67    fn drop(&mut self) {
68        waiting().remove(&self.0);
69    }
70}
71
72impl Proactor {
73    /// The proactor, with the default configuration unless
74    /// [`with_config`](Self::with_config) set one first.
75    pub fn get() -> &'static Proactor {
76        PROACTOR.get_or_init(|| Proactor {
77            config: NucleiConfig::default(),
78        })
79    }
80
81    /// Sets the proactor's configuration, and returns the proactor. It is
82    /// the default every runtime starts from;
83    /// [`Builder::io_uring`](crate::Builder::io_uring) gives one runtime
84    /// ring settings of its own, checked when it is built.
85    ///
86    /// # Panics
87    ///
88    /// When the proactor was configured already, or used (which configures
89    /// it with the defaults): a configuration is set once, before any IO.
90    /// The signature returns the proactor itself, so it has no error to
91    /// answer; `Builder::io_uring` answers its error instead.
92    pub fn with_config(config: NucleiConfig) -> &'static Proactor {
93        let mut config = Some(config);
94        let proactor = PROACTOR.get_or_init(|| Proactor {
95            config: config.take().unwrap_or_default(),
96        });
97        assert!(
98            config.is_none(),
99            "the proactor is configured already: call `with_config` once, before any IO"
100        );
101        proactor
102    }
103
104    /// The configuration in place.
105    pub fn config(&self) -> &NucleiConfig {
106        &self.config
107    }
108
109    /// Ends every [`wait`](Self::wait) in progress at once.
110    pub fn wake(&self) {
111        wake_all(waiting());
112    }
113
114    /// Drives the calling thread's shard once: runs its ready tasks, waits
115    /// for IO (up to `duration`; `None` until some comes or [`wake`]),
116    /// handling at most `max_event_size` events, and runs the tasks the IO
117    /// woke. Returns how many IO events it handled.
118    ///
119    /// [`wake`]: Self::wake
120    ///
121    /// # Errors
122    ///
123    /// When the thread's shard cannot start, or its poller fails.
124    pub fn wait(&self, max_event_size: usize, duration: Option<Duration>) -> io::Result<usize> {
125        if let Some(events) =
126            context::current(|current| current.map(|s| s.turn(max_event_size, duration)))
127        {
128            return events;
129        }
130        let shard = context::own_shard()?;
131        let _current = context::enter(&shard);
132        #[cfg(port)]
133        shard.take_port_back();
134        let _waiting = InWait::new(&shard.shared);
135        shard.turn(max_event_size, duration)
136    }
137
138    /// The backend the calling thread's shard runs its IO on; on a thread
139    /// that runs none, the one of every runtime that names none (the
140    /// platform's poller when the backend named cannot be had: the shards
141    /// then fail with the reason).
142    pub fn backend() -> IoBackend {
143        match context::current(|shard| shard.map(|s| s.io_backend())) {
144            Some(running) => running,
145            None => crate::rt::backend::chosen().map_or(POLLER, |b| {
146                if b.rings() {
147                    IoBackend::IoUring
148                } else {
149                    POLLER
150                }
151            }),
152        }
153    }
154}
155
156/// The platform's readiness poller.
157#[cfg(any(target_os = "linux", target_os = "android", target_os = "illumos"))]
158pub(crate) const POLLER: IoBackend = IoBackend::Epoll;
159#[cfg(any(
160    target_os = "macos",
161    target_os = "ios",
162    target_os = "freebsd",
163    target_os = "netbsd",
164    target_os = "openbsd",
165    target_os = "dragonfly",
166))]
167pub(crate) const POLLER: IoBackend = IoBackend::Kqueue;
168#[cfg(target_os = "windows")]
169pub(crate) const POLLER: IoBackend = IoBackend::Iocp;
170
171/// Runs `future` to its end on the calling thread, driving its shard's IO
172/// meanwhile: [`block_on`](crate::runtime::block_on).
173pub fn drive<T>(future: impl Future<Output = T>) -> T {
174    crate::runtime::block_on(future)
175}
176
177#[cfg(test)]
178mod tests {
179    use std::sync::Arc;
180    use std::time::{Duration, Instant};
181
182    use super::{Waiting, wake_all};
183    use crate::rt::context;
184
185    /// A wake that comes after the thread went into the wait, but before
186    /// its shard waits on the poller, still ends that wait.
187    #[test]
188    fn a_wake_before_the_poller_wait_ends_it() {
189        let shard = context::own_shard().unwrap();
190        let _current = context::enter(&shard);
191        let waiting = Waiting::with_capacity(64);
192        waiting.insert(0, Arc::clone(&shard.shared));
193        wake_all(&waiting);
194        let start = Instant::now();
195        shard.turn(64, Some(Duration::from_secs(10))).unwrap();
196        assert!(
197            start.elapsed() < Duration::from_secs(5),
198            "the wake was lost: the wait ran {:?}",
199            start.elapsed()
200        );
201    }
202}
nuclei/proactor.rs

nuclei/
proactor.rs
