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    pub const fn new() -> Self {
39        Self {
40            pool: Array::new(),
41            used: [const { AtomicBool::new(false) }; N],
42        }
43    }
44    /// Hold a place in the pool
45    pub fn register(&'_ self) -> Result<WakerToken<'_, T, N>, ()> {
46        for (i, used) in self.used.iter().enumerate() {
47            if used
48                .compare_exchange(false, true, Ordering::SeqCst, Ordering::Relaxed)
49                .is_ok()
50            {
51                return Ok(WakerToken {
52                    pool: self,
53                    index: i,
54                });
55            }
56        }
57        Err(())
58    }
59    /// Wake a waiter, and remove it.
60    ///
61    /// Returns false if the pool is empty.
62    pub fn wake_one(&self) -> bool {
63        loop {
64            if let Some(waker) = self.pool.pop() {
65                if waker.wake() {
66                    return true;
67                }
68            } else {
69                return false;
70            }
71        }
72    }
73    /// Wake all waiter, and remove it.
74    ///
75    /// returns the number of had waked
76    pub fn wake_all(&self) -> usize {
77        let mut num = 0;
78        while let Some(waker) = self.pool.pop() {
79            waker.wake();
80            num += 1;
81        }
82        num
83    }
84
85    /// Retains only the elements specified by the predicate.
86    ///
87    /// In other words, remove all elements e such that f(&e) returns false.
88    /// This method operates in place, visiting each element exactly once in the original order,
89    /// but not preserves the order of the retained elements.
90    pub fn retain(&self, mut f: impl FnMut(&WakerEntity<T>) -> bool) {
91        for e in self.pool.iter(false) {
92            if !f(&*e) {
93                // Remove it
94                e.remove();
95            }
96        }
97    }
98}
async_ach_waker/pool.rs

async_ach_waker/
pool.rs
