workerbee/
builder.rs

1/// Builder for create worker instance.
2
3
4use std::{io, thread};
5use std::time::Duration;
6use std::sync::Arc;
7
8use tokio::runtime::{Builder as RuntimeBuilder, Handle};
9use tokio::net::{TcpListener, TcpStream};
10use tokio::sync::oneshot;
11
12use crate::worker::Worker;
13use crate::dropper::Dropper;
14
15pub struct Builder {
16    /// The id of this worker
17    id: String,
18
19    /// worker builder
20    builder: RuntimeBuilder,
21}
22
23impl Default for Builder {
24    fn default() -> Self {
25        Self {
26            id: "default-worker".to_string(),
27            builder: RuntimeBuilder::new_multi_thread(),
28        }
29    }
30}
31
32impl Builder {
33    /// Sets the number of worker threads the Runtime will use.
34    ///
35    /// This can be any number above 0. The default value is the number of cores available to the system.
36    pub fn worker_threads(&mut self, val: usize) -> &mut Self {
37        self.builder.worker_threads(val);
38        self
39    }
40
41    /// Specifies the limit for additional threads spawned by the Runtime.
42    ///
43    /// These threads are used for blocking operations like tasks spawned through spawn_blocking,
44    /// they are not always active and will exit if left idle for too long, You can change this timeout duration
45    /// with thread_keep_alive. The default value is 512.
46    pub fn max_blocking_threads(&mut self, val: usize) -> &mut Self {
47        self.builder.max_blocking_threads(val);
48        self
49    }
50
51    /// Sets a custom timeout for a thread in the blocking pool.
52    ///
53    /// By default, the timeout for a thread is set to 10 seconds.
54    pub fn thread_keep_alive(&mut self, duration: Duration) -> &mut Self {
55        self.builder.thread_keep_alive(duration);
56        self
57    }
58
59    /// Sets name of threads spawned by the Runtime thread pool
60    pub fn thread_name(&mut self, val: impl Into<String>) -> &mut Self {
61        self.id = val.into();
62        self
63    }
64
65    pub fn build(&mut self) -> Result<Worker, io::Error> {
66        let worker = self
67            .builder
68            .enable_all()
69            .thread_name(self.id.clone())
70            .build()?;
71
72        let handle = worker.handle().clone();
73        let (send_stop, recv_stop) = oneshot::channel();
74        // Block the runtime to shutdown.
75        let _ = thread::Builder::new()
76            .name(format!("{}-blocker", self.id))
77            .spawn(move || worker.block_on(recv_stop));
78
79        Ok(Worker {
80            handle,
81            _dropper: Arc::new(Dropper {
82                close: Some(send_stop),
83            }),
84        })
85    }
86}