ax_task/thread/handle/
wake_batch.rs

1//! Allocation-free task-context wake batching.
2
3use alloc::{rc::Rc, sync::Arc};
4use core::{marker::PhantomData, mem::ManuallyDrop, ptr, sync::atomic::Ordering};
5
6use super::{ThreadCore, ThreadWakeHandle};
7
8/// A task-context wake list backed by nodes embedded in each thread core.
9///
10/// Queue owners add handles while holding their own metadata lock, release
11/// that lock, and then call [`Self::wake_all`]. This is the same ownership
12/// split as Linux `wake_q`: unlink and select atomically under the domain lock,
13/// perform scheduler wakes afterwards, and never allocate in between.
14///
15/// A thread can occur at most once in all live batches. Duplicate insertion is
16/// coalesced and returns `false`. The batch is deliberately neither `Send` nor
17/// `Sync`; the task context that selected the waiters must drain it.
18#[must_use = "selected threads must be woken after releasing the domain lock"]
19pub struct ThreadWakeBatch {
20    head: *const ThreadCore,
21    tail: *const ThreadCore,
22    len: usize,
23    _task_context: PhantomData<Rc<()>>,
24}
25
26impl ThreadWakeBatch {
27    /// Creates an empty batch without allocating.
28    pub const fn new() -> Self {
29        Self {
30            head: ptr::null(),
31            tail: ptr::null(),
32            len: 0,
33            _task_context: PhantomData,
34        }
35    }
36
37    /// Adds a wake handle, returning `false` when this thread is already in a
38    /// live batch.
39    pub fn push(&mut self, wake: ThreadWakeHandle) -> bool {
40        let core = &wake.core;
41        if core
42            .wake_batch_linked
43            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
44            .is_err()
45        {
46            return false;
47        }
48
49        let raw = Self::into_raw(wake);
50        unsafe {
51            // SAFETY: `raw` owns one strong reference transferred from `wake`.
52            // The successful linked transition gives this batch exclusive
53            // access to the embedded link until `pop` clears it.
54            (*raw)
55                .wake_batch_next
56                .store(ptr::null_mut(), Ordering::Relaxed);
57            if self.tail.is_null() {
58                self.head = raw;
59            } else {
60                (*self.tail)
61                    .wake_batch_next
62                    .store(raw.cast_mut(), Ordering::Release);
63            }
64        }
65        self.tail = raw;
66        self.len += 1;
67        true
68    }
69
70    /// Returns the number of unique threads selected by this batch.
71    pub const fn len(&self) -> usize {
72        self.len
73    }
74
75    /// Returns whether no thread has been selected.
76    pub const fn is_empty(&self) -> bool {
77        self.len == 0
78    }
79
80    /// Wakes all selected threads in FIFO order.
81    pub fn wake_all(mut self) -> usize {
82        let count = self.len;
83        while let Some(wake) = self.pop() {
84            let _result = wake.wake();
85        }
86        count
87    }
88
89    fn into_raw(wake: ThreadWakeHandle) -> *const ThreadCore {
90        let mut wake = ManuallyDrop::new(wake);
91        let core = unsafe {
92            // SAFETY: `wake` will not run Drop. Its core ownership is moved to
93            // the returned raw Arc and reconstructed exactly once by `pop`.
94            ManuallyDrop::take(&mut wake.core)
95        };
96        let reap_signal = unsafe {
97            // SAFETY: identical ownership transfer for the auxiliary Arc. The
98            // external lease remains owned by the batch and is released when
99            // the reconstructed handle is dropped.
100            ptr::read(&wake.reap_signal)
101        };
102        drop(reap_signal);
103        Arc::into_raw(core)
104    }
105
106    unsafe fn from_raw(raw: *const ThreadCore) -> ThreadWakeHandle {
107        let core = unsafe {
108            // SAFETY: every pointer placed in the batch came from one
109            // `Arc::into_raw`, and `pop` removes it exactly once.
110            Arc::from_raw(raw)
111        };
112        let reap_signal = Arc::clone(&core.reap_signal);
113        ThreadWakeHandle {
114            core: ManuallyDrop::new(core),
115            reap_signal,
116        }
117    }
118
119    fn pop(&mut self) -> Option<ThreadWakeHandle> {
120        let raw = self.head;
121        if raw.is_null() {
122            return None;
123        }
124
125        let next = unsafe {
126            // SAFETY: the raw Arc keeps the core alive and this batch has
127            // exclusive ownership of its embedded link.
128            (*raw).wake_batch_next.load(Ordering::Acquire).cast_const()
129        };
130        self.head = next;
131        if next.is_null() {
132            self.tail = ptr::null();
133        }
134        self.len -= 1;
135        unsafe {
136            // SAFETY: clearing the link completes this batch's exclusive node
137            // ownership before reconstructing the owning wake handle.
138            (*raw)
139                .wake_batch_next
140                .store(ptr::null_mut(), Ordering::Relaxed);
141            (*raw).wake_batch_linked.store(false, Ordering::Release);
142            Some(Self::from_raw(raw))
143        }
144    }
145}
146
147impl Default for ThreadWakeBatch {
148    fn default() -> Self {
149        Self::new()
150    }
151}
152
153impl Drop for ThreadWakeBatch {
154    fn drop(&mut self) {
155        // Dropping an undrained batch is a caller bug, but ownership must still
156        // be released without invoking scheduler callbacks from an unknown
157        // lock context.
158        let was_empty = self.is_empty();
159        while let Some(wake) = self.pop() {
160            drop(wake);
161        }
162        debug_assert!(was_empty, "thread wake batch was not drained");
163    }
164}
ax_task/thread/handle/wake_batch.rs

ax_task/thread/handle/
wake_batch.rs
