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    #[inline]
76    /// Calls `wake` on the last `Waker` passed to `register`.
77    ///
78    /// The waker is consumed, so a later `wake` does nothing until `register`
79    /// is called again. If no waker is registered, this does nothing.
80    pub fn wake(&self) {
81        if let Some(waker) = self.take_local() {
82            waker.wake();
83        }
84    }
85
86    #[inline]
87    /// Calls `wake` on the last `Waker` passed to `register`.
88    ///
89    /// Returns `false` if no waker is registered. Like [`wake`](Self::wake),
90    /// this consumes the registered waker.
91    pub fn wake_checked(&self) -> bool {
92        if let Some(waker) = self.take_local() {
93            waker.wake();
94            true
95        } else {
96            false
97        }
98    }
99
100    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
101    ///
102    /// This empties the slot and converts the owned handle into a thread-safe
103    /// standard waker without cloning. If no waker is registered, returns `None`.
104    /// Use [`take_local`](Self::take_local) to retain native local operations.
105    pub fn take(&self) -> Option<Waker> {
106        self.take_local().map(ntex_rt::LocalWaker::into_waker)
107    }
108
109    /// Takes the registered wake handle, leaving the slot empty.
110    ///
111    /// The returned handle retains native local operations and cannot be sent
112    /// to another thread. If no waker is registered, returns `None`.
113    #[inline]
114    pub fn take_local(&self) -> Option<ntex_rt::LocalWaker> {
115        self.waker.take()
116    }
117
118    #[doc(hidden)]
119    /// Returns `true` if a waker is registered.
120    pub fn is_set(&self) -> bool {
121        let waker = self.waker.take();
122        let set = waker.is_some();
123        self.waker.set(waker);
124        set
125    }
126}
127
128/// Cloning creates an empty `LocalWaker`; the registered waker is not copied.
129impl Clone for LocalWaker {
130    fn clone(&self) -> Self {
131        LocalWaker::new()
132    }
133}
134
135impl fmt::Debug for LocalWaker {
136    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        write!(f, "LocalWaker")
138    }
139}
140
141/// Yields execution back to the current runtime.
142pub async fn yield_to() {
143    use std::{future::Future, pin::Pin, task::Context, task::Poll};
144
145    struct Yield {
146        completed: bool,
147    }
148
149    impl Future for Yield {
150        type Output = ();
151
152        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
153            if self.completed {
154                return Poll::Ready(());
155            }
156
157            self.completed = true;
158            cx.waker().wake_by_ref();
159
160            Poll::Pending
161        }
162    }
163
164    Yield { completed: false }.await;
165}
166
167#[cfg(test)]
168mod test {
169    use super::*;
170
171    /// STATEMENT: A task can resume after yielding.
172    ///
173    /// HOW:
174    /// - Await `yield_to` inside an asynchronous test and let the test complete.
175    #[ntex::test]
176    async fn yield_test() {
177        yield_to().await;
178    }
179
180    /// STATEMENT: Local and standard slot wake paths reschedule the current task.
181    ///
182    /// HOW:
183    /// - Register and wake through the slot, a taken local waker, and a taken standard waker.
184    /// - Require a fourth poll to finish and verify the slot is empty afterward.
185    #[ntex::test]
186    async fn local_slot_wakes_current_task() {
187        crate::spawn(async {
188            let slot = LocalWaker::new();
189            let mut polls = 0;
190            std::future::poll_fn(|cx| {
191                polls += 1;
192                if polls == 4 {
193                    return std::task::Poll::Ready(());
194                }
195                assert!(!slot.register(cx.waker()));
196                match polls {
197                    1 => slot.wake(),
198                    2 => slot.take_local().unwrap().wake(),
199                    _ => slot.take().unwrap().wake(),
200                }
201                std::task::Poll::Pending
202            })
203            .await;
204            assert!(!slot.is_set());
205        })
206        .await
207        .unwrap();
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use std::panic::{AssertUnwindSafe, catch_unwind};
214    use std::rc::Rc;
215    use std::sync::Arc;
216    use std::sync::atomic::{AtomicUsize, Ordering};
217    use std::task::{RawWaker, RawWakerVTable, Wake};
218    use std::thread::LocalKey;
219
220    use super::*;
221
222    static CLONES: AtomicUsize = AtomicUsize::new(0);
223
224    static VTABLE: RawWakerVTable = RawWakerVTable::new(
225        |p| {
226            CLONES.fetch_add(1, Ordering::Relaxed);
227            RawWaker::new(p, &VTABLE)
228        },
229        |_| {},
230        |_| {},
231        |_| {},
232    );
233
234    type Hook = Cell<Option<Box<dyn FnOnce()>>>;
235
236    // Waker callbacks may run on any thread; local test closures stay in TLS.
237    thread_local! {
238        static ON_CLONE: Hook = Cell::new(None);
239        static ON_DROP: Hook = Cell::new(None);
240        static ON_WAKE: Hook = Cell::new(None);
241    }
242
243    fn run_hook(key: &'static LocalKey<Hook>) {
244        if let Some(hook) = key.with(Cell::take) {
245            hook();
246        }
247    }
248
249    fn hook_waker() -> Waker {
250        static HOOK_VTABLE: RawWakerVTable = RawWakerVTable::new(
251            |p| {
252                run_hook(&ON_CLONE);
253                RawWaker::new(p, &HOOK_VTABLE)
254            },
255            |_| run_hook(&ON_WAKE),
256            |_| run_hook(&ON_WAKE),
257            |_| run_hook(&ON_DROP),
258        );
259        unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &HOOK_VTABLE)) }
260    }
261
262    #[derive(Default)]
263    struct WakeCounter(AtomicUsize);
264
265    impl Wake for WakeCounter {
266        fn wake(self: Arc<Self>) {
267            self.0.fetch_add(1, Ordering::Relaxed);
268        }
269    }
270
271    fn counted_waker() -> (Arc<WakeCounter>, Waker) {
272        let count = Arc::new(WakeCounter::default());
273        let waker = Waker::from(count.clone());
274        (count, waker)
275    }
276
277    /// STATEMENT: Registering an equivalent waker avoids cloning it again.
278    ///
279    /// HOW:
280    /// - Repeatedly register one raw waker, then replace it with a distinct one.
281    /// - Check clone counts and verify the taken waker matches the replacement.
282    #[test]
283    fn test_register_same_waker() {
284        static A: u8 = 0;
285        static B: u8 = 0;
286        let a = unsafe { Waker::from_raw(RawWaker::new((&raw const A).cast(), &VTABLE)) };
287        let b = unsafe { Waker::from_raw(RawWaker::new((&raw const B).cast(), &VTABLE)) };
288
289        let w = LocalWaker::new();
290        assert!(!w.register(&a));
291        assert!(w.register(&a));
292        assert!(w.register(&a));
293        assert_eq!(CLONES.load(Ordering::Relaxed), 1);
294
295        assert!(w.register(&b));
296        assert_eq!(CLONES.load(Ordering::Relaxed), 2);
297        assert!(w.take().unwrap().will_wake(&b));
298    }
299
300    /// STATEMENT: Cloning a `LocalWaker` creates an empty slot without clearing the original.
301    ///
302    /// HOW:
303    /// - Clone a registered slot and verify only the original remains set.
304    /// - Check waking the empty clone returns false and inspect the original's debug output.
305    #[test]
306    fn local_waker_clone_is_empty() {
307        let waker = LocalWaker::new();
308        waker.register(std::task::Waker::noop());
309        assert!(waker.is_set());
310        let cloned = waker.clone();
311        assert!(!cloned.is_set());
312        assert!(!cloned.wake_checked());
313        assert!(waker.is_set());
314        assert_eq!(format!("{waker:?}"), "LocalWaker");
315    }
316
317    /// STATEMENT: Owning waker conversions and wake operations consume the slot without cloning.
318    ///
319    /// HOW:
320    /// - Install a clone hook that panics and exercise four owning wake paths.
321    /// - Verify four wakes, empty slots, and an unused clone hook.
322    #[test]
323    fn owning_conversions_and_wakes_do_not_clone() {
324        ON_CLONE.set(Some(Box::new(|| panic!("unexpected clone"))));
325        let calls = Rc::new(Cell::new(0));
326        for path in 0..4 {
327            let calls = calls.clone();
328            ON_WAKE.set(Some(Box::new(move || calls.set(calls.get() + 1))));
329            let slot = LocalWaker::with(Some(hook_waker()));
330            match path {
331                0 => slot.take().unwrap().wake(),
332                1 => slot.take_local().unwrap().wake(),
333                2 => slot.wake(),
334                _ => assert!(slot.wake_checked()),
335            }
336            assert!(slot.take().is_none());
337            assert!(slot.take_local().is_none());
338            assert!(!slot.wake_checked());
339        }
340        assert_eq!(calls.get(), 4);
341        assert!(ON_CLONE.with(Cell::take).is_some());
342    }
343
344    /// STATEMENT: A replacement-waker clone panic preserves the previous registration.
345    ///
346    /// HOW:
347    /// - Register a counted waker, then catch a replacement's clone panic.
348    /// - Wake the slot and verify the original waker receives exactly one wake.
349    #[test]
350    fn clone_panic_preserves_registration() {
351        let (calls, original) = counted_waker();
352        let slot = LocalWaker::with(Some(original));
353        let replacement = hook_waker();
354        ON_CLONE.set(Some(Box::new(|| panic!("clone"))));
355        assert!(catch_unwind(AssertUnwindSafe(|| slot.register(&replacement))).is_err());
356        assert!(slot.wake_checked());
357        assert_eq!(calls.0.load(Ordering::Relaxed), 1);
358    }
359
360    /// STATEMENT: Reentrant clone callbacks retain nested registration when replacement unwinds.
361    ///
362    /// HOW:
363    /// - During replacement cloning, wake the previous waker and register a nested one.
364    /// - Exercise normal and panicking callbacks and check which registration survives.
365    #[test]
366    fn clone_reentry_preserves_nested_changes_on_unwind() {
367        for panic in [false, true] {
368            let (previous_calls, previous) = counted_waker();
369            let (nested_calls, nested) = counted_waker();
370            let slot = Rc::new(LocalWaker::with(Some(previous)));
371            let callback_slot = slot.clone();
372            ON_CLONE.set(Some(Box::new(move || {
373                assert!(callback_slot.wake_checked());
374                assert!(!callback_slot.register(&nested));
375                assert!(!panic, "clone");
376            })));
377            let replacement = hook_waker();
378            let result = catch_unwind(AssertUnwindSafe(|| slot.register(&replacement)));
379            if panic {
380                assert!(result.is_err());
381            } else {
382                assert!(result.unwrap());
383            }
384            assert_eq!(previous_calls.0.load(Ordering::Relaxed), 1);
385            assert!(slot.wake_checked());
386            assert_eq!(nested_calls.0.load(Ordering::Relaxed), usize::from(panic));
387        }
388    }
389
390    /// STATEMENT: Reentrant drop callbacks can replace a newly registered waker even when they panic.
391    ///
392    /// HOW:
393    /// - Drop the old waker during replacement and have its callback wake and replace the new one.
394    /// - Check replacement and nested wake counts with and without callback unwind.
395    #[test]
396    fn replaced_drop_reentry_and_unwind() {
397        for panic in [false, true] {
398            let (replacement_calls, replacement) = counted_waker();
399            let (nested_calls, nested) = counted_waker();
400            let slot = Rc::new(LocalWaker::with(Some(hook_waker())));
401            let callback_slot = slot.clone();
402            ON_DROP.set(Some(Box::new(move || {
403                assert!(callback_slot.wake_checked());
404                assert!(!callback_slot.register(&nested));
405                assert!(!panic, "drop");
406            })));
407            let result = catch_unwind(AssertUnwindSafe(|| slot.register(&replacement)));
408            if panic {
409                assert!(result.is_err());
410            } else {
411                assert!(result.unwrap());
412            }
413            assert_eq!(replacement_calls.0.load(Ordering::Relaxed), 1);
414            assert!(slot.wake_checked());
415            assert_eq!(nested_calls.0.load(Ordering::Relaxed), 1);
416        }
417    }
418
419    /// STATEMENT: Wake callbacks see an empty slot and can register a surviving replacement.
420    ///
421    /// HOW:
422    /// - Exercise checked and unchecked wakes with callbacks that register a nested waker.
423    /// - Repeat with panics and verify the nested registration still receives one wake.
424    #[test]
425    fn wake_reentry_and_unwind() {
426        for panic in [false, true] {
427            for checked in [false, true] {
428                let (calls, nested) = counted_waker();
429                let slot = Rc::new(LocalWaker::with(Some(hook_waker())));
430                let callback_slot = slot.clone();
431                ON_WAKE.set(Some(Box::new(move || {
432                    assert!(!callback_slot.is_set());
433                    assert!(!callback_slot.register(&nested));
434                    assert!(!panic, "wake");
435                })));
436                let result = catch_unwind(AssertUnwindSafe(|| {
437                    if checked {
438                        assert!(slot.wake_checked());
439                    } else {
440                        slot.wake();
441                    }
442                }));
443                assert_eq!(result.is_err(), panic);
444                assert!(slot.wake_checked());
445                assert_eq!(calls.0.load(Ordering::Relaxed), 1);
446            }
447        }
448    }
449
450    /// STATEMENT: Native task wakers support local and cross-thread delivery and remain safe after shutdown.
451    ///
452    /// HOW:
453    /// - Poll a native runtime task, wake it locally and from another thread, and count three polls.
454    /// - Drop the runtime and check task completion, released task state, and one final consumed wake.
455    #[test]
456    fn native_slot_delivery_and_shutdown() {
457        #[derive(Debug)]
458        struct NoopNotify;
459
460        impl ntex_rt::Notify for NoopNotify {
461            fn notify(&self) -> std::io::Result<()> {
462                Ok(())
463            }
464        }
465
466        let rt = ntex_rt::Runtime::new(Box::new(NoopNotify));
467        let slot = Rc::new(LocalWaker::new());
468        let polls = Rc::new(Cell::new(0));
469        let (task_slot, task_polls) = (slot.clone(), polls.clone());
470        let task = rt.spawn(std::future::poll_fn(move |cx| {
471            task_polls.set(task_polls.get() + 1);
472            task_slot.register(cx.waker());
473            std::task::Poll::<()>::Pending
474        }));
475        rt.poll();
476        slot.take_local().unwrap().wake();
477        rt.poll();
478        let standard = slot.take().unwrap();
479        std::thread::spawn(move || standard.wake()).join().unwrap();
480        rt.poll();
481        assert_eq!(polls.get(), 3);
482        assert_eq!(Rc::strong_count(&slot), 2);
483        drop(rt);
484        assert!(task.is_finished());
485        assert_eq!(Rc::strong_count(&slot), 1);
486        assert!(slot.wake_checked());
487        assert!(!slot.wake_checked());
488    }
489}
ntex_util/task.rs

ntex_util/
task.rs
