tokio/util/
wake_list.rs

1use core::mem::MaybeUninit;
2use core::ptr;
3use std::task::Waker;
4
5const NUM_WAKERS: usize = 32;
6
7/// A list of wakers to be woken.
8///
9/// # Invariants
10///
11/// The first `curr` elements of `inner` are initialized.
12pub(crate) struct WakeList {
13    inner: [MaybeUninit<Waker>; NUM_WAKERS],
14    curr: usize,
15}
16
17impl WakeList {
18    pub(crate) fn new() -> Self {
19        const UNINIT_WAKER: MaybeUninit<Waker> = MaybeUninit::uninit();
20
21        Self {
22            inner: [UNINIT_WAKER; NUM_WAKERS],
23            curr: 0,
24        }
25    }
26
27    #[inline]
28    pub(crate) fn can_push(&self) -> bool {
29        self.curr < NUM_WAKERS
30    }
31
32    pub(crate) fn push(&mut self, val: Waker) {
33        debug_assert!(self.can_push());
34
35        self.inner[self.curr] = MaybeUninit::new(val);
36        self.curr += 1;
37    }
38
39    pub(crate) fn wake_all(&mut self) {
40        struct DropGuard {
41            start: *mut Waker,
42            end: *mut Waker,
43        }
44
45        impl Drop for DropGuard {
46            fn drop(&mut self) {
47                // SAFETY: Both pointers are part of the same object, with `start <= end`.
48                let len = unsafe { self.end.offset_from(self.start) } as usize;
49                let slice = ptr::slice_from_raw_parts_mut(self.start, len);
50                // SAFETY: All elements in `start..len` are initialized, so we can drop them.
51                unsafe { ptr::drop_in_place(slice) };
52            }
53        }
54
55        debug_assert!(self.curr <= NUM_WAKERS);
56
57        let mut guard = {
58            let start = self.inner.as_mut_ptr().cast::<Waker>();
59            // SAFETY: The resulting pointer is in bounds or one after the length of the same object.
60            let end = unsafe { start.add(self.curr) };
61            // Transfer ownership of the wakers in `inner` to `DropGuard`.
62            self.curr = 0;
63            DropGuard { start, end }
64        };
65        while !ptr::eq(guard.start, guard.end) {
66            // SAFETY: `start` is always initialized if `start != end`.
67            let waker = unsafe { ptr::read(guard.start) };
68            // SAFETY: The resulting pointer is in bounds or one after the length of the same object.
69            guard.start = unsafe { guard.start.add(1) };
70            // If this panics, then `guard` will clean up the remaining wakers.
71            waker.wake();
72        }
73    }
74}
75
76impl Drop for WakeList {
77    fn drop(&mut self) {
78        let slice =
79            ptr::slice_from_raw_parts_mut(self.inner.as_mut_ptr().cast::<Waker>(), self.curr);
80        // SAFETY: The first `curr` elements are initialized, so we can drop them.
81        unsafe { ptr::drop_in_place(slice) };
82    }
83}
tokio/util/wake_list.rs

tokio/util/
wake_list.rs
