rustc_thread_pool/latch.rs
1use std::marker::PhantomData;
2use std::ops::Deref;
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Arc, Condvar, Mutex};
5
6use crate::job::JobRef;
7use crate::registry::{Registry, WorkerThread};
8
9/// We define various kinds of latches, which are all a primitive signaling
10/// mechanism. A latch starts as false. Eventually someone calls `set()` and
11/// it becomes true. You can test if it has been set by calling `probe()`.
12///
13/// Some kinds of latches, but not all, support a `wait()` operation
14/// that will wait until the latch is set, blocking efficiently. That
15/// is not part of the trait since it is not possibly to do with all
16/// latches.
17///
18/// The intention is that `set()` is called once, but `probe()` may be
19/// called any number of times. Once `probe()` returns true, the memory
20/// effects that occurred before `set()` become visible.
21///
22/// It'd probably be better to refactor the API into two paired types,
23/// but that's a bit of work, and this is not a public API.
24///
25/// ## Memory ordering
26///
27/// Latches need to guarantee two things:
28///
29/// - Once `probe()` returns true, all memory effects from the `set()`
30/// are visible (in other words, the set should synchronize-with
31/// the probe).
32/// - Once `set()` occurs, the next `probe()` *will* observe it. This
33/// typically requires a seq-cst ordering. See [the "tickle-then-get-sleepy" scenario in the sleep
34/// README](/src/sleep/README.md#tickle-then-get-sleepy) for details.
35pub(super) trait Latch {
36 /// Set the latch, signalling others.
37 ///
38 /// # WARNING
39 ///
40 /// Setting a latch triggers other threads to wake up and (in some
41 /// cases) complete. This may, in turn, cause memory to be
42 /// deallocated and so forth. One must be very careful about this,
43 /// and it's typically better to read all the fields you will need
44 /// to access *before* a latch is set!
45 ///
46 /// This function operates on `*const Self` instead of `&self` to allow it
47 /// to become dangling during this call. The caller must ensure that the
48 /// pointer is valid upon entry, and not invalidated during the call by any
49 /// actions other than `set` itself.
50 unsafe fn set(this: *const Self);
51}
52
53pub(super) trait AsCoreLatch {
54 fn as_core_latch(&self) -> &CoreLatch;
55}
56
57/// Latch is not set, owning thread is awake
58const UNSET: usize = 0;
59
60/// Latch is not set, owning thread is going to sleep on this latch
61/// (but has not yet fallen asleep).
62const SLEEPY: usize = 1;
63
64/// Latch is not set, owning thread is asleep on this latch and
65/// must be awoken.
66const SLEEPING: usize = 2;
67
68/// Latch is set.
69const SET: usize = 3;
70
71/// Spin latches are the simplest, most efficient kind, but they do
72/// not support a `wait()` operation. They just have a boolean flag
73/// that becomes true when `set()` is called.
74#[derive(Debug)]
75pub(super) struct CoreLatch {
76 state: AtomicUsize,
77}
78
79impl CoreLatch {
80 #[inline]
81 fn new() -> Self {
82 Self { state: AtomicUsize::new(0) }
83 }
84
85 /// Invoked by owning thread as it prepares to sleep. Returns true
86 /// if the owning thread may proceed to fall asleep, false if the
87 /// latch was set in the meantime.
88 #[inline]
89 pub(super) fn get_sleepy(&self) -> bool {
90 self.state.compare_exchange(UNSET, SLEEPY, Ordering::SeqCst, Ordering::Relaxed).is_ok()
91 }
92
93 /// Invoked by owning thread as it falls asleep sleep. Returns
94 /// true if the owning thread should block, or false if the latch
95 /// was set in the meantime.
96 #[inline]
97 pub(super) fn fall_asleep(&self) -> bool {
98 self.state.compare_exchange(SLEEPY, SLEEPING, Ordering::SeqCst, Ordering::Relaxed).is_ok()
99 }
100
101 /// Invoked by owning thread as it falls asleep sleep. Returns
102 /// true if the owning thread should block, or false if the latch
103 /// was set in the meantime.
104 #[inline]
105 pub(super) fn wake_up(&self) {
106 if !self.probe() {
107 let _ =
108 self.state.compare_exchange(SLEEPING, UNSET, Ordering::SeqCst, Ordering::Relaxed);
109 }
110 }
111
112 /// Set the latch. If this returns true, the owning thread was sleeping
113 /// and must be awoken.
114 ///
115 /// This is private because, typically, setting a latch involves
116 /// doing some wakeups; those are encapsulated in the surrounding
117 /// latch code.
118 #[inline]
119 unsafe fn set(this: *const Self) -> bool {
120 let old_state = unsafe { (*this).state.swap(SET, Ordering::AcqRel) };
121 old_state == SLEEPING
122 }
123
124 /// Test if this latch has been set.
125 #[inline]
126 pub(super) fn probe(&self) -> bool {
127 self.state.load(Ordering::Acquire) == SET
128 }
129}
130
131impl AsCoreLatch for CoreLatch {
132 #[inline]
133 fn as_core_latch(&self) -> &CoreLatch {
134 self
135 }
136}
137
138/// Spin latches are the simplest, most efficient kind, but they do
139/// not support a `wait()` operation. They just have a boolean flag
140/// that becomes true when `set()` is called.
141pub(super) struct SpinLatch<'r> {
142 core_latch: CoreLatch,
143 registry: &'r Arc<Registry>,
144 target_worker_index: usize,
145 cross: bool,
146}
147
148impl<'r> SpinLatch<'r> {
149 /// Creates a new spin latch that is owned by `thread`. This means
150 /// that `thread` is the only thread that should be blocking on
151 /// this latch -- it also means that when the latch is set, we
152 /// will wake `thread` if it is sleeping.
153 #[inline]
154 pub(super) fn new(thread: &'r WorkerThread) -> SpinLatch<'r> {
155 SpinLatch {
156 core_latch: CoreLatch::new(),
157 registry: thread.registry(),
158 target_worker_index: thread.index(),
159 cross: false,
160 }
161 }
162
163 /// Creates a new spin latch for cross-threadpool blocking. Notably, we
164 /// need to make sure the registry is kept alive after setting, so we can
165 /// safely call the notification.
166 #[inline]
167 pub(super) fn cross(thread: &'r WorkerThread) -> SpinLatch<'r> {
168 SpinLatch { cross: true, ..SpinLatch::new(thread) }
169 }
170}
171
172impl<'r> AsCoreLatch for SpinLatch<'r> {
173 #[inline]
174 fn as_core_latch(&self) -> &CoreLatch {
175 &self.core_latch
176 }
177}
178
179impl<'r> Latch for SpinLatch<'r> {
180 #[inline]
181 unsafe fn set(this: *const Self) {
182 let cross_registry;
183
184 let registry: &Registry = if unsafe { (*this).cross } {
185 // Ensure the registry stays alive while we notify it.
186 // Otherwise, it would be possible that we set the spin
187 // latch and the other thread sees it and exits, causing
188 // the registry to be deallocated, all before we get a
189 // chance to invoke `registry.notify_worker_latch_is_set`.
190 cross_registry = Arc::clone(unsafe { (*this).registry });
191 &cross_registry
192 } else {
193 // If this is not a "cross-registry" spin-latch, then the
194 // thread which is performing `set` is itself ensuring
195 // that the registry stays alive. However, that doesn't
196 // include this *particular* `Arc` handle if the waiting
197 // thread then exits, so we must completely dereference it.
198 unsafe { (*this).registry }
199 };
200 let target_worker_index = unsafe { (*this).target_worker_index };
201
202 // NOTE: Once we `set`, the target may proceed and invalidate `this`!
203 if unsafe { CoreLatch::set(&(*this).core_latch) } {
204 // Subtle: at this point, we can no longer read from
205 // `self`, because the thread owning this spin latch may
206 // have awoken and deallocated the latch. Therefore, we
207 // only use fields whose values we already read.
208 registry.notify_worker_latch_is_set(target_worker_index);
209 }
210 }
211}
212
213/// A Latch starts as false and eventually becomes true. You can block
214/// until it becomes true.
215#[derive(Debug)]
216pub(super) struct LockLatch {
217 m: Mutex<bool>,
218 v: Condvar,
219}
220
221impl LockLatch {
222 #[inline]
223 pub(super) fn new() -> LockLatch {
224 LockLatch { m: Mutex::new(false), v: Condvar::new() }
225 }
226
227 /// Block until latch is set, then resets this lock latch so it can be reused again.
228 pub(super) fn wait_and_reset(&self) {
229 let mut guard = self.m.lock().unwrap();
230 while !*guard {
231 guard = self.v.wait(guard).unwrap();
232 }
233 *guard = false;
234 }
235
236 /// Block until latch is set.
237 pub(super) fn wait(&self) {
238 let mut guard = self.m.lock().unwrap();
239 while !*guard {
240 guard = self.v.wait(guard).unwrap();
241 }
242 }
243}
244
245impl Latch for LockLatch {
246 #[inline]
247 unsafe fn set(this: *const Self) {
248 let mut guard = unsafe { (*this).m.lock().unwrap() };
249 *guard = true;
250 unsafe { (*this).v.notify_all() };
251 }
252}
253
254/// Once latches are used to implement one-time blocking, primarily
255/// for the termination flag of the threads in the pool.
256///
257/// Note: like a `SpinLatch`, once-latches are always associated with
258/// some registry that is probing them, which must be tickled when
259/// they are set. *Unlike* a `SpinLatch`, they don't themselves hold a
260/// reference to that registry. This is because in some cases the
261/// registry owns the once-latch, and that would create a cycle. So a
262/// `OnceLatch` must be given a reference to its owning registry when
263/// it is set. For this reason, it does not implement the `Latch`
264/// trait (but it doesn't have to, as it is not used in those generic
265/// contexts).
266#[derive(Debug)]
267pub(super) struct OnceLatch {
268 core_latch: CoreLatch,
269}
270
271impl OnceLatch {
272 #[inline]
273 pub(super) fn new() -> OnceLatch {
274 Self { core_latch: CoreLatch::new() }
275 }
276
277 /// Set the latch, then tickle the specific worker thread,
278 /// which should be the one that owns this latch.
279 #[inline]
280 pub(super) unsafe fn set_and_tickle_one(
281 this: *const Self,
282 registry: &Registry,
283 target_worker_index: usize,
284 ) {
285 if unsafe { CoreLatch::set(&(*this).core_latch) } {
286 registry.notify_worker_latch_is_set(target_worker_index);
287 }
288 }
289}
290
291impl AsCoreLatch for OnceLatch {
292 #[inline]
293 fn as_core_latch(&self) -> &CoreLatch {
294 &self.core_latch
295 }
296}
297
298/// Counting latches are used to implement scopes. They track a
299/// counter. Unlike other latches, calling `set()` does not
300/// necessarily make the latch be considered `set()`; instead, it just
301/// decrements the counter. The latch is only "set" (in the sense that
302/// `probe()` returns true) once the counter reaches zero.
303#[derive(Debug)]
304pub(super) struct CountLatch {
305 counter: AtomicUsize,
306 kind: CountLatchKind,
307}
308
309enum CountLatchKind {
310 /// A latch for scopes created on a rayon thread which will participate in work-
311 /// stealing while it waits for completion. This thread is not necessarily part
312 /// of the same registry as the scope itself!
313 Stealing {
314 latch: CoreLatch,
315 /// If a worker thread in registry A calls `in_place_scope` on a ThreadPool
316 /// with registry B, when a job completes in a thread of registry B, we may
317 /// need to call `notify_worker_latch_is_set()` to wake the thread in registry A.
318 /// That means we need a reference to registry A (since at that point we will
319 /// only have a reference to registry B), so we stash it here.
320 registry: Arc<Registry>,
321 /// The index of the worker to wake in `registry`
322 worker_index: usize,
323 },
324
325 /// A latch for scopes created on a non-rayon thread which will block to wait.
326 Blocking { latch: LockLatch },
327}
328
329impl std::fmt::Debug for CountLatchKind {
330 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
331 match self {
332 CountLatchKind::Stealing { latch, .. } => {
333 f.debug_tuple("Stealing").field(latch).finish()
334 }
335 CountLatchKind::Blocking { latch, .. } => {
336 f.debug_tuple("Blocking").field(latch).finish()
337 }
338 }
339 }
340}
341
342impl CountLatch {
343 pub(super) fn new(owner: Option<&WorkerThread>) -> Self {
344 Self::with_count(1, owner)
345 }
346
347 pub(super) fn with_count(count: usize, owner: Option<&WorkerThread>) -> Self {
348 Self {
349 counter: AtomicUsize::new(count),
350 kind: match owner {
351 Some(owner) => CountLatchKind::Stealing {
352 latch: CoreLatch::new(),
353 registry: Arc::clone(owner.registry()),
354 worker_index: owner.index(),
355 },
356 None => CountLatchKind::Blocking { latch: LockLatch::new() },
357 },
358 }
359 }
360
361 #[inline]
362 pub(super) fn increment(&self) {
363 let old_counter = self.counter.fetch_add(1, Ordering::Relaxed);
364 debug_assert!(old_counter != 0);
365 }
366
367 pub(super) fn wait(
368 &self,
369 owner: Option<&WorkerThread>,
370 all_jobs_started: impl FnMut() -> bool,
371 is_job: impl FnMut(&JobRef) -> bool,
372 ) {
373 match &self.kind {
374 CountLatchKind::Stealing { latch, registry, worker_index } => unsafe {
375 let owner = owner.expect("owner thread");
376 debug_assert_eq!(registry.id(), owner.registry().id());
377 debug_assert_eq!(*worker_index, owner.index());
378 owner.wait_for_jobs::<_, true>(latch, all_jobs_started, is_job, |job| {
379 owner.execute(job);
380 });
381 },
382 CountLatchKind::Blocking { latch } => latch.wait(),
383 }
384 }
385}
386
387impl Latch for CountLatch {
388 #[inline]
389 unsafe fn set(this: *const Self) {
390 if unsafe { (*this).counter.fetch_sub(1, Ordering::SeqCst) == 1 } {
391 // NOTE: Once we call `set` on the internal `latch`,
392 // the target may proceed and invalidate `this`!
393 match unsafe { &(*this).kind } {
394 CountLatchKind::Stealing { latch, registry, worker_index } => {
395 let registry = Arc::clone(registry);
396 if unsafe { CoreLatch::set(latch) } {
397 registry.notify_worker_latch_is_set(*worker_index);
398 }
399 }
400 CountLatchKind::Blocking { latch } => unsafe { LockLatch::set(latch) },
401 }
402 }
403 }
404}
405
406/// `&L` without any implication of `dereferenceable` for `Latch::set`
407pub(super) struct LatchRef<'a, L> {
408 inner: *const L,
409 marker: PhantomData<&'a L>,
410}
411
412impl<L> LatchRef<'_, L> {
413 pub(super) fn new(inner: &L) -> LatchRef<'_, L> {
414 LatchRef { inner, marker: PhantomData }
415 }
416}
417
418unsafe impl<L: Sync> Sync for LatchRef<'_, L> {}
419
420impl<L> Deref for LatchRef<'_, L> {
421 type Target = L;
422
423 fn deref(&self) -> &L {
424 // SAFETY: if we have &self, the inner latch is still alive
425 unsafe { &*self.inner }
426 }
427}
428
429impl<L: Latch> Latch for LatchRef<'_, L> {
430 #[inline]
431 unsafe fn set(this: *const Self) {
432 unsafe { L::set((*this).inner) };
433 }
434}