lwleen_rpc/mod_pool.rs
1//! ## 来自开源库 rusty_pool `Apache-2.0`
2//! - <https://crates.io/crates/rusty_pool> 0.7.0
3//! - 代码 <https://docs.rs/rusty_pool/latest/src/rusty_pool/lib.rs.html>
4//!
5//!
6//!
7//! Future 特征,即 poll(self,ctx) 函数
8//! 详细 <https://course.rs/advance/async/future-excuting.html>
9//!
10//! 其中 ctx 中,由执行环境来提供 waker(唤醒方法),例如下面(回调转异步_Future)回调后,触发 oneshot::receiver 中的唤醒
11//!
12//! ----------------------------------------------------
13//! 【执行环境】有个就绪Receiver通道,收到task(Fut)后就执行poll方法( 提供ctx中waker方法 )
14//! 1. task(Fut)新建时发送就绪Receiver通道(run方法运行), 执行poll中函数一次,获得(ctx)waker方法
15//! 2. 函数后续完成后,例如下例中回调触发 waker方法, 将task(Fut)再次发送到就绪Receiver通道(run方法运行)
16//!
17//!
18//! ```no_run
19//! // 这个是 rusty_pool 用自定义执行异步环境run方法,自定义waker
20//! // 普通闭包调用一次 异步(闭包)是特殊的闭包,调用poll即执行这个异步(闭包)
21//! struct AsyncTask {
22//! future: Mutex<Option<BoxFuture<'static, ()>>>,
23//! pool: ThreadPool,
24//! }
25//!
26//!
27//! //这个相当于 --- 自定义的执行环境(waker唤醒方法) ---
28//! //替换waker方法将 task 再次发送到就绪通道
29//! #[cfg(feature = "async")]
30//! impl ArcWake for AsyncTask {
31//! fn wake_by_ref(arc_self: &Arc<Self>) {
32//! let cloned_task = arc_self.clone();
33//! //waker再次发送任务,会再次 poll异步,即唤醒任务再次执行Fut的poll
34//! arc_self
35//! .pool
36//! .try_execute(cloned_task)
37//! .expect("failed to wake future because message could not be sent to pool");
38//! }
39//! }
40//!
41//! // - 🍁 run 中第一次执行poll后, 让出线程, 等待回调再次启用一个线程执行此任务
42//! //这个相当于 --- 自定义的执行环境(就绪Receiver通道) run方法 执行Fut的poll ---
43//! //调用Fut中的 poll 推动Fut运行poll
44//! impl Task<()> for Arc<AsyncTask> {
45//! fn run(self) {
46//! let mut future_slot = self.future.lock().expect("failed to acquire mutex");
47//! if let Some(mut future) = future_slot.take() { //处理时取走任务
48//! let waker = waker_ref(&self);
49//! let context = &mut Context::from_waker(&*waker); //提供 waker
50//! if future.as_mut().poll(context).is_pending() { //等待时任务继续存在
51//! *future_slot = Some(future);
52//! }
53//! }
54//! }
55//! }
56//!
57//! // `self`, `&self`, `&mut self`, `self: Box<Self>`, `self: Rc<Self>`, `self: Arc<Self>`, or `self: Pin<P>` (where P is one of the previous types except `Self`)
58//! ```
59#![allow(dead_code)]
60
61use futures::{ channel::oneshot };
62use futures_lite::{ future::block_on };
63use crossbeam_channel;
64
65use crate::prelude::全局_CPU数量;
66
67
68use futures::{
69 future::BoxFuture,
70 task::{waker_ref, ArcWake},
71};
72use std::future::Future;
73use std::option::Option;
74use std::sync::{
75 atomic::{AtomicUsize, Ordering},
76 Arc, Condvar, Mutex,
77};
78
79use std::task::Context;
80use std::thread;
81use std::time::Duration;
82
83const BITS: usize = std::mem::size_of::<usize>() * 8;
84/// The absolute maximum number of workers. This corresponds to the maximum value that can be stored within half the bits of usize,
85/// as two counters (total workers and idle workers) are stored in one AtomicUsize.
86pub const MAX_SIZE: usize = (1 << (BITS / 2)) - 1;
87
88type Job = Box<dyn FnOnce() + Send + 'static>;
89
90/// Trait to implement for all items that may be executed by the `ThreadPool`.
91pub trait Task<R: Send>: Send {
92 /// Execute this task and return its result.
93 fn run(self) -> R;
94
95 /// Transform this `Task` into a heap allocated `FnOnce` if possible.
96 ///
97 /// Used by [`ThreadPool::execute`](struct.ThreadPool.html#method.execute) to turn this `Task` into a `Job`
98 /// directly without having to create an additional `Job` that calls this `Task`.
99 fn into_fn(self) -> Option<Box<dyn FnOnce() -> R + Send + 'static>>;
100
101 /// Return `true` if calling [`Task::into_fn`] on this `Task` returns `Some`.
102 fn is_fn(&self) -> bool;
103}
104
105/// Implement the `Task` trait for any FnOnce closure that returns a thread-safe result.
106impl<R, F> Task<R> for F
107where
108 R: Send,
109 F: FnOnce() -> R + Send + 'static,
110{
111 fn run(self) -> R {
112 self()
113 }
114
115 fn into_fn(self) -> Option<Box<dyn FnOnce() -> R + Send + 'static>> {
116 Some(Box::new(self))
117 }
118
119 fn is_fn(&self) -> bool {
120 true
121 }
122}
123
124/// Handle returned by [`ThreadPool::evaluate`](struct.ThreadPool.html#method.evaluate) and [`ThreadPool::complete`](struct.ThreadPool.html#method.complete)
125/// that allows to block the current thread and wait for the result of a submitted task. The returned `JoinHandle` may also be sent to the [`ThreadPool`](struct.ThreadPool.html)
126/// to create a task that blocks a worker thread until the task is completed and then does something with the result. This handle communicates with the worker thread
127/// using a oneshot channel blocking the thread when [`try_await_complete()`](struct.JoinHandle.html#method.try_await_complete) is called until a message, i.e. the result of the
128/// task, is received.
129pub struct JoinHandle<T: Send> {
130 pub receiver: oneshot::Receiver<T>,
131}
132
133impl<T: Send> JoinHandle<T> {
134 /// Block the current thread until the result of the task is received.
135 ///
136 /// # Errors
137 ///
138 /// This function might return a `oneshot::Canceled` if the channel was broken
139 /// before the result was received. This is generally the case if execution of
140 /// the task panicked.
141 pub fn try_await_complete(self) -> Result<T, oneshot::Canceled> {
142 block_on(self.receiver)
143 }
144
145 /// Block the current thread until the result of the task is received.
146 ///
147 /// # Panics
148 ///
149 /// This function might panic if [`try_await_complete()`](struct.JoinHandle.html#method.try_await_complete) returns `oneshot::Canceled`.
150 /// This is generally the case if execution of the task panicked and the sender was dropped before sending a result to the receiver.
151 pub fn await_complete(self) -> T {
152 self.try_await_complete()
153 .expect("could not receive message because channel was cancelled")
154 }
155}
156
157
158struct AsyncTask {
159 future: Mutex<Option<BoxFuture<'static, ()>>>,
160 pool: ThreadPool,
161}
162
163/// Implement `ArcWake` for `AsyncTask` by re-submitting the `AsyncTask` i.e. the `Future` to the pool.
164
165impl ArcWake for AsyncTask {
166 fn wake_by_ref(arc_self: &Arc<Self>) {
167 let cloned_task = arc_self.clone();
168 arc_self
169 .pool
170 .try_execute(cloned_task)
171 .expect("failed to wake future because message could not be sent to pool");
172 }
173}
174
175/// Implement the `Task` trait for `AsyncTask` in order to make it executable for the pool by
176/// creating a waker and polling the future.
177
178impl Task<()> for Arc<AsyncTask> {
179 fn run(self) {
180 let mut future_slot = self.future.lock().expect("failed to acquire mutex");
181 if let Some(mut future) = future_slot.take() {
182 let waker = waker_ref(&self);
183 let context = &mut Context::from_waker(&*waker);
184 if future.as_mut().poll(context).is_pending() {
185 *future_slot = Some(future);
186 }
187 }
188 }
189
190 fn into_fn(self) -> Option<Box<dyn FnOnce() + Send + 'static>> {
191 None
192 }
193
194 fn is_fn(&self) -> bool {
195 false
196 }
197}
198
199// assert that Send is implemented
200trait ThreadSafe: Send {}
201
202impl<R: Send> ThreadSafe for dyn Task<R> {}
203
204impl<R: Send> ThreadSafe for JoinHandle<R> {}
205
206impl ThreadSafe for ThreadPool {}
207
208/// Self growing / shrinking `ThreadPool` implementation based on crossbeam's
209/// multi-producer multi-consumer channels that enables awaiting the result of a
210/// task and offers async support.
211///
212/// This `ThreadPool` has two different pool sizes; a core pool size filled with
213/// threads that live for as long as the channel and a max pool size which describes
214/// the maximum amount of worker threads that may live at the same time.
215/// Those additional non-core threads have a specific keep_alive time described when
216/// creating the `ThreadPool` that defines how long such threads may be idle for
217/// without receiving any work before giving up and terminating their work loop.
218///
219/// This `ThreadPool` does not spawn any threads until a task is submitted to it.
220/// Then it will create a new thread for each task until the core pool size is full.
221/// After that a new thread will only be created upon an `execute()` call if the
222/// current pool is lower than the max pool size and there are no idle threads.
223///
224/// Functions like `evaluate()` and `complete()` return a `JoinHandle` that may be used
225/// to await the result of a submitted task or future. JoinHandles may be sent to the
226/// thread pool to create a task that blocks a worker thread until it receives the
227/// result of the other task and then operates on the result. If the task panics the
228/// `JoinHandle` receives a cancellation error. This is implemented using a futures
229/// oneshot channel to communicate with the worker thread.
230///
231/// This `ThreadPool` may be used as a futures executor if the "async" feature is enabled,
232/// which is the case by default. The "async" feature includes the `spawn()` and
233/// `try_spawn()` functions which create a task that polls the future one by one and
234/// creates a waker that re-submits the future to the pool when it can make progress.
235/// Without the "async" feature, futures can simply be executed to completion using
236/// the `complete` function, which simply blocks a worker thread until the future has
237/// been polled to completion.
238///
239/// The "async" feature can be disabled if not need by adding the following to your
240/// Cargo dependency:
241/// ```toml
242/// [dependencies.rusty_pool]
243/// default-features = false
244/// version = "*"
245/// ```
246///
247/// When creating a new worker this `ThreadPool` tries to increment the worker count
248/// using a compare-and-swap mechanism, if the increment fails because the total worker
249/// count has been incremented to the specified limit (the core_size when trying to
250/// create a core thread, else the max_size) by another thread, the pool tries to create
251/// a non-core worker instead (if previously trying to create a core worker and no idle
252/// worker exists) or sends the task to the channel instead. Panicking workers are always
253/// cloned and replaced.
254///
255/// Locks are only used for the join functions to lock the `Condvar`, apart from that
256/// this `ThreadPool` implementation fully relies on crossbeam and atomic operations.
257/// This `ThreadPool` decides whether it is currently idle (and should fast-return
258/// join attempts) by comparing the total worker count to the idle worker count, which
259/// are two values stored in one `AtomicUsize` (both half the size of usize) making sure
260/// that if both are updated they may be updated in a single atomic operation.
261///
262/// The thread pool and its crossbeam channel can be destroyed by using the shutdown
263/// function, however that does not stop tasks that are already running but will
264/// terminate the thread the next time it will try to fetch work from the channel.
265/// The channel is only destroyed once all clones of the `ThreadPool` have been
266/// shut down / dropped.
267///
268/// # Usage
269/// Create a new `ThreadPool`:
270/// ```rust
271/// use rusty_pool::Builder;
272/// use rusty_pool::ThreadPool;
273/// // Create default `ThreadPool` configuration with the number of CPUs as core pool size
274/// let pool = ThreadPool::default();
275/// // Create a `ThreadPool` with default naming:
276/// use std::time::Duration;
277/// let pool2 = ThreadPool::new(5, 50, Duration::from_secs(60));
278/// // Create a `ThreadPool` with a custom name:
279/// let pool3 = ThreadPool::new_named(String::from("my_pool"), 5, 50, Duration::from_secs(60));
280/// // using the Builder struct:
281/// let pool4 = Builder::new().core_size(5).max_size(50).build();
282/// ```
283///
284/// Submit a closure for execution in the `ThreadPool`:
285/// ```rust
286/// use rusty_pool::ThreadPool;
287/// use std::thread;
288/// use std::time::Duration;
289/// let pool = ThreadPool::default();
290/// pool.execute(|| {
291/// thread::sleep(Duration::from_secs(5));
292/// print!("hello");
293/// });
294/// ```
295///
296/// Submit a task and await the result:
297/// ```rust
298/// use rusty_pool::ThreadPool;
299/// use std::thread;
300/// use std::time::Duration;
301/// let pool = ThreadPool::default();
302/// let handle = pool.evaluate(|| {
303/// thread::sleep(Duration::from_secs(5));
304/// return 4;
305/// });
306/// let result = handle.await_complete();
307/// assert_eq!(result, 4);
308/// ```
309///
310/// Spawn futures using the `ThreadPool`:
311/// ```rust
312/// async fn some_async_fn(x: i32, y: i32) -> i32 {
313/// x + y
314/// }
315///
316/// async fn other_async_fn(x: i32, y: i32) -> i32 {
317/// x - y
318/// }
319///
320/// use rusty_pool::ThreadPool;
321/// let pool = ThreadPool::default();
322///
323/// // simply complete future by blocking a worker until the future has been completed
324/// let handle = pool.complete(async {
325/// let a = some_async_fn(4, 6).await; // 10
326/// let b = some_async_fn(a, 3).await; // 13
327/// let c = other_async_fn(b, a).await; // 3
328/// some_async_fn(c, 5).await // 8
329/// });
330/// assert_eq!(handle.await_complete(), 8);
331///
332/// use std::sync::{Arc, atomic::{AtomicI32, Ordering}};
333///
334/// // spawn future and create waker that automatically re-submits itself to the threadpool if ready to make progress, this requires the "async" feature which is enabled by default
335/// let count = Arc::new(AtomicI32::new(0));
336/// let clone = count.clone();
337/// pool.spawn(async move {
338/// let a = some_async_fn(3, 6).await; // 9
339/// let b = other_async_fn(a, 4).await; // 5
340/// let c = some_async_fn(b, 7).await; // 12
341/// clone.fetch_add(c, Ordering::Relaxed);
342/// });
343/// pool.join();
344/// assert_eq!(count.load(Ordering::Relaxed), 12);
345/// ```
346///
347/// Join and shut down the `ThreadPool`:
348/// ```rust
349/// use std::thread;
350/// use std::time::Duration;
351/// use rusty_pool::ThreadPool;
352/// use std::sync::{Arc, atomic::{AtomicI32, Ordering}};
353///
354/// let pool = ThreadPool::default();
355/// for _ in 0..10 {
356/// pool.execute(|| { thread::sleep(Duration::from_secs(10)) })
357/// }
358/// // wait for all threads to become idle, i.e. all tasks to be completed including tasks added by other threads after join() is called by this thread or for the timeout to be reached
359/// pool.join_timeout(Duration::from_secs(5));
360///
361/// let count = Arc::new(AtomicI32::new(0));
362/// for _ in 0..15 {
363/// let clone = count.clone();
364/// pool.execute(move || {
365/// thread::sleep(Duration::from_secs(5));
366/// clone.fetch_add(1, Ordering::Relaxed);
367/// });
368/// }
369///
370/// // shut down and drop the only instance of this `ThreadPool` (no clones) causing the channel to be broken leading all workers to exit after completing their current work
371/// // and wait for all workers to become idle, i.e. finish their work.
372/// pool.shutdown_join();
373/// assert_eq!(count.load(Ordering::Relaxed), 15);
374/// ```
375#[derive(Clone)]
376pub struct ThreadPool {
377 core_size: usize,
378 max_size: usize,
379 keep_alive: Duration,
380 channel_data: Arc<ChannelData>,
381 worker_data: Arc<WorkerData>,
382}
383
384impl ThreadPool {
385 /// Construct a new `ThreadPool` with the specified core pool size, max pool size
386 /// and keep_alive time for non-core threads. This function does not spawn any
387 /// threads. This `ThreadPool` will receive a default name in the following format:
388 /// "rusty_pool_" + pool number.
389 ///
390 /// `core_size` specifies the amount of threads to keep alive for as long as
391 /// the `ThreadPool` exists and its channel remains connected.
392 ///
393 /// `max_size` specifies the maximum number of worker threads that may exist
394 /// at the same time.
395 ///
396 /// `keep_alive` specifies the duration for which to keep non-core pool
397 /// worker threads alive while they do not receive any work.
398 ///
399 /// # Panics
400 ///
401 /// This function will panic if max_size is 0, lower than core_size or exceeds half
402 /// the size of usize. This restriction exists because two counters (total workers and
403 /// idle counters) are stored within one AtomicUsize.
404 pub fn new(core_size: usize, max_size: usize, keep_alive: Duration) -> Self {
405 static POOL_COUNTER: AtomicUsize = AtomicUsize::new(1);
406 let name = format!(
407 "rusty_pool_{}",
408 POOL_COUNTER.fetch_add(1, Ordering::Relaxed)
409 );
410 ThreadPool::new_named(name, core_size, max_size, keep_alive)
411 }
412
413 /// Construct a new `ThreadPool` with the specified name, core pool size, max pool size
414 /// and keep_alive time for non-core threads. This function does not spawn any
415 /// threads.
416 ///
417 /// `name` the name of the `ThreadPool` that will be used as prefix for each
418 /// thread.
419 ///
420 /// `core_size` specifies the amount of threads to keep alive for as long as
421 /// the `ThreadPool` exists and its channel remains connected.
422 ///
423 /// `max_size` specifies the maximum number of worker threads that may exist
424 /// at the same time.
425 ///
426 /// `keep_alive` specifies the duration for which to keep non-core pool
427 /// worker threads alive while they do not receive any work.
428 ///
429 /// # Panics
430 ///
431 /// This function will panic if max_size is 0, lower than core_size or exceeds half
432 /// the size of usize. This restriction exists because two counters (total workers and
433 /// idle counters) are stored within one AtomicUsize.
434 pub fn new_named(
435 name: String,
436 core_size: usize,
437 max_size: usize,
438 keep_alive: Duration,
439 ) -> Self {
440 let (sender, receiver) = crossbeam_channel::unbounded();
441
442 if max_size == 0 || max_size < core_size {
443 panic!("max_size must be greater than 0 and greater or equal to the core pool size");
444 } else if max_size > MAX_SIZE {
445 panic!(
446 "max_size may not exceed {}, the maximum value that can be stored within half the bits of usize ({} -> {} bits in this case)",
447 MAX_SIZE,
448 BITS,
449 BITS / 2
450 );
451 }
452
453 let worker_data = WorkerData {
454 pool_name: name,
455 worker_count_data: WorkerCountData::default(),
456 worker_number: AtomicUsize::new(1),
457 join_notify_condvar: Condvar::new(),
458 join_notify_mutex: Mutex::new(()),
459 join_generation: AtomicUsize::new(0),
460 };
461
462 let channel_data = ChannelData { sender, receiver };
463
464 Self {
465 core_size,
466 max_size,
467 keep_alive,
468 channel_data: Arc::new(channel_data),
469 worker_data: Arc::new(worker_data),
470 }
471 }
472
473 /// Get the number of live workers, includes all workers waiting for work or executing tasks.
474 ///
475 /// This counter is incremented when creating a new worker. The value is increment just before
476 /// the worker starts executing its initial task. Incrementing the worker total might fail
477 /// if the total has already reached the specified limit (either core_size or max_size) after
478 /// being incremented by another thread, as of rusty_pool 0.5.0 failed attempts to create a worker
479 /// no longer skews the worker total as failed attempts to increment the worker total does not
480 /// increment the value at all.
481 /// This counter is decremented when a worker reaches the end of its working loop, which for non-core
482 /// threads might happen if it does not receive any work during its keep alive time,
483 /// for core threads this only happens once the channel is disconnected.
484 pub fn get_current_worker_count(&self) -> usize {
485 self.worker_data.worker_count_data.get_total_worker_count()
486 }
487
488 /// Get the number of workers currently waiting for work. Those threads are currently
489 /// polling from the crossbeam receiver. Core threads wait indefinitely and might remain
490 /// in this state until the `ThreadPool` is dropped. The remaining threads give up after
491 /// waiting for the specified keep_alive time.
492 pub fn get_idle_worker_count(&self) -> usize {
493 self.worker_data.worker_count_data.get_idle_worker_count()
494 }
495
496 /// Send a new task to the worker threads. This function is responsible for sending the message through the
497 /// channel and creating new workers if needed. If the current worker count is lower than the core pool size
498 /// this function will always create a new worker. If the current worker count is equal to or greater than
499 /// the core pool size this function only creates a new worker if the worker count is below the max pool size
500 /// and there are no idle threads.
501 ///
502 /// When attempting to increment the total worker count before creating a worker fails due to the
503 /// counter reaching the provided limit (core_size when attempting to create core thread, else
504 /// max_size) after being incremented by another thread, the pool tries to create
505 /// a non-core worker instead (if previously trying to create a core worker and no idle
506 /// worker exists) or sends the task to the channel instead. If incrementing the counter succeeded,
507 /// either because the current value of the counter matched the expected value or because the
508 /// last observed value was still below the limit, the worker starts with the provided task as
509 /// initial task and spawns its thread.
510 ///
511 /// # Panics
512 ///
513 /// This function might panic if `try_execute` returns an error when the crossbeam channel has been
514 /// closed unexpectedly.
515 /// This should never occur under normal circumstances using safe code, as shutting down the `ThreadPool`
516 /// consumes ownership and the crossbeam channel is never dropped unless dropping the `ThreadPool`.
517 pub fn execute<T: Task<()> + 'static>(&self, task: T) {
518 if self.try_execute(task).is_err() {
519 panic!("the channel of the thread pool has been closed");
520 }
521 }
522
523 /// Send a new task to the worker threads. This function is responsible for sending the message through the
524 /// channel and creating new workers if needed. If the current worker count is lower than the core pool size
525 /// this function will always create a new worker. If the current worker count is equal to or greater than
526 /// the core pool size this function only creates a new worker if the worker count is below the max pool size
527 /// and there are no idle threads.
528 ///
529 /// When attempting to increment the total worker count before creating a worker fails due to the
530 /// counter reaching the provided limit (core_size when attempting to create core thread, else
531 /// max_size) after being incremented by another thread, the pool tries to create
532 /// a non-core worker instead (if previously trying to create a core worker and no idle
533 /// worker exists) or sends the task to the channel instead. If incrementing the counter succeeded,
534 /// either because the current value of the counter matched the expected value or because the
535 /// last observed value was still below the limit, the worker starts with the provided task as
536 /// initial task and spawns its thread.
537 ///
538 /// # Errors
539 ///
540 /// This function might return `crossbeam_channel::SendError` if the sender was dropped unexpectedly.
541 pub fn try_execute<T: Task<()> + 'static>(
542 &self,
543 task: T,
544 ) -> Result<(), crossbeam_channel::SendError<Job>> {
545 if task.is_fn() {
546 self.try_execute_task(
547 task.into_fn()
548 .expect("Task::into_fn returned None despite is_fn returning true"),
549 )
550 } else {
551 self.try_execute_task(Box::new(move || {
552 task.run();
553 }))
554 }
555 }
556
557 /// Send a new task to the worker threads and return a [`JoinHandle`](struct.JoinHandle.html) that may be used to await
558 /// the result. This function is responsible for sending the message through the channel and creating new
559 /// workers if needed. If the current worker count is lower than the core pool size this function will always
560 /// create a new worker. If the current worker count is equal to or greater than the core pool size this
561 /// function only creates a new worker if the worker count is below the max pool size and there are no idle
562 /// threads.
563 ///
564 /// When attempting to increment the total worker count before creating a worker fails due to the
565 /// counter reaching the provided limit (core_size when attempting to create core thread, else
566 /// max_size) after being incremented by another thread, the pool tries to create
567 /// a non-core worker instead (if previously trying to create a core worker and no idle
568 /// worker exists) or sends the task to the channel instead. If incrementing the counter succeeded,
569 /// either because the current value of the counter matched the expected value or because the
570 /// last observed value was still below the limit, the worker starts with the provided task as
571 /// initial task and spawns its thread.
572 ///
573 /// # Panics
574 ///
575 /// This function might panic if `try_execute` returns an error when the crossbeam channel has been
576 /// closed unexpectedly.
577 /// This should never occur under normal circumstances using safe code, as shutting down the `ThreadPool`
578 /// consumes ownership and the crossbeam channel is never dropped unless dropping the `ThreadPool`.
579 pub fn evaluate<R: Send + 'static, T: Task<R> + 'static>(&self, task: T) -> JoinHandle<R> {
580 match self.try_evaluate(task) {
581 Ok(handle) => handle,
582 Err(e) => panic!("the channel of the thread pool has been closed: {:?}", e),
583 }
584 }
585
586 /// Send a new task to the worker threads and return a [`JoinHandle`](struct.JoinHandle.html) that may be used to await
587 /// the result. This function is responsible for sending the message through the channel and creating new
588 /// workers if needed. If the current worker count is lower than the core pool size this function will always
589 /// create a new worker. If the current worker count is equal to or greater than the core pool size this
590 /// function only creates a new worker if the worker count is below the max pool size and there are no idle
591 /// threads.
592 ///
593 /// When attempting to increment the total worker count before creating a worker fails due to the
594 /// counter reaching the provided limit (core_size when attempting to create core thread, else
595 /// max_size) after being incremented by another thread, the pool tries to create
596 /// a non-core worker instead (if previously trying to create a core worker and no idle
597 /// worker exists) or sends the task to the channel instead. If incrementing the counter succeeded,
598 /// either because the current value of the counter matched the expected value or because the
599 /// last observed value was still below the limit, the worker starts with the provided task as
600 /// initial task and spawns its thread.
601 ///
602 /// # Errors
603 ///
604 /// This function might return `crossbeam_channel::SendError` if the sender was dropped unexpectedly.
605 pub fn try_evaluate<R: Send + 'static, T: Task<R> + 'static>(
606 &self,
607 task: T,
608 ) -> Result<JoinHandle<R>, crossbeam_channel::SendError<Job>> {
609 let (sender, receiver) = oneshot::channel::<R>();
610 let join_handle = JoinHandle { receiver };
611 let job = || {
612 let result = task.run();
613 // if the receiver was dropped that means the caller was not interested in the result
614 let _ignored_result = sender.send(result);
615 };
616
617 let execute_attempt = self.try_execute_task(Box::new(job));
618 execute_attempt.map(|_| join_handle)
619 }
620
621 /// Send a task to the `ThreadPool` that completes the given `Future` and return a [`JoinHandle`](struct.JoinHandle.html)
622 /// that may be used to await the result. This function simply calls [`evaluate()`](struct.ThreadPool.html#method.evaluate)
623 /// with a closure that calls `block_on` with the provided future.
624 ///
625 /// # Panic
626 ///
627 /// This function panics if the task fails to be sent to the `ThreadPool` due to the channel being broken.
628 pub fn complete<R: Send + 'static>(
629 &self,
630 future: impl Future<Output = R> + 'static + Send,
631 ) -> JoinHandle<R> {
632 self.evaluate(|| block_on(future))
633 }
634
635 /// Send a task to the `ThreadPool` that completes the given `Future` and return a [`JoinHandle`](struct.JoinHandle.html)
636 /// that may be used to await the result. This function simply calls [`try_evaluate()`](struct.ThreadPool.html#method.try_evaluate)
637 /// with a closure that calls `block_on` with the provided future.
638 ///
639 /// # Errors
640 ///
641 /// This function returns `crossbeam_channel::SendError` if the task fails to be sent to the `ThreadPool` due to the channel being broken.
642 pub fn try_complete<R: Send + 'static>(
643 &self,
644 future: impl Future<Output = R> + 'static + Send,
645 ) -> Result<JoinHandle<R>, crossbeam_channel::SendError<Job>> {
646 self.try_evaluate(|| block_on(future))
647 }
648
649 /// Submit a `Future` to be polled by this `ThreadPool`. Unlike [`complete()`](struct.ThreadPool.html#method.complete) this does not
650 /// block a worker until the `Future` has been completed but polls the `Future` once at a time and creates a `Waker`
651 /// that re-submits the Future to this pool when awakened. Since `Arc<AsyncTask>` implements the [`Task`](trait.Task.html) trait this
652 /// function simply constructs the `AsyncTask` and calls [`execute()`](struct.ThreadPool.html#method.execute).
653 ///
654 /// # Panic
655 ///
656 /// This function panics if the task fails to be sent to the `ThreadPool` due to the channel being broken.
657
658 pub fn spawn(&self, future: impl Future<Output = ()> + 'static + Send) {
659 let future_task = Arc::new(AsyncTask {
660 future: Mutex::new(Some(Box::pin(future))),
661 pool: self.clone(),
662 });
663
664 self.execute(future_task)
665 }
666
667 /// Submit a `Future` to be polled by this `ThreadPool`. Unlike [`try_complete()`](struct.ThreadPool.html#method.try_complete) this does not
668 /// block a worker until the `Future` has been completed but polls the `Future` once at a time and creates a `Waker`
669 /// that re-submits the Future to this pool when awakened. Since `Arc<AsyncTask>` implements the [`Task`](trait.Task.html) trait this
670 /// function simply constructs the `AsyncTask` and calls [`try_execute()`](struct.ThreadPool.html#method.try_execute).
671 ///
672 /// # Errors
673 ///
674 /// This function returns `crossbeam_channel::SendError` if the task fails to be sent to the `ThreadPool` due to the channel being broken.
675
676 pub fn try_spawn(
677 &self,
678 future: impl Future<Output = ()> + 'static + Send,
679 ) -> Result<(), crossbeam_channel::SendError<Job>> {
680 let future_task = Arc::new(AsyncTask {
681 future: Mutex::new(Some(Box::pin(future))),
682 pool: self.clone(),
683 });
684
685 self.try_execute(future_task)
686 }
687
688 /// Create a top-level `Future` that awaits the provided `Future` and then sends the result to the
689 /// returned [`JoinHandle`](struct.JoinHandle.html). Unlike [`complete()`](struct.ThreadPool.html#method.complete) this does not
690 /// block a worker until the `Future` has been completed but polls the `Future` once at a time and creates a `Waker`
691 /// that re-submits the Future to this pool when awakened. Since `Arc<AsyncTask>` implements the [`Task`](trait.Task.html) trait this
692 /// function simply constructs the `AsyncTask` and calls [`execute()`](struct.ThreadPool.html#method.execute).
693 ///
694 /// This enables awaiting the final result outside of an async context like [`complete()`](struct.ThreadPool.html#method.complete) while still
695 /// polling the future lazily instead of eagerly blocking the worker until the future is done.
696 ///
697 /// # Panic
698 ///
699 /// This function panics if the task fails to be sent to the `ThreadPool` due to the channel being broken.
700
701 pub fn spawn_await<R: Send + 'static>(
702 &self,
703 future: impl Future<Output = R> + 'static + Send,
704 ) -> JoinHandle<R> {
705 match self.try_spawn_await(future) {
706 Ok(handle) => handle,
707 Err(e) => panic!("the channel of the thread pool has been closed: {:?}", e),
708 }
709 }
710
711 /// Create a top-level `Future` that awaits the provided `Future` and then sends the result to the
712 /// returned [`JoinHandle`](struct.JoinHandle.html). Unlike [`try_complete()`](struct.ThreadPool.html#method.try_complete) this does not
713 /// block a worker until the `Future` has been completed but polls the `Future` once at a time and creates a `Waker`
714 /// that re-submits the Future to this pool when awakened. Since `Arc<AsyncTask>` implements the [`Task`](trait.Task.html) trait this
715 /// function simply constructs the `AsyncTask` and calls [`try_execute()`](struct.ThreadPool.html#method.try_execute).
716 ///
717 /// This enables awaiting the final result outside of an async context like [`complete()`](struct.ThreadPool.html#method.complete) while still
718 /// polling the future lazily instead of eagerly blocking the worker until the future is done.
719 ///
720 /// # Errors
721 ///
722 /// This function returns `crossbeam_channel::SendError` if the task fails to be sent to the `ThreadPool` due to the channel being broken.
723
724 pub fn try_spawn_await<R: Send + 'static>(
725 &self,
726 future: impl Future<Output = R> + 'static + Send,
727 ) -> Result<JoinHandle<R>, crossbeam_channel::SendError<Job>> {
728 let (sender, receiver) = oneshot::channel::<R>();
729 let join_handle = JoinHandle { receiver };
730
731 self.try_spawn(async {
732 let result = future.await;
733 // if the receiver was dropped that means the caller was not interested in the result
734 let _ignored_result = sender.send(result);
735 })
736 .map(|_| join_handle)
737 }
738
739 #[inline]
740 fn try_execute_task(&self, task: Job) -> Result<(), crossbeam_channel::SendError<Job>> {
741 // create a new worker either if the current worker count is lower than the core pool size
742 // or if there are no idle threads and the current worker count is lower than the max pool size
743 let worker_count_data = &self.worker_data.worker_count_data;
744 let mut worker_count_val = worker_count_data.worker_count.load(Ordering::Relaxed);
745 let (mut curr_worker_count, idle_worker_count) = WorkerCountData::split(worker_count_val);
746 let mut curr_idle_count = idle_worker_count;
747
748 // always create a new worker if current pool size is below core size
749 if curr_worker_count < self.core_size {
750 let witnessed =
751 worker_count_data.try_increment_worker_total(worker_count_val, self.core_size);
752
753 // the witnessed value matched the expected value, meaning the initial exchange succeeded, or the final witnessed
754 // value is still below the coreSize, meaning the increment eventually succeeded
755 if witnessed == worker_count_val
756 || WorkerCountData::get_total_count(witnessed) < self.core_size
757 {
758 let worker = Worker::new(
759 self.channel_data.receiver.clone(),
760 Arc::clone(&self.worker_data),
761 None,
762 );
763
764 worker.start(Some(task));
765 return Ok(());
766 }
767
768 curr_worker_count = WorkerCountData::get_total_count(witnessed);
769 curr_idle_count = WorkerCountData::get_idle_count(witnessed);
770 worker_count_val = witnessed;
771 }
772
773 // create a new worker if the current worker count is below the maxSize and the pool has been observed to be busy
774 // (no idle workers) during the invocation of this function
775 if curr_worker_count < self.max_size && (idle_worker_count == 0 || curr_idle_count == 0) {
776 let witnessed =
777 worker_count_data.try_increment_worker_total(worker_count_val, self.max_size);
778
779 if witnessed == worker_count_val
780 || WorkerCountData::get_total_count(witnessed) < self.max_size
781 {
782 let worker = Worker::new(
783 self.channel_data.receiver.clone(),
784 Arc::clone(&self.worker_data),
785 Some(self.keep_alive),
786 );
787
788 worker.start(Some(task));
789 return Ok(());
790 }
791 }
792
793 self.send_task_to_channel(task)
794 }
795
796 /// Blocks the current thread until there aren't any non-idle threads anymore.
797 /// This includes work started after calling this function.
798 /// This function blocks until the next time this `ThreadPool` completes all of its work,
799 /// except if all threads are idle and the channel is empty at the time of calling this
800 /// function, in which case it will fast-return.
801 ///
802 /// This utilizes a `Condvar` that is notified by workers when they complete a job and notice
803 /// that the channel is currently empty and it was the last thread to finish the current
804 /// generation of work (i.e. when incrementing the idle worker counter brings the value
805 /// up to the total worker counter, meaning it's the last thread to become idle).
806 pub fn join(&self) {
807 self.inner_join(None);
808 }
809
810 /// Blocks the current thread until there aren't any non-idle threads anymore or until the
811 /// specified time_out Duration passes, whichever happens first.
812 /// This includes work started after calling this function.
813 /// This function blocks until the next time this `ThreadPool` completes all of its work,
814 /// (or until the time_out is reached) except if all threads are idle and the channel is
815 /// empty at the time of calling this function, in which case it will fast-return.
816 ///
817 /// This utilizes a `Condvar` that is notified by workers when they complete a job and notice
818 /// that the channel is currently empty and it was the last thread to finish the current
819 /// generation of work (i.e. when incrementing the idle worker counter brings the value
820 /// up to the total worker counter, meaning it's the last thread to become idle).
821 pub fn join_timeout(&self, time_out: Duration) {
822 self.inner_join(Some(time_out));
823 }
824
825 /// Destroy this `ThreadPool` by claiming ownership and dropping the value,
826 /// causing the `Sender` to drop thus disconnecting the channel.
827 /// Threads in this pool that are currently executing a task will finish what
828 /// they're doing until they check the channel, discovering that it has been
829 /// disconnected from the sender and thus terminate their work loop.
830 ///
831 /// If other clones of this `ThreadPool` exist the sender will remain intact
832 /// and tasks submitted to those clones will succeed, this includes pending
833 /// `AsyncTask` instances as they hold an owned clone of the `ThreadPool`
834 /// to re-submit awakened futures.
835 pub fn shutdown(self) {
836 drop(self);
837 }
838
839 /// Destroy this `ThreadPool` by claiming ownership and dropping the value,
840 /// causing the `Sender` to drop thus disconnecting the channel.
841 /// Threads in this pool that are currently executing a task will finish what
842 /// they're doing until they check the channel, discovering that it has been
843 /// disconnected from the sender and thus terminate their work loop.
844 ///
845 /// If other clones of this `ThreadPool` exist the sender will remain intact
846 /// and tasks submitted to those clones will succeed, this includes pending
847 /// `AsyncTask` instances as they hold an owned clone of the `ThreadPool`
848 /// to re-submit awakened futures.
849 ///
850 /// This function additionally joins all workers after dropping the pool to
851 /// wait for all work to finish.
852 /// Blocks the current thread until there aren't any non-idle threads anymore.
853 /// This function blocks until this `ThreadPool` completes all of its work,
854 /// except if all threads are idle and the channel is empty at the time of
855 /// calling this function, in which case the join will fast-return.
856 /// If other live clones of this `ThreadPool` exist this behaves the same as
857 /// calling [`join`](struct.ThreadPool.html#method.join) on a live `ThreadPool` as tasks submitted
858 /// to one of the clones will be joined as well.
859 ///
860 /// The join utilizes a `Condvar` that is notified by workers when they complete a job and notice
861 /// that the channel is currently empty and it was the last thread to finish the current
862 /// generation of work (i.e. when incrementing the idle worker counter brings the value
863 /// up to the total worker counter, meaning it's the last thread to become idle).
864 pub fn shutdown_join(self) {
865 self.inner_shutdown_join(None);
866 }
867
868 /// Destroy this `ThreadPool` by claiming ownership and dropping the value,
869 /// causing the `Sender` to drop thus disconnecting the channel.
870 /// Threads in this pool that are currently executing a task will finish what
871 /// they're doing until they check the channel, discovering that it has been
872 /// disconnected from the sender and thus terminate their work loop.
873 ///
874 /// If other clones of this `ThreadPool` exist the sender will remain intact
875 /// and tasks submitted to those clones will succeed, this includes pending
876 /// `AsyncTask` instances as they hold an owned clone of the `ThreadPool`
877 /// to re-submit awakened futures.
878 ///
879 /// This function additionally joins all workers after dropping the pool to
880 /// wait for all work to finish.
881 /// Blocks the current thread until there aren't any non-idle threads anymore or until the
882 /// specified time_out Duration passes, whichever happens first.
883 /// This function blocks until this `ThreadPool` completes all of its work,
884 /// (or until the time_out is reached) except if all threads are idle and the channel is
885 /// empty at the time of calling this function, in which case the join will fast-return.
886 /// If other live clones of this `ThreadPool` exist this behaves the same as
887 /// calling [`join`](struct.ThreadPool.html#method.join) on a live `ThreadPool` as tasks submitted
888 /// to one of the clones will be joined as well.
889 ///
890 /// The join utilizes a `Condvar` that is notified by workers when they complete a job and notice
891 /// that the channel is currently empty and it was the last thread to finish the current
892 /// generation of work (i.e. when incrementing the idle worker counter brings the value
893 /// up to the total worker counter, meaning it's the last thread to become idle).
894 pub fn shutdown_join_timeout(self, timeout: Duration) {
895 self.inner_shutdown_join(Some(timeout));
896 }
897
898 /// Return the name of this pool, used as prefix for each worker thread.
899 pub fn get_name(&self) -> &str {
900 &self.worker_data.pool_name
901 }
902
903 /// Starts all core workers by creating core idle workers until the total worker count reaches the core count.
904 ///
905 /// Returns immediately if the current worker count is already >= core size.
906 pub fn start_core_threads(&self) {
907 let worker_count_data = &self.worker_data.worker_count_data;
908
909 let core_size = self.core_size;
910 let mut curr_worker_count = worker_count_data.worker_count.load(Ordering::Relaxed);
911 if WorkerCountData::get_total_count(curr_worker_count) >= core_size {
912 return;
913 }
914
915 loop {
916 let witnessed = worker_count_data.try_increment_worker_count(
917 curr_worker_count,
918 INCREMENT_TOTAL | INCREMENT_IDLE,
919 core_size,
920 );
921
922 if WorkerCountData::get_total_count(witnessed) >= core_size {
923 return;
924 }
925
926 let worker = Worker::new(
927 self.channel_data.receiver.clone(),
928 Arc::clone(&self.worker_data),
929 None,
930 );
931
932 worker.start(None);
933 curr_worker_count = witnessed;
934 }
935 }
936
937 #[inline]
938 fn send_task_to_channel(&self, task: Job) -> Result<(), crossbeam_channel::SendError<Job>> {
939 self.channel_data.sender.send(task)?;
940
941 Ok(())
942 }
943
944 #[inline]
945 fn inner_join(&self, time_out: Option<Duration>) {
946 ThreadPool::_do_join(&self.worker_data, &self.channel_data.receiver, time_out);
947 }
948
949 #[inline]
950 fn inner_shutdown_join(self, timeout: Option<Duration>) {
951 let current_worker_data = self.worker_data.clone();
952 let receiver = self.channel_data.receiver.clone();
953 drop(self);
954 ThreadPool::_do_join(¤t_worker_data, &receiver, timeout);
955 }
956
957 #[inline]
958 fn _do_join(
959 current_worker_data: &Arc<WorkerData>,
960 receiver: &crossbeam_channel::Receiver<Job>,
961 time_out: Option<Duration>,
962 ) {
963 // no thread is currently doing any work, return
964 if ThreadPool::is_idle(current_worker_data, receiver) {
965 return;
966 }
967
968 let join_generation = current_worker_data.join_generation.load(Ordering::SeqCst);
969 let guard = current_worker_data
970 .join_notify_mutex
971 .lock()
972 .expect("could not get join notify mutex lock");
973
974 match time_out {
975 Some(time_out) => {
976 let _ret_guard = current_worker_data
977 .join_notify_condvar
978 .wait_timeout_while(guard, time_out, |_| {
979 join_generation
980 == current_worker_data.join_generation.load(Ordering::Relaxed)
981 && !ThreadPool::is_idle(current_worker_data, receiver)
982 })
983 .expect("could not wait for join condvar");
984 }
985 None => {
986 let _ret_guard = current_worker_data
987 .join_notify_condvar
988 .wait_while(guard, |_| {
989 join_generation
990 == current_worker_data.join_generation.load(Ordering::Relaxed)
991 && !ThreadPool::is_idle(current_worker_data, receiver)
992 })
993 .expect("could not wait for join condvar");
994 }
995 };
996
997 // increment generation if current thread is first thread to be awakened from wait in current generation
998 let _ = current_worker_data.join_generation.compare_exchange(
999 join_generation,
1000 join_generation.wrapping_add(1),
1001 Ordering::SeqCst,
1002 Ordering::SeqCst,
1003 );
1004 }
1005
1006 #[inline]
1007 fn is_idle(
1008 current_worker_data: &Arc<WorkerData>,
1009 receiver: &crossbeam_channel::Receiver<Job>,
1010 ) -> bool {
1011 let (current_worker_count, current_idle_count) =
1012 current_worker_data.worker_count_data.get_both();
1013 current_idle_count == current_worker_count && receiver.is_empty()
1014 }
1015}
1016
1017impl Default for ThreadPool {
1018 /// create default ThreadPool with the core pool size being equal to the number of cpus
1019 /// and the max_size being twice the core size with a 60 second timeout
1020 fn default() -> Self {
1021 let num_cpus = *全局_CPU数量;
1022 ThreadPool::new(
1023 num_cpus,
1024 std::cmp::max(num_cpus, num_cpus * 2),
1025 Duration::from_secs(60),
1026 )
1027 }
1028}
1029
1030/// A helper struct to aid creating a new `ThreadPool` using default values where no value was
1031/// explicitly specified.
1032#[derive(Default)]
1033pub struct Builder {
1034 name: Option<String>,
1035 core_size: Option<usize>,
1036 max_size: Option<usize>,
1037 keep_alive: Option<Duration>,
1038}
1039
1040impl Builder {
1041 /// Create a new `Builder`.
1042 pub fn new() -> Builder {
1043 Builder::default()
1044 }
1045
1046 /// Specify the name of the `ThreadPool` that will be used as prefix for the name of each worker thread.
1047 /// By default the name is "rusty_pool_x" with x being a static pool counter.
1048 pub fn name(mut self, name: String) -> Builder {
1049 self.name = Some(name);
1050 self
1051 }
1052
1053 /// Specify the core pool size for the `ThreadPool`. The core pool size is the number of threads that stay alive
1054 /// for the entire lifetime of the `ThreadPool` or, to be more precise, its channel. These threads are spawned if
1055 /// a task is submitted to the `ThreadPool` and the current worker count is below the core pool size.
1056 pub fn core_size(mut self, size: usize) -> Builder {
1057 self.core_size = Some(size);
1058 self
1059 }
1060
1061 /// Specify the maximum pool size this `ThreadPool` may scale up to. This numbers represents the maximum number
1062 /// of threads that may be alive at the same time within this pool. Additional threads above the core pool size
1063 /// only remain idle for the duration specified by the `keep_alive` parameter before terminating. If the core pool
1064 /// is full, the current pool size is below the max size and there are no idle threads then additional threads
1065 /// will be spawned.
1066 pub fn max_size(mut self, size: usize) -> Builder {
1067 self.max_size = Some(size);
1068 self
1069 }
1070
1071 /// Specify the duration for which additional threads outside the core pool remain alive while not receiving any
1072 /// work before giving up and terminating.
1073 pub fn keep_alive(mut self, keep_alive: Duration) -> Builder {
1074 self.keep_alive = Some(keep_alive);
1075 self
1076 }
1077
1078 /// Build the `ThreadPool` using the parameters previously supplied to this `Builder` using the number of CPUs as
1079 /// default core size if none provided, twice the core size as max size if none provided, 60 seconds keep_alive
1080 /// if none provided and the default naming (rusty_pool_{pool_number}) if none provided.
1081 /// This function calls [`ThreadPool::new`](struct.ThreadPool.html#method.new) or
1082 /// [`ThreadPool::new_named`](struct.ThreadPool.html#method.new_named) depending on whether a name was provided.
1083 ///
1084 /// # Panics
1085 ///
1086 /// Building might panic if the `max_size` is 0 or lower than `core_size` or exceeds half
1087 /// the size of usize. This restriction exists because two counters (total workers and
1088 /// idle counters) are stored within one AtomicUsize.
1089 pub fn build(self) -> ThreadPool {
1090 use std::cmp::{max, min};
1091
1092 let core_size = self.core_size.unwrap_or_else(|| {
1093 let num_cpus = num_cpus::get();
1094 if let Some(max_size) = self.max_size {
1095 min(MAX_SIZE, min(num_cpus, max_size))
1096 } else {
1097 min(MAX_SIZE, num_cpus)
1098 }
1099 });
1100 // handle potential overflow: try using twice the core_size or return core_size
1101 let max_size = self
1102 .max_size
1103 .unwrap_or_else(|| min(MAX_SIZE, max(core_size, core_size * 2)));
1104 let keep_alive = self.keep_alive.unwrap_or_else(|| Duration::from_secs(60));
1105
1106 if let Some(name) = self.name {
1107 ThreadPool::new_named(name, core_size, max_size, keep_alive)
1108 } else {
1109 ThreadPool::new(core_size, max_size, keep_alive)
1110 }
1111 }
1112}
1113
1114#[derive(Clone)]
1115struct Worker {
1116 receiver: crossbeam_channel::Receiver<Job>,
1117 worker_data: Arc<WorkerData>,
1118 keep_alive: Option<Duration>,
1119}
1120
1121impl Worker {
1122 fn new(
1123 receiver: crossbeam_channel::Receiver<Job>,
1124 worker_data: Arc<WorkerData>,
1125 keep_alive: Option<Duration>,
1126 ) -> Self {
1127 Worker {
1128 receiver,
1129 worker_data,
1130 keep_alive,
1131 }
1132 }
1133
1134 fn start(self, task: Option<Job>) {
1135 let worker_name = format!(
1136 "{}_thread_{}",
1137 self.worker_data.pool_name,
1138 self.worker_data
1139 .worker_number
1140 .fetch_add(1, Ordering::Relaxed)
1141 );
1142
1143 thread::Builder::new()
1144 .name(worker_name)
1145 .spawn(move || {
1146 let mut sentinel = Sentinel::new(&self);
1147
1148 if let Some(task) = task {
1149 self.exec_task_and_notify(&mut sentinel, task);
1150 }
1151
1152 loop {
1153 // the two functions return different error types, but since the error type doesn't matter it is mapped to unit to make them compatible
1154 let received_task: Result<Job, _> = match self.keep_alive {
1155 Some(keep_alive) => self.receiver.recv_timeout(keep_alive).map_err(|_| ()),
1156 None => self.receiver.recv().map_err(|_| ()),
1157 };
1158
1159 match received_task {
1160 Ok(task) => {
1161 // mark current as no longer idle and execute task
1162 self.worker_data.worker_count_data.decrement_worker_idle();
1163 self.exec_task_and_notify(&mut sentinel, task);
1164 }
1165 Err(_) => {
1166 // either channel was broken because the sender disconnected or, if can_timeout is true, the Worker has not received any work during
1167 // its keep_alive period and will now terminate, break working loop
1168 break;
1169 }
1170 }
1171 }
1172
1173 // can decrement both at once as the thread only gets here from an idle state
1174 // (if waiting for work and receiving an error)
1175 self.worker_data.worker_count_data.decrement_both();
1176 })
1177 .expect("could not spawn thread");
1178 }
1179
1180 #[inline]
1181 fn exec_task_and_notify(&self, sentinel: &mut Sentinel, task: Job) {
1182 sentinel.is_working = true;
1183 task();
1184 sentinel.is_working = false;
1185 // can already mark as idle as this thread will continue the work loop
1186 self.mark_idle_and_notify_joiners_if_no_work();
1187 }
1188
1189 #[inline]
1190 fn mark_idle_and_notify_joiners_if_no_work(&self) {
1191 let (old_total_count, old_idle_count) = self
1192 .worker_data
1193 .worker_count_data
1194 .increment_worker_idle_ret_both();
1195 // if the last task was the last one in the current generation,
1196 // i.e. if incrementing the idle count leads to the idle count
1197 // being equal to the total worker count, notify joiners
1198 if old_total_count == old_idle_count + 1 && self.receiver.is_empty() {
1199 let _lock = self
1200 .worker_data
1201 .join_notify_mutex
1202 .lock()
1203 .expect("could not get join notify mutex lock");
1204 self.worker_data.join_notify_condvar.notify_all();
1205 }
1206 }
1207}
1208
1209/// Type that exists to manage worker exit on panic.
1210///
1211/// This type is constructed once per `Worker` and implements `Drop` to handle proper worker exit
1212/// in case the worker panics when executing the current task or anywhere else in its work loop.
1213/// If the `Sentinel` is dropped at the end of the worker's work loop and the current thread is
1214/// panicking, handle worker exit the same way as if the task completed normally (if the worker
1215/// panicked while executing a submitted task) then clone the worker and start it with an initial
1216/// task of `None`.
1217struct Sentinel<'s> {
1218 is_working: bool,
1219 worker_ref: &'s Worker,
1220}
1221
1222impl Sentinel<'_> {
1223 fn new(worker_ref: &Worker) -> Sentinel<'_> {
1224 Sentinel {
1225 is_working: false,
1226 worker_ref,
1227 }
1228 }
1229}
1230
1231impl Drop for Sentinel<'_> {
1232 fn drop(&mut self) {
1233 if thread::panicking() {
1234 if self.is_working {
1235 // worker thread panicked in the process of executing a submitted task,
1236 // run the same logic as if the task completed normally and mark it as
1237 // idle, since a clone of this worker will start the work loop as idle
1238 // thread
1239 self.worker_ref.mark_idle_and_notify_joiners_if_no_work();
1240 }
1241
1242 let worker = self.worker_ref.clone();
1243 worker.start(None);
1244 }
1245 }
1246}
1247
1248const WORKER_IDLE_MASK: usize = MAX_SIZE;
1249const INCREMENT_TOTAL: usize = 1 << (BITS / 2);
1250const INCREMENT_IDLE: usize = 1;
1251
1252/// Struct that stores and handles an `AtomicUsize` that stores the total worker count
1253/// in the higher half of bits and the idle worker count in the lower half of bits.
1254/// This allows to to increment / decrement both counters in a single atomic operation.
1255#[derive(Default)]
1256struct WorkerCountData {
1257 worker_count: AtomicUsize,
1258}
1259
1260impl WorkerCountData {
1261 fn get_total_worker_count(&self) -> usize {
1262 let curr_val = self.worker_count.load(Ordering::Relaxed);
1263 WorkerCountData::get_total_count(curr_val)
1264 }
1265
1266 fn get_idle_worker_count(&self) -> usize {
1267 let curr_val = self.worker_count.load(Ordering::Relaxed);
1268 WorkerCountData::get_idle_count(curr_val)
1269 }
1270
1271 fn get_both(&self) -> (usize, usize) {
1272 let curr_val = self.worker_count.load(Ordering::Relaxed);
1273 WorkerCountData::split(curr_val)
1274 }
1275
1276 // keep for testing and completion's sake
1277 #[allow(dead_code)]
1278 fn increment_both(&self) -> (usize, usize) {
1279 let old_val = self
1280 .worker_count
1281 .fetch_add(INCREMENT_TOTAL | INCREMENT_IDLE, Ordering::Relaxed);
1282 WorkerCountData::split(old_val)
1283 }
1284
1285 fn decrement_both(&self) -> (usize, usize) {
1286 let old_val = self
1287 .worker_count
1288 .fetch_sub(INCREMENT_TOTAL | INCREMENT_IDLE, Ordering::Relaxed);
1289 WorkerCountData::split(old_val)
1290 }
1291
1292 fn try_increment_worker_total(&self, expected: usize, max_total: usize) -> usize {
1293 self.try_increment_worker_count(expected, INCREMENT_TOTAL, max_total)
1294 }
1295
1296 fn try_increment_worker_count(
1297 &self,
1298 mut expected: usize,
1299 increment: usize,
1300 max_total: usize,
1301 ) -> usize {
1302 loop {
1303 match self.worker_count.compare_exchange_weak(
1304 expected,
1305 expected + increment,
1306 Ordering::Relaxed,
1307 Ordering::Relaxed,
1308 ) {
1309 Ok(witnessed) => return witnessed,
1310 Err(witnessed) if WorkerCountData::get_total_count(witnessed) >= max_total => {
1311 return witnessed
1312 }
1313 Err(witnessed) => expected = witnessed,
1314 }
1315 }
1316 }
1317
1318 // keep for testing and completion's sake
1319 #[allow(dead_code)]
1320 fn increment_worker_total(&self) -> usize {
1321 let old_val = self
1322 .worker_count
1323 .fetch_add(INCREMENT_TOTAL, Ordering::Relaxed);
1324 WorkerCountData::get_total_count(old_val)
1325 }
1326
1327 // keep for testing and completion's sake
1328 #[allow(dead_code)]
1329 fn increment_worker_total_ret_both(&self) -> (usize, usize) {
1330 let old_val = self
1331 .worker_count
1332 .fetch_add(INCREMENT_TOTAL, Ordering::Relaxed);
1333 WorkerCountData::split(old_val)
1334 }
1335
1336 // keep for testing and completion's sake
1337 #[allow(dead_code)]
1338 fn decrement_worker_total(&self) -> usize {
1339 let old_val = self
1340 .worker_count
1341 .fetch_sub(INCREMENT_TOTAL, Ordering::Relaxed);
1342 WorkerCountData::get_total_count(old_val)
1343 }
1344
1345 // keep for testing and completion's sake
1346 #[allow(dead_code)]
1347 fn decrement_worker_total_ret_both(&self) -> (usize, usize) {
1348 let old_val = self
1349 .worker_count
1350 .fetch_sub(INCREMENT_TOTAL, Ordering::Relaxed);
1351 WorkerCountData::split(old_val)
1352 }
1353
1354 // keep for testing and completion's sake
1355 #[allow(dead_code)]
1356 fn increment_worker_idle(&self) -> usize {
1357 let old_val = self
1358 .worker_count
1359 .fetch_add(INCREMENT_IDLE, Ordering::Relaxed);
1360 WorkerCountData::get_idle_count(old_val)
1361 }
1362
1363 fn increment_worker_idle_ret_both(&self) -> (usize, usize) {
1364 let old_val = self
1365 .worker_count
1366 .fetch_add(INCREMENT_IDLE, Ordering::Relaxed);
1367 WorkerCountData::split(old_val)
1368 }
1369
1370 fn decrement_worker_idle(&self) -> usize {
1371 let old_val = self
1372 .worker_count
1373 .fetch_sub(INCREMENT_IDLE, Ordering::Relaxed);
1374 WorkerCountData::get_idle_count(old_val)
1375 }
1376
1377 // keep for testing and completion's sake
1378 #[allow(dead_code)]
1379 fn decrement_worker_idle_ret_both(&self) -> (usize, usize) {
1380 let old_val = self
1381 .worker_count
1382 .fetch_sub(INCREMENT_IDLE, Ordering::Relaxed);
1383 WorkerCountData::split(old_val)
1384 }
1385
1386 #[inline]
1387 fn split(val: usize) -> (usize, usize) {
1388 let total_count = val >> (BITS / 2);
1389 let idle_count = val & WORKER_IDLE_MASK;
1390 (total_count, idle_count)
1391 }
1392
1393 #[inline]
1394 fn get_total_count(val: usize) -> usize {
1395 val >> (BITS / 2)
1396 }
1397
1398 #[inline]
1399 fn get_idle_count(val: usize) -> usize {
1400 val & WORKER_IDLE_MASK
1401 }
1402}
1403
1404/// struct containing data shared between workers
1405struct WorkerData {
1406 pool_name: String,
1407 worker_count_data: WorkerCountData,
1408 worker_number: AtomicUsize,
1409 join_notify_condvar: Condvar,
1410 join_notify_mutex: Mutex<()>,
1411 join_generation: AtomicUsize,
1412}
1413
1414struct ChannelData {
1415 sender: crossbeam_channel::Sender<Job>,
1416 receiver: crossbeam_channel::Receiver<Job>,
1417}
1418