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

compio_executor/
lib.rs
