relay_threading/
builder.rs

1use std::any::Any;
2use std::future::Future;
3use std::io;
4use std::sync::Arc;
5
6use crate::pool::{AsyncPool, Thread};
7use crate::pool::{CustomSpawn, DefaultSpawn, ThreadSpawn};
8
9/// Type alias for a thread safe closure that is used for panic handling across the code.
10pub(crate) type PanicHandler = dyn Fn(Box<dyn Any + Send>) + Send + Sync;
11
12/// [`AsyncPoolBuilder`] provides a flexible way to configure and build an [`AsyncPool`] for executing
13/// asynchronous tasks concurrently on dedicated threads.
14///
15/// This builder enables you to customize the number of threads, concurrency limits, thread naming,
16/// and panic handling strategies.
17pub struct AsyncPoolBuilder<S = DefaultSpawn> {
18    pub(crate) runtime: tokio::runtime::Handle,
19    pub(crate) pool_name: Option<&'static str>,
20    pub(crate) thread_name: Option<Box<dyn FnMut(usize) -> String>>,
21    pub(crate) thread_panic_handler: Option<Arc<PanicHandler>>,
22    pub(crate) task_panic_handler: Option<Arc<PanicHandler>>,
23    pub(crate) spawn_handler: S,
24    pub(crate) num_threads: usize,
25    pub(crate) max_concurrency: usize,
26}
27
28impl AsyncPoolBuilder<DefaultSpawn> {
29    /// Initializes a new [`AsyncPoolBuilder`] with default settings.
30    ///
31    /// The builder is tied to the provided [`tokio::runtime::Handle`] and prepares to configure an [`AsyncPool`].
32    pub fn new(runtime: tokio::runtime::Handle) -> AsyncPoolBuilder<DefaultSpawn> {
33        AsyncPoolBuilder {
34            runtime,
35            pool_name: None,
36            thread_name: None,
37            thread_panic_handler: None,
38            task_panic_handler: None,
39            spawn_handler: DefaultSpawn,
40            num_threads: 1,
41            max_concurrency: 1,
42        }
43    }
44}
45
46impl<S> AsyncPoolBuilder<S>
47where
48    S: ThreadSpawn,
49{
50    /// Specifies a custom name for this pool.
51    pub fn pool_name(mut self, pool_name: &'static str) -> Self {
52        self.pool_name = Some(pool_name);
53        self
54    }
55
56    /// Specifies a custom naming convention for threads in the [`AsyncPool`].
57    ///
58    /// The provided closure receives the thread's index and returns a name,
59    /// which can be useful for debugging and logging.
60    pub fn thread_name<F>(mut self, thread_name: F) -> Self
61    where
62        F: FnMut(usize) -> String + 'static,
63    {
64        self.thread_name = Some(Box::new(thread_name));
65        self
66    }
67
68    /// Sets a custom panic handler for threads in the [`AsyncPool`].
69    ///
70    /// If a thread panics, the provided handler will be invoked so that you can perform
71    /// custom error handling or cleanup.
72    pub fn thread_panic_handler<F>(mut self, panic_handler: F) -> Self
73    where
74        F: Fn(Box<dyn Any + Send>) + Send + Sync + 'static,
75    {
76        self.thread_panic_handler = Some(Arc::new(panic_handler));
77        self
78    }
79
80    /// Sets a custom panic handler for tasks executed by the [`AsyncPool`].
81    ///
82    /// This handler is used to manage panics that occur during task execution, allowing for graceful
83    /// error handling.
84    pub fn task_panic_handler<F>(mut self, panic_handler: F) -> Self
85    where
86        F: Fn(Box<dyn Any + Send>) + Send + Sync + 'static,
87    {
88        self.task_panic_handler = Some(Arc::new(panic_handler));
89        self
90    }
91
92    /// Configures a custom thread spawning procedure for the [`AsyncPool`].
93    ///
94    /// This method allows you to adjust thread settings (e.g. naming, stack size) before thread creation,
95    /// making it possible to apply application-specific configurations.
96    pub fn spawn_handler<F>(self, spawn_handler: F) -> AsyncPoolBuilder<CustomSpawn<F>>
97    where
98        F: FnMut(Thread) -> io::Result<()>,
99    {
100        AsyncPoolBuilder {
101            runtime: self.runtime,
102            pool_name: self.pool_name,
103            thread_name: self.thread_name,
104            thread_panic_handler: self.thread_panic_handler,
105            task_panic_handler: self.task_panic_handler,
106            spawn_handler: CustomSpawn::new(spawn_handler),
107            num_threads: self.num_threads,
108            max_concurrency: self.max_concurrency,
109        }
110    }
111
112    /// Sets the number of worker threads for the [`AsyncPool`].
113    ///
114    /// This determines how many dedicated threads will be available for running tasks concurrently.
115    pub fn num_threads(mut self, num_threads: usize) -> Self {
116        self.num_threads = num_threads;
117        self
118    }
119
120    /// Sets the maximum number of concurrent tasks per thread in the [`AsyncPool`].
121    ///
122    /// This controls how many futures can be polled simultaneously on each worker thread.
123    pub fn max_concurrency(mut self, max_concurrency: usize) -> Self {
124        self.max_concurrency = max_concurrency;
125        self
126    }
127
128    /// Constructs an [`AsyncPool`] based on the configured settings.
129    ///
130    /// Finalizing the builder sets up dedicated worker threads and configures the executor
131    /// to enforce the specified concurrency limits.
132    pub fn build<F>(self) -> Result<AsyncPool<F>, io::Error>
133    where
134        F: Future<Output = ()> + Send + 'static,
135    {
136        AsyncPool::new(self)
137    }
138}