1use crate::job::{ArcJob, StackJob};
2use crate::latch::LatchRef;
3use crate::registry::{Registry, WorkerThread};
4use crate::scope::ScopeLatch;
5use std::fmt;
6use std::marker::PhantomData;
7use std::sync::Arc;
8
9mod test;
10
11/// Executes `op` within every thread in the current threadpool. If this is
12/// called from a non-Rayon thread, it will execute in the global threadpool.
13/// Any attempts to use `join`, `scope`, or parallel iterators will then operate
14/// within that threadpool. When the call has completed on each thread, returns
15/// a vector containing all of their return values.
16///
17/// For more information, see the [`ThreadPool::broadcast()`][m] method.
18///
19/// [m]: struct.ThreadPool.html#method.broadcast
20pub fn broadcast<OP, R>(op: OP) -> Vec<R>
21where
22 OP: Fn(BroadcastContext<'_>) -> R + Sync,
23 R: Send,
24{
25 // We assert that current registry has not terminated.
26 unsafe { broadcast_in(op, &Registry::current()) }
27}
28
29/// Spawns an asynchronous task on every thread in this thread-pool. This task
30/// will run in the implicit, global scope, which means that it may outlast the
31/// current stack frame -- therefore, it cannot capture any references onto the
32/// stack (you will likely need a `move` closure).
33///
34/// For more information, see the [`ThreadPool::spawn_broadcast()`][m] method.
35///
36/// [m]: struct.ThreadPool.html#method.spawn_broadcast
37pub fn spawn_broadcast<OP>(op: OP)
38where
39 OP: Fn(BroadcastContext<'_>) + Send + Sync + 'static,
40{
41 // We assert that current registry has not terminated.
42 unsafe { spawn_broadcast_in(op, &Registry::current()) }
43}
44
45/// Provides context to a closure called by `broadcast`.
46pub struct BroadcastContext<'a> {
47 worker: &'a WorkerThread,
48
49 /// Make sure to prevent auto-traits like `Send` and `Sync`.
50 _marker: PhantomData<&'a mut dyn Fn()>,
51}
52
53impl<'a> BroadcastContext<'a> {
54 pub(super) fn with<R>(f: impl FnOnce(BroadcastContext<'_>) -> R) -> R {
55 let worker_thread = WorkerThread::current();
56 assert!(!worker_thread.is_null());
57 f(BroadcastContext {
58 worker: unsafe { &*worker_thread },
59 _marker: PhantomData,
60 })
61 }
62
63 /// Our index amongst the broadcast threads (ranges from `0..self.num_threads()`).
64 #[inline]
65 pub fn index(&self) -> usize {
66 self.worker.index()
67 }
68
69 /// The number of threads receiving the broadcast in the thread pool.
70 ///
71 /// # Future compatibility note
72 ///
73 /// Future versions of Rayon might vary the number of threads over time, but
74 /// this method will always return the number of threads which are actually
75 /// receiving your particular `broadcast` call.
76 #[inline]
77 pub fn num_threads(&self) -> usize {
78 self.worker.registry().num_threads()
79 }
80}
81
82impl<'a> fmt::Debug for BroadcastContext<'a> {
83 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
84 fmt&mut DebugStruct<'_, '_>.debug_struct("BroadcastContext")
85 .field("index", &self.index())
86 .field("num_threads", &self.num_threads())
87 .field(name:"pool_id", &self.worker.registry().id())
88 .finish()
89 }
90}
91
92/// Execute `op` on every thread in the pool. It will be executed on each
93/// thread when they have nothing else to do locally, before they try to
94/// steal work from other threads. This function will not return until all
95/// threads have completed the `op`.
96///
97/// Unsafe because `registry` must not yet have terminated.
98pub(super) unsafe fn broadcast_in<OP, R>(op: OP, registry: &Arc<Registry>) -> Vec<R>
99where
100 OP: Fn(BroadcastContext<'_>) -> R + Sync,
101 R: Send,
102{
103 let f: impl Fn(bool) -> R = move |injected: bool| {
104 debug_assert!(injected);
105 BroadcastContext::with(&op)
106 };
107
108 let n_threads: usize = registry.num_threads();
109 let current_thread: Option<&WorkerThread> = WorkerThread::current().as_ref();
110 let tlv: Tlv = crate::tlv::get();
111 let latch: ScopeLatch = ScopeLatch::with_count(count:n_threads, owner:current_thread);
112 let jobs: Vec<_> = (0..n_threads)
113 .map(|_| StackJob::new(tlv, &f, latch:LatchRef::new(&latch)))
114 .collect();
115 let job_refs: impl Iterator = jobs.iter().map(|job: &StackJob, …, …>| job.as_job_ref());
116
117 registry.inject_broadcast(injected_jobs:job_refs);
118
119 // Wait for all jobs to complete, then collect the results, maybe propagating a panic.
120 latch.wait(owner:current_thread);
121 jobs.into_iter().map(|job: StackJob, …, …>| job.into_result()).collect()
122}
123
124/// Execute `op` on every thread in the pool. It will be executed on each
125/// thread when they have nothing else to do locally, before they try to
126/// steal work from other threads. This function returns immediately after
127/// injecting the jobs.
128///
129/// Unsafe because `registry` must not yet have terminated.
130pub(super) unsafe fn spawn_broadcast_in<OP>(op: OP, registry: &Arc<Registry>)
131where
132 OP: Fn(BroadcastContext<'_>) + Send + Sync + 'static,
133{
134 let job: Arc> = ArcJob::new({
135 let registry: Arc = Arc::clone(self:registry);
136 move || {
137 registry.catch_unwind(|| BroadcastContext::with(&op));
138 registry.terminate(); // (*) permit registry to terminate now
139 }
140 });
141
142 let n_threads: usize = registry.num_threads();
143 let job_refs: impl Iterator = (0..n_threads).map(|_| {
144 // Ensure that registry cannot terminate until this job has executed
145 // on each thread. This ref is decremented at the (*) above.
146 registry.increment_terminate_count();
147
148 ArcJob::as_static_job_ref(&job)
149 });
150
151 registry.inject_broadcast(injected_jobs:job_refs);
152}
153