compio_dispatcher/
lib.rs

1//! Multithreading dispatcher.
2
3#![allow(unused_features)]
4#![warn(missing_docs)]
5#![deny(rustdoc::broken_intra_doc_links)]
6#![doc(
7    html_logo_url = "https://github.com/compio-rs/compio-logo/raw/refs/heads/master/generated/colored-bold.svg"
8)]
9#![doc(
10    html_favicon_url = "https://github.com/compio-rs/compio-logo/raw/refs/heads/master/generated/colored-bold.svg"
11)]
12
13use std::{
14    collections::HashSet,
15    io,
16    num::NonZeroUsize,
17    panic::resume_unwind,
18    thread::{JoinHandle, available_parallelism},
19};
20
21use compio_driver::{AsyncifyPool, DispatchError, Dispatchable, ProactorBuilder};
22use compio_runtime::{JoinHandle as CompioJoinHandle, Runtime, SpawnMeta};
23use flume::{Sender, unbounded};
24use futures_channel::oneshot;
25
26/// A closure to spawn, and the [`SpawnMeta`] of the `dispatch` call it came
27/// from.
28struct Spawning {
29    task: Box<dyn Spawnable + Send>,
30    meta: SpawnMeta,
31}
32
33trait Spawnable {
34    fn spawn(self: Box<Self>, handle: &Runtime, meta: SpawnMeta) -> CompioJoinHandle<()>;
35}
36
37/// Concrete type for the closure we're sending to worker threads
38struct Concrete<F, R> {
39    callback: oneshot::Sender<R>,
40    func: F,
41}
42
43impl<F, R> Concrete<F, R> {
44    pub fn new(func: F) -> (Self, oneshot::Receiver<R>) {
45        let (tx, rx) = oneshot::channel();
46        (Self { callback: tx, func }, rx)
47    }
48}
49
50impl<F, R> Spawnable for Concrete<F, R>
51where
52    F: AsyncFnOnce() -> R + Send + 'static,
53    R: Send + 'static,
54{
55    fn spawn(self: Box<Self>, handle: &Runtime, meta: SpawnMeta) -> CompioJoinHandle<()> {
56        let Concrete { callback, func } = *self;
57        handle.spawn_at(
58            async move {
59                let res = func().await;
60                callback.send(res).ok();
61            },
62            meta,
63        )
64    }
65}
66
67impl<F, R> Dispatchable for Concrete<F, R>
68where
69    F: FnOnce() -> R + Send + 'static,
70    R: Send + 'static,
71{
72    fn run(self: Box<Self>) {
73        let Concrete { callback, func } = *self;
74        let res = func();
75        callback.send(res).ok();
76    }
77}
78
79/// The dispatcher. It manages the threads and dispatches the tasks.
80#[derive(Debug)]
81pub struct Dispatcher {
82    sender: Sender<Spawning>,
83    threads: Vec<JoinHandle<()>>,
84    pool: AsyncifyPool,
85}
86
87impl Dispatcher {
88    /// Create the dispatcher with specified number of threads.
89    #[track_caller]
90    pub(crate) fn new_impl(builder: DispatcherBuilder) -> io::Result<Self> {
91        let DispatcherBuilder {
92            nthreads,
93            concurrent,
94            stack_size,
95            mut thread_affinity,
96            mut names,
97            mut proactor_builder,
98        } = builder;
99        proactor_builder.force_reuse_thread_pool();
100        let pool = proactor_builder.create_or_get_thread_pool();
101        let (sender, receiver) = unbounded::<Spawning>();
102        // Captured out here, since `#[track_caller]` does not reach into the
103        // closures the threads run, and every worker belongs to this call.
104        let meta = SpawnMeta::capture().named("dispatcher::worker");
105
106        let threads = (0..nthreads)
107            .map({
108                |index| {
109                    let proactor_builder = proactor_builder.clone();
110                    let receiver = receiver.clone();
111
112                    let thread_builder = std::thread::Builder::new();
113                    let thread_builder = if let Some(s) = stack_size {
114                        thread_builder.stack_size(s)
115                    } else {
116                        thread_builder
117                    };
118                    let thread_builder = if let Some(f) = &mut names {
119                        thread_builder.name(f(index))
120                    } else {
121                        thread_builder
122                    };
123
124                    let cpus = if let Some(f) = &mut thread_affinity {
125                        f(index)
126                    } else {
127                        HashSet::new()
128                    };
129                    thread_builder.spawn(move || {
130                        Runtime::builder()
131                            .with_proactor(proactor_builder)
132                            .thread_affinity(cpus)
133                            .build()
134                            .expect("cannot create compio runtime")
135                            .block_on_at(
136                                async move {
137                                    while let Ok(Spawning { task: f, meta }) =
138                                        receiver.recv_async().await
139                                    {
140                                        let task = Runtime::with_current(|rt| f.spawn(rt, meta));
141                                        if concurrent {
142                                            task.detach()
143                                        } else {
144                                            task.await.ok();
145                                        }
146                                    }
147                                },
148                                meta,
149                            );
150                    })
151                }
152            })
153            .collect::<io::Result<Vec<_>>>()?;
154
155        Ok(Self {
156            sender,
157            threads,
158            pool,
159        })
160    }
161
162    /// Create the dispatcher with default config.
163    #[track_caller]
164    pub fn new() -> io::Result<Self> {
165        Self::builder().build()
166    }
167
168    /// Create a builder to build a dispatcher.
169    pub fn builder() -> DispatcherBuilder {
170        DispatcherBuilder::default()
171    }
172
173    /// Dispatch a task to the threads
174    ///
175    /// The provided `f` should be [`Send`] because it will be send to another
176    /// thread before calling. The returned [`Future`] need not to be [`Send`]
177    /// because it will be executed on only one thread.
178    ///
179    /// # Error
180    ///
181    /// If all threads have panicked, this method will return an error with the
182    /// sent closure.
183    #[track_caller]
184    pub fn dispatch<F, R>(&self, f: F) -> Result<oneshot::Receiver<R>, DispatchError<F>>
185    where
186        F: AsyncFnOnce() -> R + Send + 'static,
187        R: Send + 'static,
188    {
189        let (concrete, rx) = Concrete::new(f);
190
191        let meta = SpawnMeta::capture().named("dispatch");
192        match self.sender.send(Spawning {
193            task: Box::new(concrete),
194            meta,
195        }) {
196            Ok(_) => Ok(rx),
197            Err(err) => {
198                // SAFETY: We know the dispatchable we sent has type
199                // `Concrete<F, R>`
200                let recovered =
201                    unsafe { Box::from_raw(Box::into_raw(err.0.task) as *mut Concrete<F, R>) };
202                Err(DispatchError(recovered.func))
203            }
204        }
205    }
206
207    /// Dispatch a blocking task to the threads.
208    ///
209    /// Blocking pool of the dispatcher will be obtained from the proactor
210    /// builder. So any configuration of the proactor's blocking pool will be
211    /// applied to the dispatcher.
212    ///
213    /// # Error
214    ///
215    /// If all threads are busy and the thread pool is full, this method will
216    /// return an error with the original closure. The limit can be configured
217    /// with [`DispatcherBuilder::proactor_builder`] and
218    /// [`ProactorBuilder::thread_pool_limit`].
219    pub fn dispatch_blocking<Fn, R>(&self, f: Fn) -> Result<oneshot::Receiver<R>, DispatchError<Fn>>
220    where
221        Fn: FnOnce() -> R + Send + 'static,
222        R: Send + 'static,
223    {
224        let (concrete, rx) = Concrete::new(f);
225
226        self.pool
227            .dispatch(concrete)
228            .map_err(|e| DispatchError(e.0.func))?;
229
230        Ok(rx)
231    }
232
233    /// Stop the dispatcher and wait for the threads to complete. If there is a
234    /// thread panicked, this method will resume the panic.
235    pub async fn join(self) -> io::Result<()> {
236        drop(self.sender);
237        let (tx, rx) = oneshot::channel::<Vec<_>>();
238        if let Err(f) = self.pool.dispatch({
239            move || {
240                let results = self
241                    .threads
242                    .into_iter()
243                    .map(|thread| thread.join())
244                    .collect();
245                tx.send(results).ok();
246            }
247        }) {
248            std::thread::spawn(f.0);
249        }
250        let results = rx
251            .await
252            .map_err(|_| io::Error::other("the join task cancelled unexpectedly"))?;
253        for res in results {
254            res.unwrap_or_else(|e| resume_unwind(e));
255        }
256        Ok(())
257    }
258}
259
260/// A builder for [`Dispatcher`].
261pub struct DispatcherBuilder {
262    nthreads: usize,
263    concurrent: bool,
264    stack_size: Option<usize>,
265    thread_affinity: Option<Box<dyn FnMut(usize) -> HashSet<usize>>>,
266    names: Option<Box<dyn FnMut(usize) -> String>>,
267    proactor_builder: ProactorBuilder,
268}
269
270impl DispatcherBuilder {
271    /// Create a builder with default settings.
272    pub fn new() -> Self {
273        Self {
274            nthreads: available_parallelism().map(|n| n.get()).unwrap_or(1),
275            concurrent: true,
276            stack_size: None,
277            thread_affinity: None,
278            names: None,
279            proactor_builder: ProactorBuilder::new(),
280        }
281    }
282
283    /// If execute tasks concurrently. Default to be `true`.
284    ///
285    /// When set to `false`, tasks are executed sequentially without any
286    /// concurrency within the thread.
287    pub fn concurrent(mut self, concurrent: bool) -> Self {
288        self.concurrent = concurrent;
289        self
290    }
291
292    /// Set the number of worker threads of the dispatcher. The default value is
293    /// the CPU number. If the CPU number could not be retrieved, the
294    /// default value is 1.
295    pub fn worker_threads(mut self, nthreads: NonZeroUsize) -> Self {
296        self.nthreads = nthreads.get();
297        self
298    }
299
300    /// Set the size of stack of the worker threads.
301    pub fn stack_size(mut self, s: usize) -> Self {
302        self.stack_size = Some(s);
303        self
304    }
305
306    /// Set the thread affinity for the dispatcher.
307    pub fn thread_affinity(mut self, f: impl FnMut(usize) -> HashSet<usize> + 'static) -> Self {
308        self.thread_affinity = Some(Box::new(f));
309        self
310    }
311
312    /// Provide a function to assign names to the worker threads.
313    pub fn thread_names(mut self, f: impl (FnMut(usize) -> String) + 'static) -> Self {
314        self.names = Some(Box::new(f) as _);
315        self
316    }
317
318    /// Set the proactor builder for the inner runtimes.
319    pub fn proactor_builder(mut self, builder: ProactorBuilder) -> Self {
320        self.proactor_builder = builder;
321        self
322    }
323
324    /// Build the [`Dispatcher`].
325    #[track_caller]
326    pub fn build(self) -> io::Result<Dispatcher> {
327        Dispatcher::new_impl(self)
328    }
329}
330
331impl Default for DispatcherBuilder {
332    fn default() -> Self {
333        Self::new()
334    }
335}
compio_dispatcher/lib.rs

compio_dispatcher/
lib.rs
