compio_executor/
lib.rs

1//! Executor for compio runtime.
2
3#![cfg_attr(docsrs, feature(doc_cfg))]
4#![allow(unused_features)]
5#![warn(missing_docs)]
6#![deny(rustdoc::broken_intra_doc_links)]
7#![doc(
8    html_logo_url = "https://github.com/compio-rs/compio-logo/raw/refs/heads/master/generated/colored-bold.svg"
9)]
10#![doc(
11    html_favicon_url = "https://github.com/compio-rs/compio-logo/raw/refs/heads/master/generated/colored-bold.svg"
12)]
13
14use std::{any::Any, fmt::Debug, ptr::NonNull, task::Waker};
15
16use crate::queue::{TaskId, TaskQueue};
17
18pub mod console;
19mod join_handle;
20mod queue;
21mod task;
22mod util;
23mod waker;
24
25use compio_log::{instrument, trace};
26use compio_send_wrapper::SendWrapper;
27pub use console::SpawnMeta;
28use crossbeam_queue::ArrayQueue;
29pub use join_handle::{JoinError, JoinHandle, ResumeUnwind};
30use util::panic_guard;
31
32cfg_select! {
33    loom => {
34        use loom::{cell::UnsafeCell, hint, sync::atomic::*, thread::yield_now};
35    }
36    _ => {
37        use std::{hint, sync::atomic::*, thread::yield_now};
38
39        #[repr(transparent)]
40        struct UnsafeCell<T>(std::cell::UnsafeCell<T>);
41
42        impl<T> UnsafeCell<T> {
43            pub fn new(value: T) -> Self {
44                Self(std::cell::UnsafeCell::new(value))
45            }
46
47            #[inline(always)]
48            pub fn with_mut<F, R>(&self, f: F) -> R
49            where
50                F: FnOnce(*mut T) -> R,
51            {
52                f(self.0.get())
53            }
54
55            #[inline(always)]
56            pub fn with<F, R>(&self, f: F) -> R
57            where
58                F: FnOnce(*const T) -> R,
59            {
60                f(self.0.get())
61            }
62        }
63    }
64}
65
66pub(crate) type PanicResult<T> = Result<T, Panic>;
67pub(crate) type Panic = Box<dyn Any + Send + 'static>;
68
69/// A dual-queue executor optimized for singlethreaded usecase, with support for
70/// multithreaded wakes.
71///
72/// Same-thread wakes ([`Waker::wake`]) will schedule tasks within the queue
73/// directly; cross-thread wakes will send task id's to a channel, and
74/// piggybacked to singlethreaded wakes or ticks. This ensures maximum
75/// performance for singlethreaded scenario at the trade-off of worse tail
76/// latency for multithreaded wake-ups.
77///
78/// Optionally, all [`Waker`]s generated from this executor can contain an extra
79/// data, parameterized as `E`.
80///
81/// [`Waker`]: std::task::Waker
82/// [`Waker::wake`]: std::task::Waker::wake
83#[derive(Debug)]
84pub struct Executor {
85    ptr: NonNull<Shared>,
86    config: ExecutorConfig,
87}
88
89/// Configuration for [`Executor`].
90#[derive(Debug, Clone)]
91pub struct ExecutorConfig {
92    /// The size of the sync queue, which holds task id's for cross-thread
93    /// wakes.
94    ///
95    /// This is fixed and will create backpressure when full.
96    pub sync_queue_size: usize,
97
98    /// The size of the local queues, which hold tasks for same-thread
99    /// execution.
100    ///
101    /// This is dynamically resized to avoid blocking.
102    pub local_queue_size: usize,
103
104    /// The maximum number of hot tasks to run in each tick.
105    pub max_interval: u32,
106
107    /// A waker to be woken when a task is scheduled.
108    ///
109    /// This is useful for waking up drivers that switch to kernel state when
110    /// idle.
111    pub waker: Option<Waker>,
112}
113
114impl Default for ExecutorConfig {
115    fn default() -> Self {
116        Self {
117            sync_queue_size: 64,
118            local_queue_size: 64,
119            max_interval: 61,
120            waker: None,
121        }
122    }
123}
124
125pub(crate) struct Shared {
126    waker: Option<Waker>,
127    sync: ArrayQueue<TaskId>,
128    pending: AtomicUsize,
129    queue: SendWrapper<TaskQueue>,
130}
131
132impl Shared {
133    /// Drain all pending cross-thread wakes into the local hot `queue`.
134    ///
135    /// Skips the expensive [`ArrayQueue::pop`] entirely when nothing has been
136    /// pushed, using a single relaxed-ish load of [`Shared::pending`] instead
137    /// of crossbeam's `SeqCst` empty check.
138    #[inline]
139    pub(crate) fn drain_sync(&self, queue: &TaskQueue) {
140        if self.pending.load(Ordering::Acquire) == 0 {
141            return;
142        }
143
144        let mut drained: usize = 0;
145        while let Some(id) = self.sync.pop() {
146            queue.make_hot(id);
147            drained += 1;
148        }
149
150        if drained != 0 {
151            self.pending.fetch_sub(drained, Ordering::Release);
152        }
153    }
154}
155
156impl Executor {
157    /// Create a new executor.
158    pub fn new() -> Self {
159        Self::with_config(ExecutorConfig::default())
160    }
161
162    /// Create a new executor with config.
163    pub fn with_config(mut config: ExecutorConfig) -> Self {
164        let ptr = Box::into_raw(Box::new(Shared {
165            waker: config.waker.take(),
166            sync: ArrayQueue::new(config.sync_queue_size),
167            pending: AtomicUsize::new(0),
168            queue: SendWrapper::new(TaskQueue::new(config.local_queue_size)),
169        }));
170
171        Self {
172            config,
173            ptr: unsafe { NonNull::new_unchecked(ptr) },
174        }
175    }
176
177    /// Spawn a future onto the executor.
178    #[track_caller]
179    pub fn spawn<F: Future + 'static>(&self, fut: F) -> JoinHandle<F::Output> {
180        self.spawn_at(fut, SpawnMeta::capture())
181    }
182
183    /// Spawn a future onto the executor, attributing it to `meta`.
184    ///
185    /// This is only useful for wrappers around [`spawn`] that want
186    /// [`tokio-console`] to blame their own caller instead of themselves;
187    /// [`SpawnMeta`] is a zero-sized no-op without the `console` feature.
188    ///
189    /// [`spawn`]: Self::spawn
190    /// [`tokio-console`]: crate::console
191    pub fn spawn_at<F: Future + 'static>(&self, fut: F, meta: SpawnMeta) -> JoinHandle<F::Output> {
192        let shared = self.shared();
193        let tracker = shared.queue.tracker();
194        // SAFETY: Executor cannot be sent to ther thread
195        let queue = unsafe { shared.queue.get_unchecked() };
196        let task = queue.insert(self.ptr, tracker, fut, meta);
197
198        JoinHandle::new(task)
199    }
200
201    /// Retrieve all sync tasks, schedule those to the tail of `hot` queue
202    /// and run at most [`max_interval`] tasks.
203    ///
204    /// Running start with `hot` tasks, then `cold` ones. Finished tasks will
205    /// be pushed back to tail of `cold` queue.
206    ///
207    /// Return whether there are still hot tasks after the tick.
208    ///
209    /// [`max_interval`]: ExecutorConfig::max_interval
210    pub fn tick(&self) -> bool {
211        let queue = self.queue();
212
213        self.shared().drain_sync(queue);
214
215        for id in queue.iter_hot().take(self.config.max_interval as _) {
216            queue.make_cold(id);
217            let task = queue.take(id).expect("Task was not reset back");
218            let res = unsafe { task.run() };
219            if res.is_ready() {
220                // SAFETY: We're removing it soon, so drop will only be called once.
221                // The shared pointer is kept valid until the Executor is dropped,
222                // to avoid use-after-free issues with concurrent wakers.
223                unsafe { task.drop() };
224                queue.remove(id);
225            } else {
226                queue.reset(id, task);
227            }
228        }
229
230        queue.has_hot()
231    }
232
233    /// Check if there's still scheduled task that needs to be ran.
234    #[doc(hidden)]
235    pub fn has_task(&self) -> bool {
236        self.queue().hot_head().is_some()
237    }
238
239    /// Clear the executor, drop all tasks.
240    ///
241    /// This should be called only in context of the runtime, if any future may
242    /// use it. Any panic happened during dropping the future will cause the
243    /// process to abort. If this was not called before dropping, all tasks will
244    /// be leakded.
245    pub fn clear(&self) {
246        instrument!(compio_log::Level::TRACE, "Executor::drop");
247        trace!("Dropping Executor");
248
249        while self.shared().sync.pop().is_some() {}
250        unsafe { self.queue().clear() };
251    }
252
253    #[inline(always)]
254    fn shared(&self) -> &Shared {
255        unsafe { self.ptr.as_ref() }
256    }
257
258    #[inline(always)]
259    fn queue(&self) -> &TaskQueue {
260        // SAFETY: Executor is single threaded
261        unsafe { self.shared().queue.get_unchecked() }
262    }
263}
264
265impl Drop for Executor {
266    fn drop(&mut self) {
267        self.clear();
268        unsafe { drop(Box::from_raw(self.ptr.as_ptr())) };
269    }
270}
271
272impl Default for Executor {
273    fn default() -> Self {
274        Self::new()
275    }
276}
compio_executor/lib.rs

compio_executor/
lib.rs
