async_ach_waker/
pool.rs

1use crate::WakerEntity;
2use ach_array::Array;
3use core::sync::atomic::{AtomicBool, Ordering};
4
5pub struct WakerToken<'a, T, const N: usize> {
6    pool: &'a WakerPool<T, N>,
7    index: usize,
8}
9impl<'a, T, const N: usize> WakerToken<'a, T, N> {
10    /// Swap waker
11    ///
12    /// wake it immediately, if is waking.
13    pub fn swap(&self, waker: WakerEntity<T>) {
14        if let Err(e) = self.pool.pool[self.index].try_replace(waker) {
15            e.input.wake();
16        }
17    }
18}
19impl<'a, T, const N: usize> Drop for WakerToken<'a, T, N> {
20    fn drop(&mut self) {
21        if let Ok(wake) = self.pool.pool[self.index].get() {
22            wake.remove();
23        }
24        let _ = self.pool.used[self.index].compare_exchange(
25            true,
26            false,
27            Ordering::SeqCst,
28            Ordering::Relaxed,
29        );
30    }
31}
32
33pub struct WakerPool<T, const N: usize> {
34    pool: Array<WakerEntity<T>, N>,
35    used: [AtomicBool; N],
36}
37impl<T, const N: usize> WakerPool<T, N> {
38    const FALSE: AtomicBool = AtomicBool::new(false);
39    pub const fn new() -> Self {
40        Self {
41            pool: Array::new(),
42            used: [Self::FALSE; N],
43        }
44    }
45    /// Hold a place in the pool
46    pub fn register(&self) -> Result<WakerToken<T, N>, ()> {
47        for (i, used) in self.used.iter().enumerate() {
48            if used
49                .compare_exchange(false, true, Ordering::SeqCst, Ordering::Relaxed)
50                .is_ok()
51            {
52                return Ok(WakerToken {
53                    pool: self,
54                    index: i,
55                });
56            }
57        }
58        Err(())
59    }
60    /// Wake a waiter, and remove it.
61    ///
62    /// Returns false if the pool is empty.
63    pub fn wake_one(&self) -> bool {
64        loop {
65            if let Some(waker) = self.pool.pop() {
66                if waker.wake() {
67                    return true;
68                }
69            } else {
70                return false;
71            }
72        }
73    }
74    /// Wake all waiter, and remove it.
75    ///
76    /// returns the number of had waked
77    pub fn wake_all(&self) -> usize {
78        let mut num = 0;
79        while let Some(waker) = self.pool.pop() {
80            waker.wake();
81            num += 1;
82        }
83        num
84    }
85
86    /// Retains only the elements specified by the predicate.
87    ///
88    /// In other words, remove all elements e such that f(&e) returns false.
89    /// This method operates in place, visiting each element exactly once in the original order,
90    /// but not preserves the order of the retained elements.
91    pub fn retain(&self, mut f: impl FnMut(&WakerEntity<T>) -> bool) {
92        for e in self.pool.iter(false) {
93            if !f(&*e) {
94                // Remove it
95                e.remove();
96            }
97        }
98    }
99}
async_ach_waker/pool.rs

async_ach_waker/
pool.rs
