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::SegQueue;
29pub use join_handle::{JoinError, JoinHandle, ResumeUnwind};
30use util::panic_guard;
31
32cfg_select! {
33    loom => {
34        use loom::{cell::UnsafeCell, hint, sync::atomic::*};
35    }
36    _ => {
37        use std::{hint, sync::atomic::*};
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 local queues, which hold tasks for same-thread
93    /// execution.
94    ///
95    /// This is dynamically resized to avoid blocking.
96    pub local_queue_size: usize,
97
98    /// The maximum number of hot tasks to run in each tick.
99    pub max_interval: u32,
100
101    /// A waker to be woken when a task is scheduled.
102    ///
103    /// This is useful for waking up drivers that switch to kernel state when
104    /// idle.
105    pub waker: Option<Waker>,
106}
107
108impl Default for ExecutorConfig {
109    fn default() -> Self {
110        Self {
111            local_queue_size: 64,
112            max_interval: 61,
113            waker: None,
114        }
115    }
116}
117
118pub(crate) struct Shared {
119    waker: Option<Waker>,
120    sync: SegQueue<TaskId>,
121    pending: AtomicUsize,
122    queue: SendWrapper<TaskQueue>,
123}
124
125impl Shared {
126    /// Drain all pending cross-thread wakes into the local hot `queue`.
127    ///
128    /// Skips the expensive [`SegQueue::pop`] entirely when nothing has been
129    /// pushed, using a single relaxed-ish load of [`Shared::pending`] instead
130    /// of crossbeam's `SeqCst` empty check.
131    #[inline]
132    pub(crate) fn drain_sync(&self, queue: &TaskQueue) {
133        if self.pending.load(Ordering::Acquire) == 0 {
134            return;
135        }
136
137        let mut drained: usize = 0;
138        while let Some(id) = self.sync.pop() {
139            queue.make_hot(id);
140            drained += 1;
141        }
142
143        if drained != 0 {
144            self.pending.fetch_sub(drained, Ordering::Release);
145        }
146    }
147}
148
149impl Executor {
150    /// Create a new executor.
151    pub fn new() -> Self {
152        Self::with_config(ExecutorConfig::default())
153    }
154
155    /// Create a new executor with config.
156    pub fn with_config(mut config: ExecutorConfig) -> Self {
157        let ptr = Box::into_raw(Box::new(Shared {
158            waker: config.waker.take(),
159            sync: SegQueue::new(),
160            pending: AtomicUsize::new(0),
161            queue: SendWrapper::new(TaskQueue::new(config.local_queue_size)),
162        }));
163
164        Self {
165            config,
166            ptr: unsafe { NonNull::new_unchecked(ptr) },
167        }
168    }
169
170    /// Spawn a future onto the executor.
171    #[track_caller]
172    pub fn spawn<F: Future + 'static>(&self, fut: F) -> JoinHandle<F::Output> {
173        self.spawn_at(fut, SpawnMeta::capture())
174    }
175
176    /// Spawn a future onto the executor, attributing it to `meta`.
177    ///
178    /// This is only useful for wrappers around [`spawn`] that want
179    /// [`tokio-console`] to blame their own caller instead of themselves;
180    /// [`SpawnMeta`] is a zero-sized no-op without the `console` feature.
181    ///
182    /// [`spawn`]: Self::spawn
183    /// [`tokio-console`]: crate::console
184    pub fn spawn_at<F: Future + 'static>(&self, fut: F, meta: SpawnMeta) -> JoinHandle<F::Output> {
185        let shared = self.shared();
186        let tracker = shared.queue.tracker();
187        // SAFETY: Executor cannot be sent to ther thread
188        let queue = unsafe { shared.queue.get_unchecked() };
189        let task = queue.insert(self.ptr, tracker, fut, meta);
190
191        JoinHandle::new(task)
192    }
193
194    /// Retrieve all sync tasks, schedule those to the tail of `hot` queue
195    /// and run at most [`max_interval`] tasks.
196    ///
197    /// Running start with `hot` tasks, then `cold` ones. Finished tasks will
198    /// be pushed back to tail of `cold` queue.
199    ///
200    /// Return whether there are still hot tasks after the tick.
201    ///
202    /// [`max_interval`]: ExecutorConfig::max_interval
203    pub fn tick(&self) -> bool {
204        let queue = self.queue();
205
206        self.shared().drain_sync(queue);
207
208        for _ in 0..self.config.max_interval {
209            let Some(id) = queue.hot_head() else {
210                return false;
211            };
212
213            queue.make_cold(id);
214
215            let task = queue.take(id).expect("Task was not reset back");
216            // SAFETY: Only the owning executor polls this task. Completed
217            // tasks are removed from the queue before they can be run again.
218            let res = unsafe { task.run() };
219
220            if res.is_ready() {
221                // SAFETY: We're removing it soon, so drop will only be called
222                // once. The shared pointer is kept valid until
223                // the Executor is dropped,
224                // to avoid use-after-free issues with concurrent wakers.
225                unsafe { task.drop() };
226                queue.remove(id);
227            } else {
228                queue.reset(id, task);
229            }
230        }
231
232        queue.has_hot()
233    }
234
235    /// Check if there's still scheduled task that needs to be ran.
236    #[doc(hidden)]
237    pub fn has_task(&self) -> bool {
238        self.queue().has_hot()
239    }
240
241    /// Returns the number of alive tasks in the executor.
242    ///
243    /// A task is alive from the moment it is spawned until the [`tick`] where
244    /// it completes or where its cancellation is processed. It does not
245    /// matter whether the task is scheduled to run or waiting to be woken.
246    ///
247    /// [`tick`]: Self::tick
248    pub fn num_alive_tasks(&self) -> usize {
249        self.queue().len()
250    }
251
252    /// Clear the executor, drop all tasks.
253    ///
254    /// This should be called only in context of the runtime, if any future may
255    /// use it. Any panic happened during dropping the future will cause the
256    /// process to abort. If this was not called before dropping, all tasks will
257    /// be leakded.
258    pub fn clear(&self) {
259        instrument!(compio_log::Level::TRACE, "Executor::drop");
260        trace!("Dropping Executor");
261
262        while self.shared().sync.pop().is_some() {}
263        unsafe { self.queue().clear() };
264    }
265
266    #[inline(always)]
267    fn shared(&self) -> &Shared {
268        unsafe { self.ptr.as_ref() }
269    }
270
271    #[inline(always)]
272    fn queue(&self) -> &TaskQueue {
273        // SAFETY: Executor is single threaded
274        unsafe { self.shared().queue.get_unchecked() }
275    }
276}
277
278impl Drop for Executor {
279    fn drop(&mut self) {
280        self.clear();
281        unsafe { drop(Box::from_raw(self.ptr.as_ptr())) };
282    }
283}
284
285impl Default for Executor {
286    fn default() -> Self {
287        Self::new()
288    }
289}
compio_executor/lib.rs

compio_executor/
lib.rs
