bendns' repos / rust-analyzer
Unnamed repository; edit this file 'description' to name the repository.
Diffstat (limited to 'crates/stdx/src/thread/pool.rs')
-rw-r--r--crates/stdx/src/thread/pool.rs51
1 files changed, 47 insertions, 4 deletions
diff --git a/crates/stdx/src/thread/pool.rs b/crates/stdx/src/thread/pool.rs
index a8de4db624..8d76c5fd1f 100644
--- a/crates/stdx/src/thread/pool.rs
+++ b/crates/stdx/src/thread/pool.rs
@@ -8,6 +8,7 @@
//! the threading utilities in [`crate::thread`].
use std::{
+ marker::PhantomData,
panic::{self, UnwindSafe},
sync::{
Arc,
@@ -16,8 +17,9 @@ use std::{
};
use crossbeam_channel::{Receiver, Sender};
+use crossbeam_utils::sync::WaitGroup;
-use super::{Builder, JoinHandle, ThreadIntent};
+use crate::thread::{Builder, JoinHandle, ThreadIntent};
pub struct Pool {
// `_handles` is never read: the field is present
@@ -79,9 +81,6 @@ impl Pool {
Self { _handles: handles.into_boxed_slice(), extant_tasks, job_sender }
}
- /// # Panics
- ///
- /// Panics if job panics
pub fn spawn<F>(&self, intent: ThreadIntent, f: F)
where
F: FnOnce() + Send + UnwindSafe + 'static,
@@ -97,6 +96,17 @@ impl Pool {
self.job_sender.send(job).unwrap();
}
+ pub fn scoped<'pool, 'scope, F, R>(&'pool self, f: F) -> R
+ where
+ F: FnOnce(&Scope<'pool, 'scope>) -> R,
+ {
+ let wg = WaitGroup::new();
+ let scope = Scope { pool: self, wg, _marker: PhantomData };
+ let r = f(&scope);
+ scope.wg.wait();
+ r
+ }
+
#[must_use]
pub fn len(&self) -> usize {
self.extant_tasks.load(Ordering::SeqCst)
@@ -107,3 +117,36 @@ impl Pool {
self.len() == 0
}
}
+
+pub struct Scope<'pool, 'scope> {
+ pool: &'pool Pool,
+ wg: WaitGroup,
+ _marker: PhantomData<fn(&'scope ()) -> &'scope ()>,
+}
+
+impl<'scope> Scope<'_, 'scope> {
+ pub fn spawn<F>(&self, intent: ThreadIntent, f: F)
+ where
+ F: 'scope + FnOnce() + Send + UnwindSafe,
+ {
+ let wg = self.wg.clone();
+ let f = Box::new(move || {
+ if cfg!(debug_assertions) {
+ intent.assert_is_used_on_current_thread();
+ }
+ f();
+ drop(wg);
+ });
+
+ let job = Job {
+ requested_intent: intent,
+ f: unsafe {
+ std::mem::transmute::<
+ Box<dyn 'scope + FnOnce() + Send + UnwindSafe>,
+ Box<dyn 'static + FnOnce() + Send + UnwindSafe>,
+ >(f)
+ },
+ };
+ self.pool.job_sender.send(job).unwrap();
+ }
+}