maitake_sync/util/
wake_batch.rs

1use super::CheckedMaybeUninit;
2use core::{ptr, task::Waker};
3
4/// A utility for waking multiple tasks in a batch, without reallocating.
5///
6/// This type is essentially an array of [`Waker`]s to which multiple tasks'
7/// [`Waker`]s can be added, until the array fills up. Once the array is full,
8/// all the tasks in the batch can be woken by calling [`WakeBatch::wake_all`],
9/// and the array refilled with new tasks.
10///
11/// This is useful when a lock must be held to remove [`Waker`]s from a queue
12/// (e.g. a [`cordyceps::List`]), but the lock can be released before the tasks
13/// are actually woken. Doing this repeatedly, rather than holding the lock for
14/// the entire wake process, may improve latency for other tasks that are
15/// attempting to access the lock. Additionally, it may avoid deadlocks that
16/// occur when a woken task will attempt to access the lock itself.
17pub(crate) struct WakeBatch {
18    init: usize,
19    wakers: [CheckedMaybeUninit<Waker>; MAX_WAKERS],
20}
21
22// 16 seems like a decent size for a stack array, could make this 32 (or a const
23// generic).
24//
25// when running loom tests, make the max much lower, so we can exercise behavior
26// involving multiple lock acquisitions.
27const MAX_WAKERS: usize = if cfg!(loom) { 2 } else { 16 };
28
29impl WakeBatch {
30    #[must_use]
31    pub(crate) const fn new() -> Self {
32        const INIT: CheckedMaybeUninit<Waker> = CheckedMaybeUninit::uninit();
33        Self {
34            init: 0,
35            wakers: [INIT; MAX_WAKERS],
36        }
37    }
38
39    /// Returns `true` if there is room for one or more additional [`Waker`] in
40    /// this batch.
41    ///
42    /// When this method returns `false`, [`WakeBatch::wake_all`] should be
43    /// called to wake all the wakers currently in the batch, and then this
44    /// method will return `true` again.
45    #[inline]
46    pub(crate) fn can_add_waker(&self) -> bool {
47        self.init < MAX_WAKERS
48    }
49
50    /// Adds a [`Waker`] to the batch, returning `true` if the batch still has
51    /// capacity for an additional waker.
52    ///
53    /// When this method returns `false`, the most recently added waker used the
54    /// last slot in the batch, and [`WakeBatch::wake_all`] must be called
55    /// before continuing to add new [`Waker`]s.
56    pub(crate) fn add_waker(&mut self, waker: Waker) -> bool {
57        debug_assert!(self.can_add_waker());
58        unsafe {
59            self.wakers.get_unchecked_mut(self.init).write(waker);
60        }
61        self.init += 1;
62        self.can_add_waker()
63    }
64
65    /// Wake all the tasks whose [`Waker`]s are currently in the batch.
66    pub(crate) fn wake_all(&mut self) {
67        let init = self.init;
68        self.init = 0;
69        for waker in self.wakers[..init].iter_mut() {
70            unsafe {
71                waker.as_mut_ptr().read().wake();
72            }
73        }
74    }
75}
76
77impl Drop for WakeBatch {
78    fn drop(&mut self) {
79        let slice =
80            ptr::slice_from_raw_parts_mut(self.wakers.as_mut_ptr() as *mut Waker, self.init);
81        unsafe { ptr::drop_in_place(slice) };
82    }
83}
maitake_sync/util/wake_batch.rs

maitake_sync/util/
wake_batch.rs
