1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

use core::ptr;
use core::ptr::NonNull;
use core::sync::atomic::{AtomicPtr, Ordering};
use core::task::Waker;

/// Utility struct to register and wake a waker.
pub struct AtomicWaker {
    waker: AtomicPtr<()>,
}

impl AtomicWaker {
    /// Create a new `AtomicWaker`.
    pub const fn new() -> Self {
        Self {
            waker: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Register a waker. Overwrites the previous waker, if any.
    pub fn register(&self, w: &Waker) {
        self.waker.store(w.as_turbo_ptr().as_ptr() as _, Ordering::Release);
    }

    /// Wake the registered waker, if any.
    pub fn wake(&self) {
        if let Some(ptr) = NonNull::new(self.waker.load(Ordering::Acquire)) {
            unsafe { Waker::from_turbo_ptr(ptr) }.wake();
        }
    }
}