rocketmq_common/
thread_pool.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp;
16use std::future::Future;
17use std::sync::atomic::AtomicUsize;
18use std::sync::atomic::Ordering;
19use std::time::Duration;
20
21use anyhow::bail;
22use anyhow::Context;
23use tokio::runtime::Handle;
24use tokio::task::JoinHandle;
25
26pub struct TokioExecutorService {
27    inner: tokio::runtime::Runtime,
28}
29
30impl Default for TokioExecutorService {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl TokioExecutorService {
37    pub fn shutdown(self) {
38        self.inner.shutdown_background();
39    }
40
41    pub fn shutdown_timeout(self, timeout: Duration) {
42        self.inner.shutdown_timeout(timeout);
43    }
44}
45
46impl TokioExecutorService {
47    pub fn new() -> TokioExecutorService {
48        Self::try_new().unwrap_or_else(|error| panic!("failed to create TokioExecutorService: {error:#}"))
49    }
50
51    pub fn try_new() -> anyhow::Result<TokioExecutorService> {
52        Ok(TokioExecutorService {
53            inner: tokio::runtime::Builder::new_multi_thread()
54                .worker_threads(num_cpus::get())
55                .enable_all()
56                .build()
57                .context("failed to create TokioExecutorService runtime")?,
58        })
59    }
60
61    pub fn new_with_config(
62        thread_num: usize,
63        thread_prefix: Option<impl Into<String>>,
64        keep_alive: Duration,
65        max_blocking_threads: usize,
66    ) -> TokioExecutorService {
67        Self::try_new_with_config(thread_num, thread_prefix, keep_alive, max_blocking_threads)
68            .unwrap_or_else(|error| panic!("failed to create TokioExecutorService: {error:#}"))
69    }
70
71    pub fn try_new_with_config(
72        thread_num: usize,
73        thread_prefix: Option<impl Into<String>>,
74        keep_alive: Duration,
75        max_blocking_threads: usize,
76    ) -> anyhow::Result<TokioExecutorService> {
77        validate_runtime_config(thread_num, max_blocking_threads)?;
78        let thread_prefix_inner = if let Some(thread_prefix) = thread_prefix {
79            thread_prefix.into()
80        } else {
81            "rocketmq-thread-".to_string()
82        };
83        Ok(TokioExecutorService {
84            inner: tokio::runtime::Builder::new_multi_thread()
85                .worker_threads(thread_num)
86                .thread_keep_alive(keep_alive)
87                .max_blocking_threads(max_blocking_threads)
88                .thread_name_fn(move || {
89                    static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
90                    let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
91                    format!("{thread_prefix_inner}{id}")
92                })
93                .enable_all()
94                .build()
95                .context("failed to create configured TokioExecutorService runtime")?,
96        })
97    }
98}
99
100impl TokioExecutorService {
101    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
102    where
103        F: Future + Send + 'static,
104        F::Output: Send + 'static,
105    {
106        self.inner.spawn(future)
107    }
108
109    pub fn get_handle(&self) -> &Handle {
110        self.inner.handle()
111    }
112
113    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
114        self.inner.block_on(future)
115    }
116}
117
118pub struct FuturesExecutorService {
119    inner: futures::executor::ThreadPool,
120}
121impl FuturesExecutorService {
122    pub fn spawn<F>(&self, future: F)
123    where
124        F: Future<Output = ()> + Send + 'static,
125    {
126        self.inner.spawn_ok(future);
127    }
128}
129
130#[derive(Debug, Default)]
131pub struct FuturesExecutorServiceBuilder {
132    pool_size: usize,
133    stack_size: usize,
134    thread_name_prefix: Option<String>,
135}
136
137impl FuturesExecutorServiceBuilder {
138    pub fn new() -> FuturesExecutorServiceBuilder {
139        FuturesExecutorServiceBuilder {
140            pool_size: cmp::max(1, num_cpus::get()),
141            stack_size: 0,
142            thread_name_prefix: None,
143        }
144    }
145
146    pub fn pool_size(mut self, pool_size: usize) -> Self {
147        self.pool_size = pool_size;
148        self
149    }
150
151    pub fn stack_size(mut self, stack_size: usize) -> Self {
152        self.stack_size = stack_size;
153        self
154    }
155
156    pub fn create(&mut self) -> anyhow::Result<FuturesExecutorService> {
157        let name_prefix = self.thread_name_prefix.as_deref().unwrap_or("Default-Executor");
158        let thread_pool = futures::executor::ThreadPool::builder()
159            .stack_size(self.stack_size)
160            .pool_size(self.pool_size)
161            .name_prefix(name_prefix)
162            .create()
163            .context("failed to create futures executor thread pool")?;
164        Ok(FuturesExecutorService { inner: thread_pool })
165    }
166}
167
168pub struct ScheduledExecutorService {
169    inner: tokio::runtime::Runtime,
170}
171
172impl Default for ScheduledExecutorService {
173    fn default() -> Self {
174        Self::new()
175    }
176}
177impl ScheduledExecutorService {
178    pub fn new() -> ScheduledExecutorService {
179        Self::try_new().unwrap_or_else(|error| panic!("failed to create ScheduledExecutorService: {error:#}"))
180    }
181
182    pub fn try_new() -> anyhow::Result<ScheduledExecutorService> {
183        Ok(ScheduledExecutorService {
184            inner: tokio::runtime::Builder::new_multi_thread()
185                .worker_threads(num_cpus::get())
186                .enable_all()
187                .build()
188                .context("failed to create ScheduledExecutorService runtime")?,
189        })
190    }
191
192    pub fn new_with_config(
193        thread_num: usize,
194        thread_prefix: Option<impl Into<String>>,
195        keep_alive: Duration,
196        max_blocking_threads: usize,
197    ) -> ScheduledExecutorService {
198        Self::try_new_with_config(thread_num, thread_prefix, keep_alive, max_blocking_threads)
199            .unwrap_or_else(|error| panic!("failed to create ScheduledExecutorService: {error:#}"))
200    }
201
202    pub fn try_new_with_config(
203        thread_num: usize,
204        thread_prefix: Option<impl Into<String>>,
205        keep_alive: Duration,
206        max_blocking_threads: usize,
207    ) -> anyhow::Result<ScheduledExecutorService> {
208        validate_runtime_config(thread_num, max_blocking_threads)?;
209        let thread_prefix_inner = if let Some(thread_prefix) = thread_prefix {
210            thread_prefix.into()
211        } else {
212            "rocketmq-thread-".to_string()
213        };
214        Ok(ScheduledExecutorService {
215            inner: tokio::runtime::Builder::new_multi_thread()
216                .worker_threads(thread_num)
217                .thread_keep_alive(keep_alive)
218                .max_blocking_threads(max_blocking_threads)
219                .thread_name_fn(move || {
220                    static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
221                    let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
222                    format!("{thread_prefix_inner}{id}")
223                })
224                .enable_all()
225                .build()
226                .context("failed to create configured ScheduledExecutorService runtime")?,
227        })
228    }
229
230    pub fn schedule_at_fixed_rate<F>(&self, mut task: F, initial_delay: Option<Duration>, period: Duration)
231    where
232        F: FnMut() + Send + 'static,
233    {
234        self.inner.spawn(async move {
235            // initial delay
236
237            if let Some(initial_delay_inner) = initial_delay {
238                tokio::time::sleep(initial_delay_inner).await;
239            }
240
241            loop {
242                // record current execution time
243                let current_execution_time = tokio::time::Instant::now();
244                // execute task
245                task();
246                // Calculate the time of the next execution
247                let next_execution_time = current_execution_time + period;
248
249                // Wait until the next execution
250                let delay = next_execution_time.saturating_duration_since(tokio::time::Instant::now());
251                tokio::time::sleep(delay).await;
252            }
253        });
254    }
255}
256
257fn validate_runtime_config(thread_num: usize, max_blocking_threads: usize) -> anyhow::Result<()> {
258    if thread_num == 0 {
259        bail!("thread_num must be greater than zero");
260    }
261    if max_blocking_threads == 0 {
262        bail!("max_blocking_threads must be greater than zero");
263    }
264    Ok(())
265}
266
267#[cfg(test)]
268mod runtime_config_tests {
269    use super::*;
270
271    #[test]
272    fn tokio_executor_try_new_with_config_rejects_invalid_thread_counts() {
273        let error = match TokioExecutorService::try_new_with_config(0, Some("test-"), Duration::from_secs(1), 1) {
274            Ok(_) => panic!("zero worker threads should be rejected before tokio builder panics"),
275            Err(error) => error,
276        };
277        assert!(error.to_string().contains("thread_num must be greater than zero"));
278
279        let error = match TokioExecutorService::try_new_with_config(1, Some("test-"), Duration::from_secs(1), 0) {
280            Ok(_) => panic!("zero blocking threads should be rejected before tokio builder panics"),
281            Err(error) => error,
282        };
283        assert!(error
284            .to_string()
285            .contains("max_blocking_threads must be greater than zero"));
286    }
287
288    #[test]
289    fn scheduled_executor_try_new_with_config_rejects_invalid_thread_counts() {
290        let error = match ScheduledExecutorService::try_new_with_config(0, Some("test-"), Duration::from_secs(1), 1) {
291            Ok(_) => panic!("zero worker threads should be rejected before tokio builder panics"),
292            Err(error) => error,
293        };
294        assert!(error.to_string().contains("thread_num must be greater than zero"));
295
296        let error = match ScheduledExecutorService::try_new_with_config(1, Some("test-"), Duration::from_secs(1), 0) {
297            Ok(_) => panic!("zero blocking threads should be rejected before tokio builder panics"),
298            Err(error) => error,
299        };
300        assert!(error
301            .to_string()
302            .contains("max_blocking_threads must be greater than zero"));
303    }
304}
305
306impl ScheduledExecutorService {
307    pub fn shutdown(self) {
308        self.inner.shutdown_background();
309    }
310
311    pub fn shutdown_timeout(self, timeout: Duration) {
312        self.inner.shutdown_timeout(timeout);
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use std::sync::mpsc;
319    use std::time::Duration;
320
321    use super::FuturesExecutorServiceBuilder;
322
323    #[test]
324    fn futures_executor_builder_create_returns_executor() {
325        let mut builder = FuturesExecutorServiceBuilder::new().pool_size(1).stack_size(0);
326        let executor = builder.create().expect("executor should be created");
327        let (tx, rx) = mpsc::channel();
328
329        executor.spawn(async move {
330            tx.send(42).expect("test receiver should be alive");
331        });
332
333        assert_eq!(rx.recv_timeout(Duration::from_secs(5)).unwrap(), 42);
334    }
335}
rocketmq_common/thread_pool.rs

rocketmq_common/
thread_pool.rs
