ntex_util/
task.rs

1//! Task wake-up and scheduling utilities.
2use std::{cell::Cell, fmt, task::Waker};
3
4/// A single-threaded synchronization primitive for task wake-up.
5///
6/// Sometimes the task interested in a given event will change over time.
7/// A `LocalWaker` can coordinate notifications with the consumer
8/// potentially "updating" the underlying task to wake up. This is useful in
9/// scenarios where a computation completes in another task and wants to
10/// notify the consumer, but the consumer is in the process of being migrated to
11/// a new logical task.
12///
13/// Consumers should call `register` before checking the result of a computation
14/// and producers should call `wake` after producing the computation (this
15/// differs from the usual `thread::park` pattern). It is also permitted for
16/// `wake` to be called **before** `register`. This results in a no-op.
17///
18/// A single `LocalWaker` may be reused for any number of calls to `register` or
19/// `wake`.
20///
21/// The slot and its owned local handle must remain on their acquiring thread.
22/// Use [`take`](Self::take) to transfer a registered standard waker instead.
23///
24/// ```compile_fail
25/// fn assert_send<T: Send>() {}
26/// assert_send::<ntex_util::task::LocalWaker>();
27/// ```
28///
29/// ```compile_fail
30/// fn assert_sync<T: Sync>() {}
31/// assert_sync::<ntex_util::task::LocalWaker>();
32/// ```
33#[derive(Default)]
34pub struct LocalWaker {
35    waker: Cell<Option<ntex_rt::LocalWaker>>,
36}
37
38#[cfg(target_pointer_width = "64")]
39const _: () = assert!(size_of::<LocalWaker>() == 16);
40
41impl LocalWaker {
42    /// Creates an empty `LocalWaker`.
43    pub fn new() -> Self {
44        LocalWaker::with(None)
45    }
46
47    /// Creates a `LocalWaker` with an optional initial waker.
48    pub fn with(waker: Option<Waker>) -> Self {
49        LocalWaker {
50            waker: Cell::new(waker.map(ntex_rt::LocalWaker::from)),
51        }
52    }
53
54    #[inline]
55    /// Registers the waker to be notified on calls to `wake`.
56    ///
57    /// Returns `true` if a waker was already registered.
58    ///
59    /// If cloning the replacement panics, the previous registration is retained
60    /// unless the clone callback itself changes the slot.
61    pub fn register(&self, waker: &Waker) -> bool {
62        let prev = self.waker.take();
63        let registered = prev.is_some();
64        let same = prev.as_ref().is_some_and(|prev| prev.will_wake(waker));
65        self.waker.set(prev);
66        if !same {
67            // Clone and drop callbacks may reenter the slot. Keep the old
68            // registration until cloning succeeds, then publish before dropping.
69            let next = ntex_rt::LocalWaker::from_waker(waker);
70            drop(self.waker.replace(Some(next)));
71        }
72        registered
73    }
74
75    /// Registers `waker` in the slot returned by `slot`, like [`register`](Self::register).
76    ///
77    /// Waker clone and drop callbacks may reenter the owner of the slot and
78    /// move or mutate it. `slot` is called again after every callback, so it
79    /// must derive a fresh reference from the owner each time. No reference
80    /// to the slot is kept while a callback runs.
81    pub(crate) fn register_in<'a>(slot: impl Fn() -> &'a LocalWaker, waker: &Waker) -> bool {
82        let prev = slot().waker.take();
83        let registered = prev.is_some();
84        let same = prev.as_ref().is_some_and(|prev| prev.will_wake(waker));
85        slot().waker.set(prev);
86        if !same {
87            let next = ntex_rt::LocalWaker::from_waker(waker);
88            let prev = slot().waker.replace(Some(next));
89            drop(prev);
90        }
91        registered
92    }
93
94    #[inline]
95    /// Calls `wake` on the last `Waker` passed to `register`.
96    ///
97    /// The waker is consumed, so a later `wake` does nothing until `register`
98    /// is called again. If no waker is registered, this does nothing.
99    pub fn wake(&self) {
100        if let Some(waker) = self.take_local() {
101            waker.wake();
102        }
103    }
104
105    #[inline]
106    /// Calls `wake` on the last `Waker` passed to `register`.
107    ///
108    /// Returns `false` if no waker is registered. Like [`wake`](Self::wake),
109    /// this consumes the registered waker.
110    pub fn wake_checked(&self) -> bool {
111        if let Some(waker) = self.take_local() {
112            waker.wake();
113            true
114        } else {
115            false
116        }
117    }
118
119    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
120    ///
121    /// This empties the slot and converts the owned handle into a thread-safe
122    /// standard waker without cloning. If no waker is registered, returns `None`.
123    /// Use [`take_local`](Self::take_local) to retain native local operations.
124    pub fn take(&self) -> Option<Waker> {
125        self.take_local().map(ntex_rt::LocalWaker::into_waker)
126    }
127
128    /// Takes the registered wake handle, leaving the slot empty.
129    ///
130    /// The returned handle retains native local operations and cannot be sent
131    /// to another thread. If no waker is registered, returns `None`.
132    #[inline]
133    pub fn take_local(&self) -> Option<ntex_rt::LocalWaker> {
134        self.waker.take()
135    }
136
137    #[doc(hidden)]
138    /// Returns `true` if a waker is registered.
139    pub fn is_set(&self) -> bool {
140        let waker = self.waker.take();
141        let set = waker.is_some();
142        self.waker.set(waker);
143        set
144    }
145}
146
147/// Cloning creates an empty `LocalWaker`; the registered waker is not copied.
148impl Clone for LocalWaker {
149    fn clone(&self) -> Self {
150        LocalWaker::new()
151    }
152}
153
154impl fmt::Debug for LocalWaker {
155    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
156        write!(f, "LocalWaker")
157    }
158}
159
160/// Yields execution back to the current runtime.
161pub async fn yield_to() {
162    use std::{future::Future, pin::Pin, task::Context, task::Poll};
163
164    struct Yield {
165        completed: bool,
166    }
167
168    impl Future for Yield {
169        type Output = ();
170
171        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
172            if self.completed {
173                return Poll::Ready(());
174            }
175
176            self.completed = true;
177            cx.waker().wake_by_ref();
178
179            Poll::Pending
180        }
181    }
182
183    Yield { completed: false }.await;
184}
185
186#[cfg(test)]
187mod test {
188    use super::*;
189
190    /// STATEMENT: A task can resume after yielding.
191    ///
192    /// HOW:
193    /// - Await `yield_to` inside an asynchronous test and let the test complete.
194    #[ntex::test]
195    async fn yield_test() {
196        yield_to().await;
197    }
198
199    /// STATEMENT: Local and standard slot wake paths reschedule the current task.
200    ///
201    /// HOW:
202    /// - Register and wake through the slot, a taken local waker, and a taken standard waker.
203    /// - Require a fourth poll to finish and verify the slot is empty afterward.
204    #[ntex::test]
205    async fn local_slot_wakes_current_task() {
206        crate::spawn(async {
207            let slot = LocalWaker::new();
208            let mut polls = 0;
209            std::future::poll_fn(|cx| {
210                polls += 1;
211                if polls == 4 {
212                    return std::task::Poll::Ready(());
213                }
214                assert!(!slot.register(cx.waker()));
215                match polls {
216                    1 => slot.wake(),
217                    2 => slot.take_local().unwrap().wake(),
218                    _ => slot.take().unwrap().wake(),
219                }
220                std::task::Poll::Pending
221            })
222            .await;
223            assert!(!slot.is_set());
224        })
225        .await
226        .unwrap();
227    }
228}
229
230/// Wakers whose callbacks run one-shot hooks, for reentrancy tests.
231#[cfg(test)]
232pub(crate) mod hooks {
233    use std::cell::Cell;
234    use std::task::{RawWaker, RawWakerVTable, Waker};
235    use std::thread::LocalKey;
236
237    pub(crate) type Hook = Cell<Option<Box<dyn FnOnce()>>>;
238
239    // Waker callbacks may run on any thread; local test closures stay in TLS.
240    thread_local! {
241        pub(crate) static ON_CLONE: Hook = Cell::new(None);
242        pub(crate) static ON_DROP: Hook = Cell::new(None);
243        pub(crate) static ON_WAKE: Hook = Cell::new(None);
244    }
245
246    fn run_hook(key: &'static LocalKey<Hook>) {
247        if let Some(hook) = key.with(Cell::take) {
248            hook();
249        }
250    }
251
252    /// Returns a waker that runs the pending `ON_CLONE`, `ON_WAKE` or `ON_DROP` hook.
253    pub(crate) fn hook_waker() -> Waker {
254        static HOOK_VTABLE: RawWakerVTable = RawWakerVTable::new(
255            |p| {
256                run_hook(&ON_CLONE);
257                RawWaker::new(p, &HOOK_VTABLE)
258            },
259            |_| run_hook(&ON_WAKE),
260            |_| run_hook(&ON_WAKE),
261            |_| run_hook(&ON_DROP),
262        );
263        unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &HOOK_VTABLE)) }
264    }
265
266    /// Wake counter for wakers created by [`counted_waker`].
267    #[derive(Default)]
268    pub(crate) struct WakeCount(pub(crate) std::sync::atomic::AtomicUsize);
269
270    impl std::task::Wake for WakeCount {
271        fn wake(self: std::sync::Arc<Self>) {
272            self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
273        }
274    }
275
276    impl WakeCount {
277        pub(crate) fn get(&self) -> usize {
278            self.0.load(std::sync::atomic::Ordering::Relaxed)
279        }
280    }
281
282    /// Returns a distinct waker and the counter of its wake-ups.
283    pub(crate) fn counted_waker() -> (std::sync::Arc<WakeCount>, Waker) {
284        let count = std::sync::Arc::new(WakeCount::default());
285        (count.clone(), Waker::from(count))
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use std::panic::{AssertUnwindSafe, catch_unwind};
292    use std::rc::Rc;
293    use std::sync::Arc;
294    use std::sync::atomic::{AtomicUsize, Ordering};
295    use std::task::{RawWaker, RawWakerVTable, Wake};
296
297    use super::hooks::{ON_CLONE, ON_DROP, ON_WAKE, hook_waker};
298    use super::*;
299
300    static CLONES: AtomicUsize = AtomicUsize::new(0);
301
302    static VTABLE: RawWakerVTable = RawWakerVTable::new(
303        |p| {
304            CLONES.fetch_add(1, Ordering::Relaxed);
305            RawWaker::new(p, &VTABLE)
306        },
307        |_| {},
308        |_| {},
309        |_| {},
310    );
311
312    #[derive(Default)]
313    struct WakeCounter(AtomicUsize);
314
315    impl Wake for WakeCounter {
316        fn wake(self: Arc<Self>) {
317            self.0.fetch_add(1, Ordering::Relaxed);
318        }
319    }
320
321    fn counted_waker() -> (Arc<WakeCounter>, Waker) {
322        let count = Arc::new(WakeCounter::default());
323        let waker = Waker::from(count.clone());
324        (count, waker)
325    }
326
327    /// STATEMENT: Registering an equivalent waker avoids cloning it again.
328    ///
329    /// HOW:
330    /// - Repeatedly register one raw waker, then replace it with a distinct one.
331    /// - Check clone counts and verify the taken waker matches the replacement.
332    #[test]
333    fn test_register_same_waker() {
334        static A: u8 = 0;
335        static B: u8 = 0;
336        let a = unsafe { Waker::from_raw(RawWaker::new((&raw const A).cast(), &VTABLE)) };
337        let b = unsafe { Waker::from_raw(RawWaker::new((&raw const B).cast(), &VTABLE)) };
338
339        let w = LocalWaker::new();
340        assert!(!w.register(&a));
341        assert!(w.register(&a));
342        assert!(w.register(&a));
343        assert_eq!(CLONES.load(Ordering::Relaxed), 1);
344
345        assert!(w.register(&b));
346        assert_eq!(CLONES.load(Ordering::Relaxed), 2);
347        assert!(w.take().unwrap().will_wake(&b));
348    }
349
350    /// STATEMENT: Cloning a `LocalWaker` creates an empty slot without clearing the original.
351    ///
352    /// HOW:
353    /// - Clone a registered slot and verify only the original remains set.
354    /// - Check waking the empty clone returns false and inspect the original's debug output.
355    #[test]
356    fn local_waker_clone_is_empty() {
357        let waker = LocalWaker::new();
358        waker.register(std::task::Waker::noop());
359        assert!(waker.is_set());
360        let cloned = waker.clone();
361        assert!(!cloned.is_set());
362        assert!(!cloned.wake_checked());
363        assert!(waker.is_set());
364        assert_eq!(format!("{waker:?}"), "LocalWaker");
365    }
366
367    /// STATEMENT: Owning waker conversions and wake operations consume the slot without cloning.
368    ///
369    /// HOW:
370    /// - Install a clone hook that panics and exercise four owning wake paths.
371    /// - Verify four wakes, empty slots, and an unused clone hook.
372    #[test]
373    fn owning_conversions_and_wakes_do_not_clone() {
374        ON_CLONE.set(Some(Box::new(|| panic!("unexpected clone"))));
375        let calls = Rc::new(Cell::new(0));
376        for path in 0..4 {
377            let calls = calls.clone();
378            ON_WAKE.set(Some(Box::new(move || calls.set(calls.get() + 1))));
379            let slot = LocalWaker::with(Some(hook_waker()));
380            match path {
381                0 => slot.take().unwrap().wake(),
382                1 => slot.take_local().unwrap().wake(),
383                2 => slot.wake(),
384                _ => assert!(slot.wake_checked()),
385            }
386            assert!(slot.take().is_none());
387            assert!(slot.take_local().is_none());
388            assert!(!slot.wake_checked());
389        }
390        assert_eq!(calls.get(), 4);
391        assert!(ON_CLONE.with(Cell::take).is_some());
392    }
393
394    /// STATEMENT: A replacement-waker clone panic preserves the previous registration.
395    ///
396    /// HOW:
397    /// - Register a counted waker, then catch a replacement's clone panic.
398    /// - Wake the slot and verify the original waker receives exactly one wake.
399    #[test]
400    fn clone_panic_preserves_registration() {
401        let (calls, original) = counted_waker();
402        let slot = LocalWaker::with(Some(original));
403        let replacement = hook_waker();
404        ON_CLONE.set(Some(Box::new(|| panic!("clone"))));
405        assert!(catch_unwind(AssertUnwindSafe(|| slot.register(&replacement))).is_err());
406        assert!(slot.wake_checked());
407        assert_eq!(calls.0.load(Ordering::Relaxed), 1);
408    }
409
410    /// STATEMENT: Reentrant clone callbacks retain nested registration when replacement unwinds.
411    ///
412    /// HOW:
413    /// - During replacement cloning, wake the previous waker and register a nested one.
414    /// - Exercise normal and panicking callbacks and check which registration survives.
415    #[test]
416    fn clone_reentry_preserves_nested_changes_on_unwind() {
417        for panic in [false, true] {
418            let (previous_calls, previous) = counted_waker();
419            let (nested_calls, nested) = counted_waker();
420            let slot = Rc::new(LocalWaker::with(Some(previous)));
421            let callback_slot = slot.clone();
422            ON_CLONE.set(Some(Box::new(move || {
423                assert!(callback_slot.wake_checked());
424                assert!(!callback_slot.register(&nested));
425                assert!(!panic, "clone");
426            })));
427            let replacement = hook_waker();
428            let result = catch_unwind(AssertUnwindSafe(|| slot.register(&replacement)));
429            if panic {
430                assert!(result.is_err());
431            } else {
432                assert!(result.unwrap());
433            }
434            assert_eq!(previous_calls.0.load(Ordering::Relaxed), 1);
435            assert!(slot.wake_checked());
436            assert_eq!(nested_calls.0.load(Ordering::Relaxed), usize::from(panic));
437        }
438    }
439
440    /// STATEMENT: Reentrant drop callbacks can replace a newly registered waker even when they panic.
441    ///
442    /// HOW:
443    /// - Drop the old waker during replacement and have its callback wake and replace the new one.
444    /// - Check replacement and nested wake counts with and without callback unwind.
445    #[test]
446    fn replaced_drop_reentry_and_unwind() {
447        for panic in [false, true] {
448            let (replacement_calls, replacement) = counted_waker();
449            let (nested_calls, nested) = counted_waker();
450            let slot = Rc::new(LocalWaker::with(Some(hook_waker())));
451            let callback_slot = slot.clone();
452            ON_DROP.set(Some(Box::new(move || {
453                assert!(callback_slot.wake_checked());
454                assert!(!callback_slot.register(&nested));
455                assert!(!panic, "drop");
456            })));
457            let result = catch_unwind(AssertUnwindSafe(|| slot.register(&replacement)));
458            if panic {
459                assert!(result.is_err());
460            } else {
461                assert!(result.unwrap());
462            }
463            assert_eq!(replacement_calls.0.load(Ordering::Relaxed), 1);
464            assert!(slot.wake_checked());
465            assert_eq!(nested_calls.0.load(Ordering::Relaxed), 1);
466        }
467    }
468
469    /// STATEMENT: Wake callbacks see an empty slot and can register a surviving replacement.
470    ///
471    /// HOW:
472    /// - Exercise checked and unchecked wakes with callbacks that register a nested waker.
473    /// - Repeat with panics and verify the nested registration still receives one wake.
474    #[test]
475    fn wake_reentry_and_unwind() {
476        for panic in [false, true] {
477            for checked in [false, true] {
478                let (calls, nested) = counted_waker();
479                let slot = Rc::new(LocalWaker::with(Some(hook_waker())));
480                let callback_slot = slot.clone();
481                ON_WAKE.set(Some(Box::new(move || {
482                    assert!(!callback_slot.is_set());
483                    assert!(!callback_slot.register(&nested));
484                    assert!(!panic, "wake");
485                })));
486                let result = catch_unwind(AssertUnwindSafe(|| {
487                    if checked {
488                        assert!(slot.wake_checked());
489                    } else {
490                        slot.wake();
491                    }
492                }));
493                assert_eq!(result.is_err(), panic);
494                assert!(slot.wake_checked());
495                assert_eq!(calls.0.load(Ordering::Relaxed), 1);
496            }
497        }
498    }
499
500    /// STATEMENT: Native task wakers support local and cross-thread delivery and remain safe after shutdown.
501    ///
502    /// HOW:
503    /// - Poll a native runtime task, wake it locally and from another thread, and count three polls.
504    /// - Drop the runtime and check task completion, released task state, and one final consumed wake.
505    #[test]
506    fn native_slot_delivery_and_shutdown() {
507        #[derive(Debug)]
508        struct NoopNotify;
509
510        impl ntex_rt::Notify for NoopNotify {
511            fn notify(&self) -> std::io::Result<()> {
512                Ok(())
513            }
514        }
515
516        let rt = ntex_rt::Runtime::new(Box::new(NoopNotify));
517        let slot = Rc::new(LocalWaker::new());
518        let polls = Rc::new(Cell::new(0));
519        let (task_slot, task_polls) = (slot.clone(), polls.clone());
520        let task = rt.spawn(std::future::poll_fn(move |cx| {
521            task_polls.set(task_polls.get() + 1);
522            task_slot.register(cx.waker());
523            std::task::Poll::<()>::Pending
524        }));
525        rt.poll();
526        slot.take_local().unwrap().wake();
527        rt.poll();
528        let standard = slot.take().unwrap();
529        std::thread::spawn(move || standard.wake()).join().unwrap();
530        rt.poll();
531        assert_eq!(polls.get(), 3);
532        assert_eq!(Rc::strong_count(&slot), 2);
533        drop(rt);
534        assert!(task.is_finished());
535        assert_eq!(Rc::strong_count(&slot), 1);
536        assert!(slot.wake_checked());
537        assert!(!slot.wake_checked());
538    }
539}
ntex_util/task.rs

ntex_util/
task.rs
