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

compio_dispatcher/
lib.rs
