notify_future/
lib.rs

1use std::future::Future;
2use std::mem::MaybeUninit;
3use std::pin::Pin;
4use std::sync::atomic::Ordering;
5use std::task::{Context, Poll};
6
7use atomic_waker::AtomicWaker;
8
9#[cfg(test)]
10use loom::cell::UnsafeCell;
11#[cfg(test)]
12use loom::sync::{atomic::AtomicU8, Arc};
13#[cfg(not(test))]
14use std::cell::UnsafeCell;
15#[cfg(not(test))]
16use std::sync::{atomic::AtomicU8, Arc};
17
18const PENDING: u8 = 0;
19const WRITING: u8 = 1;
20const READY: u8 = 2;
21const TAKEN: u8 = 3;
22const CANCELED: u8 = 4;
23
24struct NotifyFutureState<RESULT> {
25    state: AtomicU8,
26    waker: AtomicWaker,
27    result: UnsafeCell<MaybeUninit<RESULT>>,
28}
29
30// SAFETY: `state` serializes all writes, reads, and drops of `result`. Values
31// cross thread boundaries only by being written by the notifier and read by the
32// waiter, so `RESULT: Send` is sufficient.
33unsafe impl<RESULT: Send> Send for NotifyFutureState<RESULT> {}
34unsafe impl<RESULT: Send> Sync for NotifyFutureState<RESULT> {}
35
36impl<RESULT> NotifyFutureState<RESULT> {
37    pub fn new() -> Arc<NotifyFutureState<RESULT>> {
38        Arc::new(NotifyFutureState {
39            state: AtomicU8::new(PENDING),
40            waker: AtomicWaker::new(),
41            result: UnsafeCell::new(MaybeUninit::uninit()),
42        })
43    }
44
45    pub fn set_complete(state: &Arc<NotifyFutureState<RESULT>>, result: RESULT) {
46        state.set_complete_inner(result);
47    }
48
49    fn set_complete_inner(&self, result: RESULT) {
50        if self
51            .state
52            .compare_exchange(PENDING, WRITING, Ordering::AcqRel, Ordering::Acquire)
53            .is_err()
54        {
55            return;
56        }
57
58        self.write_result(result);
59
60        self.state.store(READY, Ordering::Release);
61        self.waker.wake();
62    }
63
64    fn poll_result(&self, cx: &mut Context<'_>, already_taken_message: &str) -> Poll<RESULT> {
65        match self.state.load(Ordering::Acquire) {
66            READY => return Poll::Ready(self.take_result(already_taken_message)),
67            TAKEN => panic!("{already_taken_message}"),
68            _ => {}
69        }
70
71        self.waker.register(cx.waker());
72
73        match self.state.load(Ordering::Acquire) {
74            READY => Poll::Ready(self.take_result(already_taken_message)),
75            TAKEN => panic!("{already_taken_message}"),
76            _ => Poll::Pending,
77        }
78    }
79
80    fn take_result(&self, already_taken_message: &str) -> RESULT {
81        match self
82            .state
83            .compare_exchange(READY, TAKEN, Ordering::AcqRel, Ordering::Acquire)
84        {
85            Ok(_) => self.read_result(),
86            Err(TAKEN) => panic!("{already_taken_message}"),
87            Err(_) => panic!("Notify result was not ready"),
88        }
89    }
90
91    fn cancel(&self) {
92        let _ = self
93            .state
94            .compare_exchange(PENDING, CANCELED, Ordering::AcqRel, Ordering::Acquire);
95        let _ = self.waker.take();
96    }
97
98    pub fn is_canceled(&self) -> bool {
99        self.state.load(Ordering::Acquire) == CANCELED
100    }
101}
102
103#[cfg(not(test))]
104impl<RESULT> NotifyFutureState<RESULT> {
105    fn write_result(&self, result: RESULT) {
106        // SAFETY: the PENDING -> WRITING transition gives this notifier unique
107        // write access. Readers only read after observing READY.
108        unsafe {
109            self.result.get().write(MaybeUninit::new(result));
110        }
111    }
112
113    fn read_result(&self) -> RESULT {
114        // SAFETY: READY means the result was initialized, and the successful
115        // READY -> TAKEN transition gives this caller unique ownership.
116        unsafe { (*self.result.get()).as_ptr().read() }
117    }
118
119    fn drop_result(&mut self) {
120        // SAFETY: READY means the result is initialized. `&mut self` means
121        // no other Arc handles remain, so no other thread can take it now.
122        unsafe {
123            self.result.get_mut().assume_init_drop();
124        }
125    }
126}
127
128#[cfg(test)]
129impl<RESULT> NotifyFutureState<RESULT> {
130    fn write_result(&self, result: RESULT) {
131        // SAFETY: the PENDING -> WRITING transition gives this notifier unique
132        // write access. Readers only read after observing READY.
133        self.result.with_mut(|ptr| unsafe {
134            ptr.write(MaybeUninit::new(result));
135        });
136    }
137
138    fn read_result(&self) -> RESULT {
139        // SAFETY: READY means the result was initialized, and the successful
140        // READY -> TAKEN transition gives this caller unique ownership.
141        self.result.with(|ptr| unsafe { (*ptr).as_ptr().read() })
142    }
143
144    fn drop_result(&mut self) {
145        // SAFETY: READY means the result is initialized. `&mut self` means
146        // no other Arc handles remain, so no other thread can take it now.
147        self.result.get_mut().with(|ptr| unsafe {
148            (*ptr).assume_init_drop();
149        });
150    }
151}
152
153impl<RESULT> Drop for NotifyFutureState<RESULT> {
154    fn drop(&mut self) {
155        if self.state.load(Ordering::Acquire) == READY {
156            self.drop_result();
157        }
158    }
159}
160
161#[deprecated(since = "0.2.1", note = "Please use Notify instead")]
162pub struct NotifyFuture<RESULT> {
163    state: Arc<NotifyFutureState<RESULT>>,
164}
165
166#[allow(deprecated)]
167impl<RESULT> Clone for NotifyFuture<RESULT> {
168    fn clone(&self) -> Self {
169        Self {
170            state: self.state.clone(),
171        }
172    }
173}
174
175#[allow(deprecated)]
176impl<RESULT> NotifyFuture<RESULT> {
177    pub fn new() -> Self {
178        Self {
179            state: NotifyFutureState::new(),
180        }
181    }
182
183    pub fn set_complete(&self, result: RESULT) {
184        NotifyFutureState::set_complete(&self.state, result);
185    }
186}
187
188#[allow(deprecated)]
189impl<RESULT> Default for NotifyFuture<RESULT> {
190    fn default() -> Self {
191        Self::new()
192    }
193}
194
195#[allow(deprecated)]
196impl<RESULT> Future for NotifyFuture<RESULT> {
197    type Output = RESULT;
198
199    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
200        self.state.poll_result(
201            cx,
202            "NotifyFuture was awaited by more than one task. Use Notify::new() instead",
203        )
204    }
205}
206
207pub struct Notify<RESULT> {
208    state: Arc<NotifyFutureState<RESULT>>,
209}
210
211impl<RESULT> Notify<RESULT> {
212    pub fn new() -> (Self, NotifyWaiter<RESULT>) {
213        let state = NotifyFutureState::new();
214        (
215            Self {
216                state: state.clone(),
217            },
218            NotifyWaiter::new(state),
219        )
220    }
221
222    pub fn notify(self, result: RESULT) {
223        NotifyFutureState::set_complete(&self.state, result);
224    }
225
226    pub fn is_canceled(&self) -> bool {
227        self.state.is_canceled()
228    }
229}
230
231pub struct NotifyWaiter<RESULT> {
232    state: Arc<NotifyFutureState<RESULT>>,
233}
234
235impl<RESULT> NotifyWaiter<RESULT> {
236    pub(crate) fn new(state: Arc<NotifyFutureState<RESULT>>) -> Self {
237        Self { state }
238    }
239}
240
241impl<RESULT> Drop for NotifyWaiter<RESULT> {
242    fn drop(&mut self) {
243        self.state.cancel();
244    }
245}
246
247impl<RESULT> Future for NotifyWaiter<RESULT> {
248    type Output = RESULT;
249
250    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
251        self.state
252            .poll_result(cx, "NotifyWaiter was polled after completion")
253    }
254}
255
256#[cfg(test)]
257#[allow(deprecated)]
258mod test {
259    use std::future::Future;
260    use std::pin::Pin;
261    use std::sync::{
262        atomic::{AtomicUsize as StdAtomicUsize, Ordering},
263        Arc as StdArc,
264    };
265    use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
266
267    use loom::sync::{atomic::AtomicUsize, Arc};
268    use loom::thread;
269
270    use crate::{Notify, NotifyFuture};
271
272    struct DropCounter {
273        drops: Arc<AtomicUsize>,
274    }
275
276    impl Drop for DropCounter {
277        fn drop(&mut self) {
278            self.drops.fetch_add(1, Ordering::SeqCst);
279        }
280    }
281
282    fn noop_waker() -> &'static Waker {
283        Waker::noop()
284    }
285
286    fn counting_waker(wakes: StdArc<StdAtomicUsize>) -> Waker {
287        unsafe fn clone(data: *const ()) -> RawWaker {
288            let wakes = StdArc::<StdAtomicUsize>::from_raw(data.cast::<StdAtomicUsize>());
289            let cloned = wakes.clone();
290            let _ = StdArc::into_raw(wakes);
291            RawWaker::new(StdArc::into_raw(cloned).cast::<()>(), &VTABLE)
292        }
293
294        unsafe fn wake(data: *const ()) {
295            let wakes = StdArc::<StdAtomicUsize>::from_raw(data.cast::<StdAtomicUsize>());
296            wakes.fetch_add(1, Ordering::SeqCst);
297        }
298
299        unsafe fn wake_by_ref(data: *const ()) {
300            let wakes = StdArc::<StdAtomicUsize>::from_raw(data.cast::<StdAtomicUsize>());
301            wakes.fetch_add(1, Ordering::SeqCst);
302            let _ = StdArc::into_raw(wakes);
303        }
304
305        unsafe fn drop(data: *const ()) {
306            let _ = StdArc::<StdAtomicUsize>::from_raw(data.cast::<StdAtomicUsize>());
307        }
308
309        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
310
311        let raw = RawWaker::new(StdArc::into_raw(wakes).cast::<()>(), &VTABLE);
312        unsafe { Waker::from_raw(raw) }
313    }
314
315    #[test]
316    fn notify_future_ready_after_set_complete() {
317        loom::model(|| {
318            let mut notify_future = NotifyFuture::<u32>::new();
319            let notifier = notify_future.clone();
320
321            notifier.set_complete(1);
322
323            let mut cx = Context::from_waker(noop_waker());
324            let value = match Pin::new(&mut notify_future).poll(&mut cx) {
325                Poll::Ready(value) => value,
326                Poll::Pending => panic!("future remained pending after set_complete"),
327            };
328
329            assert_eq!(value, 1);
330        });
331    }
332
333    #[test]
334    fn notify_waiter_ready_after_notify() {
335        loom::model(|| {
336            let (notify, mut waiter) = Notify::<u32>::new();
337
338            notify.notify(1);
339
340            let mut cx = Context::from_waker(noop_waker());
341            let value = match Pin::new(&mut waiter).poll(&mut cx) {
342                Poll::Ready(value) => value,
343                Poll::Pending => panic!("waiter remained pending after notify"),
344            };
345
346            assert_eq!(value, 1);
347        });
348    }
349
350    #[test]
351    fn notify_waiter_drop_before_ready_is_canceled() {
352        loom::model(|| {
353            let (notify, waiter) = Notify::<u32>::new();
354            drop(waiter);
355            assert!(notify.is_canceled());
356        });
357    }
358
359    #[test]
360    fn notify_waiter_wakes_registered_waker() {
361        loom::model(|| {
362            let (notify, mut waiter) = Notify::<u32>::new();
363            let wakes = StdArc::new(StdAtomicUsize::new(0));
364            let waker = counting_waker(wakes.clone());
365            let mut cx = Context::from_waker(&waker);
366
367            assert!(matches!(Pin::new(&mut waiter).poll(&mut cx), Poll::Pending));
368
369            notify.notify(1);
370
371            assert_eq!(wakes.load(Ordering::SeqCst), 1);
372            let value = match Pin::new(&mut waiter).poll(&mut cx) {
373                Poll::Ready(value) => value,
374                Poll::Pending => panic!("waiter remained pending after notify"),
375            };
376            assert_eq!(value, 1);
377        });
378    }
379
380    #[test]
381    fn notify_after_waiter_drop_discards_result() {
382        loom::model(|| {
383            let drops = Arc::new(AtomicUsize::new(0));
384            let (notify, waiter) = Notify::<DropCounter>::new();
385
386            drop(waiter);
387            assert!(notify.is_canceled());
388
389            notify.notify(DropCounter {
390                drops: drops.clone(),
391            });
392
393            assert_eq!(drops.load(Ordering::SeqCst), 1);
394        });
395    }
396
397    #[test]
398    fn waiter_drop_after_ready_does_not_cancel() {
399        loom::model(|| {
400            let state = super::NotifyFutureState::new();
401            let notify = Notify {
402                state: state.clone(),
403            };
404            let waiter = super::NotifyWaiter::new(state.clone());
405
406            notify.notify(1);
407            drop(waiter);
408
409            assert!(!state.is_canceled());
410        });
411    }
412
413    #[test]
414    fn repeated_set_complete_keeps_first_value() {
415        loom::model(|| {
416            let mut notify_future = NotifyFuture::<u32>::new();
417            let notifier = notify_future.clone();
418
419            notifier.set_complete(1);
420            notifier.set_complete(2);
421
422            let mut cx = Context::from_waker(noop_waker());
423            let value = match Pin::new(&mut notify_future).poll(&mut cx) {
424                Poll::Ready(value) => value,
425                Poll::Pending => panic!("future remained pending after set_complete"),
426            };
427
428            assert_eq!(value, 1);
429        });
430    }
431
432    #[test]
433    fn ready_value_is_dropped_if_waiter_is_dropped() {
434        loom::model(|| {
435            let drops = Arc::new(AtomicUsize::new(0));
436            let (notify, waiter) = Notify::<DropCounter>::new();
437
438            notify.notify(DropCounter {
439                drops: drops.clone(),
440            });
441            drop(waiter);
442
443            assert_eq!(drops.load(Ordering::SeqCst), 1);
444        });
445    }
446
447    #[test]
448    fn taken_value_is_not_dropped_twice() {
449        loom::model(|| {
450            let drops = Arc::new(AtomicUsize::new(0));
451            let (notify, mut waiter) = Notify::<DropCounter>::new();
452
453            notify.notify(DropCounter {
454                drops: drops.clone(),
455            });
456
457            let mut cx = Context::from_waker(noop_waker());
458            let value = match Pin::new(&mut waiter).poll(&mut cx) {
459                Poll::Ready(value) => value,
460                Poll::Pending => panic!("waiter remained pending after notify"),
461            };
462
463            assert_eq!(drops.load(Ordering::SeqCst), 0);
464            drop(value);
465            assert_eq!(drops.load(Ordering::SeqCst), 1);
466        });
467    }
468
469    #[test]
470    fn notify_racing_waiter_drop_drops_value_once() {
471        loom::model(|| {
472            let drops = Arc::new(AtomicUsize::new(0));
473            let (notify, waiter) = Notify::<DropCounter>::new();
474
475            let notify_thread = {
476                let drops = drops.clone();
477                thread::spawn(move || {
478                    notify.notify(DropCounter { drops });
479                })
480            };
481
482            let drop_thread = thread::spawn(move || {
483                drop(waiter);
484            });
485
486            notify_thread.join().unwrap();
487            drop_thread.join().unwrap();
488
489            assert_eq!(drops.load(Ordering::SeqCst), 1);
490        });
491    }
492
493    #[test]
494    fn poll_racing_notify_observes_one_ready_value() {
495        loom::model(|| {
496            let drops = Arc::new(AtomicUsize::new(0));
497            let (notify, mut waiter) = Notify::<DropCounter>::new();
498
499            let notify_thread = {
500                let drops = drops.clone();
501                thread::spawn(move || {
502                    notify.notify(DropCounter { drops });
503                })
504            };
505
506            let mut cx = Context::from_waker(noop_waker());
507
508            match Pin::new(&mut waiter).poll(&mut cx) {
509                Poll::Ready(value) => drop(value),
510                Poll::Pending => {
511                    notify_thread.join().unwrap();
512                    let value = match Pin::new(&mut waiter).poll(&mut cx) {
513                        Poll::Ready(value) => value,
514                        Poll::Pending => panic!("waiter remained pending after notify"),
515                    };
516                    drop(value);
517                    assert_eq!(drops.load(Ordering::SeqCst), 1);
518                    return;
519                }
520            }
521
522            notify_thread.join().unwrap();
523            assert_eq!(drops.load(Ordering::SeqCst), 1);
524        });
525    }
526
527    #[test]
528    fn repeated_notify_future_set_complete_drops_both_inputs_once() {
529        loom::model(|| {
530            let drops = Arc::new(AtomicUsize::new(0));
531            let mut notify_future = NotifyFuture::<DropCounter>::new();
532            let first_notifier = notify_future.clone();
533            let second_notifier = notify_future.clone();
534
535            let first_thread = {
536                let drops = drops.clone();
537                thread::spawn(move || {
538                    first_notifier.set_complete(DropCounter { drops });
539                })
540            };
541
542            let second_thread = {
543                let drops = drops.clone();
544                thread::spawn(move || {
545                    second_notifier.set_complete(DropCounter { drops });
546                })
547            };
548
549            first_thread.join().unwrap();
550            second_thread.join().unwrap();
551
552            let mut cx = Context::from_waker(noop_waker());
553            let value = match Pin::new(&mut notify_future).poll(&mut cx) {
554                Poll::Ready(value) => value,
555                Poll::Pending => panic!("future remained pending after set_complete"),
556            };
557
558            assert_eq!(drops.load(Ordering::SeqCst), 1);
559            drop(value);
560            assert_eq!(drops.load(Ordering::SeqCst), 2);
561        });
562    }
563
564    #[test]
565    #[should_panic(expected = "NotifyWaiter was polled after completion")]
566    fn notify_waiter_panics_when_polled_after_completion() {
567        loom::model(|| {
568            let (notify, mut waiter) = Notify::<u32>::new();
569
570            notify.notify(1);
571
572            let mut cx = Context::from_waker(noop_waker());
573            let value = match Pin::new(&mut waiter).poll(&mut cx) {
574                Poll::Ready(value) => value,
575                Poll::Pending => panic!("waiter remained pending after notify"),
576            };
577            assert_eq!(value, 1);
578
579            let _ = Pin::new(&mut waiter).poll(&mut cx);
580        });
581    }
582
583    #[test]
584    #[should_panic(expected = "NotifyFuture was awaited by more than one task")]
585    fn notify_future_panics_when_polled_after_result_taken() {
586        loom::model(|| {
587            let mut first_waiter = NotifyFuture::<u32>::new();
588            let notifier = first_waiter.clone();
589            let mut second_waiter = first_waiter.clone();
590
591            notifier.set_complete(1);
592
593            let mut cx = Context::from_waker(noop_waker());
594            let value = match Pin::new(&mut first_waiter).poll(&mut cx) {
595                Poll::Ready(value) => value,
596                Poll::Pending => panic!("future remained pending after set_complete"),
597            };
598            assert_eq!(value, 1);
599
600            let _ = Pin::new(&mut second_waiter).poll(&mut cx);
601        });
602    }
603}
notify_future/lib.rs

notify_future/
lib.rs
