async_pool/
lib.rs

1//! Statically allocated pool providing a std-like Box,
2//! allowing to asynchronously await for a pool slot to become available.
3//!
4//! It is tailored to be used with no-std async runtimes, like [Embassy](https://embassy.dev/), but
5//! can also be used in std environments (check examples).
6//!
7//! The most common use-case is sharing large memory regions on constrained
8//! devices (e.g. microcontrollers), where multiple tasks may need to use the
9//! memory for buffering an I/O or performing calculations, and having
10//! separate static buffers would be too costly.
11//!
12//! It is important to know that waiting forever for a memory slot to be
13//! available may dead-lock your code if done wrong. With that in mind,
14//! you should consider using a timeout when allocating asynchronously (e.g. [embassy_time::with_timeout](https://docs.rs/embassy-time/0.3.2/embassy_time/fn.with_timeout.html)).
15//!
16//! #### Dependencies
17//!
18//! This crate requires a critical section implementation. Check [critical-section](https://crates.io/crates/critical-section).
19//!
20//! #### Example
21//!
22//! ```
23//! use async_pool::{pool, Box};
24//!
25//!struct Buffer([u8; 256]);
26//!
27//!// A maximum of 2 Packet instances can be allocated at a time.
28//!// A maximum of 3 futures can be waiting at a time.
29//!pool!(BufferPool: [Buffer; 2], 3);
30//!
31//!async fn run() {
32//!    // Allocate non-blocking (will return None if no data slot is available)
33//!    let box1 = Box::<BufferPool>::new(Buffer([0; 256]));
34//!
35//!    // Allocate asynchronously (will wait if no data slot is available)
36//!    // This can return None if all future slots are taken
37//!    let box2 = Box::<BufferPool>::new_async(Buffer([0; 256])).await;
38//!}
39//! ```
40#![cfg_attr(not(test), no_std)]
41
42mod atomic_bitset;
43
44use core::cell::UnsafeCell;
45use core::future::{poll_fn, Future};
46use core::hash::{Hash, Hasher};
47use core::mem::MaybeUninit;
48use core::ops::{Deref, DerefMut};
49use core::task::Poll;
50use core::{cmp, mem, ptr::NonNull};
51use embassy_sync::waitqueue::AtomicWaker;
52use portable_atomic::AtomicU32;
53
54use crate::atomic_bitset::AtomicBitset;
55
56/// Implementation detail. Not covered by semver guarantees.
57#[doc(hidden)]
58pub trait PoolStorage<T> {
59    fn alloc(&self) -> Option<NonNull<T>>;
60    fn alloc_async(&self) -> impl Future<Output = Option<NonNull<T>>>;
61    unsafe fn free(&self, p: NonNull<T>);
62}
63
64/// Implementation detail. Not covered by semver guarantees.
65#[doc(hidden)]
66pub struct PoolStorageImpl<T, const N: usize, const K: usize, const WN: usize, const WK: usize>
67where
68    [AtomicU32; K]: Sized,
69    [AtomicU32; WK]: Sized,
70{
71    used: AtomicBitset<N, K>,
72    data: [UnsafeCell<MaybeUninit<T>>; N],
73    wakers_used: AtomicBitset<WN, WK>,
74    wakers: [AtomicWaker; WN],
75}
76
77unsafe impl<T, const N: usize, const K: usize, const WN: usize, const WK: usize> Send
78    for PoolStorageImpl<T, N, K, WN, WK>
79{
80}
81unsafe impl<T, const N: usize, const K: usize, const WN: usize, const WK: usize> Sync
82    for PoolStorageImpl<T, N, K, WN, WK>
83{
84}
85
86impl<T, const N: usize, const K: usize, const WN: usize, const WK: usize>
87    PoolStorageImpl<T, N, K, WN, WK>
88where
89    [AtomicU32; K]: Sized,
90    [AtomicU32; WK]: Sized,
91{
92    const UNINIT: UnsafeCell<MaybeUninit<T>> = UnsafeCell::new(MaybeUninit::uninit());
93
94    const WAKER: AtomicWaker = AtomicWaker::new();
95
96    pub const fn new() -> Self {
97        Self {
98            used: AtomicBitset::new(),
99            data: [Self::UNINIT; N],
100            wakers_used: AtomicBitset::new(),
101            wakers: [Self::WAKER; WN],
102        }
103    }
104}
105
106impl<T, const N: usize, const K: usize, const WN: usize, const WK: usize> PoolStorage<T>
107    for PoolStorageImpl<T, N, K, WN, WK>
108where
109    [AtomicU32; K]: Sized,
110    [AtomicU32; WK]: Sized,
111{
112    /// Returns an item from the data pool, if available.
113    /// Returns None if the data pool is full.
114    fn alloc(&self) -> Option<NonNull<T>> {
115        let n = self.used.alloc()?;
116        let p = self.data[n].get() as *mut T;
117        Some(unsafe { NonNull::new_unchecked(p) })
118    }
119
120    /// Wait until an item is available in the data pool, then return it.
121    /// Returns None if the waker pool is full.
122    fn alloc_async(&self) -> impl Future<Output = Option<NonNull<T>>> {
123        let mut waker_slot = None;
124        poll_fn(move |cx| {
125            // Check if there is a free slot in the data pool
126            if let Some(n) = self.used.alloc() {
127                let p = self.data[n].get() as *mut T;
128                return Poll::Ready(Some(unsafe { NonNull::new_unchecked(p) }));
129            }
130
131            // Try to allocate a waker slot if necessary
132            if waker_slot.is_none() {
133                waker_slot = self.wakers_used.alloc_droppable();
134            }
135
136            match &waker_slot {
137                Some(bit) => {
138                    self.wakers[bit.inner()].register(cx.waker());
139                    Poll::Pending
140                }
141                None => Poll::Ready(None), // No waker slots available
142            }
143        })
144    }
145
146    /// safety: p must be a pointer obtained from self.alloc that hasn't been freed yet.
147    unsafe fn free(&self, p: NonNull<T>) {
148        let origin = self.data.as_ptr() as *mut T;
149        let n = p.as_ptr().offset_from(origin);
150        assert!(n >= 0);
151        assert!((n as usize) < N);
152        self.used.free(n as usize);
153
154        // Wake up any wakers waiting for a slot
155        for waker in self.wakers.iter() {
156            waker.wake();
157        }
158    }
159}
160
161pub trait Pool: 'static {
162    type Item: 'static;
163
164    /// Implementation detail. Not covered by semver guarantees.
165    #[doc(hidden)]
166    type Storage: PoolStorage<Self::Item>;
167
168    /// Implementation detail. Not covered by semver guarantees.
169    #[doc(hidden)]
170    fn get() -> &'static Self::Storage;
171}
172
173pub struct Box<P: Pool> {
174    ptr: NonNull<P::Item>,
175}
176
177impl<P: Pool> Box<P> {
178    /// Returns an item from the data pool, if available.
179    /// Returns None if the data pool is full.
180    pub fn new(item: P::Item) -> Option<Self> {
181        let p = match P::get().alloc() {
182            Some(p) => p,
183            None => return None,
184        };
185        unsafe { p.as_ptr().write(item) };
186        Some(Self { ptr: p })
187    }
188
189    /// Wait until an item is available in the data pool, then return it.
190    /// Returns None if the waker pool is full.
191    pub async fn new_async(item: P::Item) -> Option<Self> {
192        let p = match P::get().alloc_async().await {
193            Some(p) => p,
194            None => return None,
195        };
196        unsafe { p.as_ptr().write(item) };
197        Some(Self { ptr: p })
198    }
199
200    pub fn into_raw(b: Self) -> NonNull<P::Item> {
201        let res = b.ptr;
202        mem::forget(b);
203        res
204    }
205
206    pub unsafe fn from_raw(ptr: NonNull<P::Item>) -> Self {
207        Self { ptr }
208    }
209}
210
211impl<P: Pool> Drop for Box<P> {
212    fn drop(&mut self) {
213        unsafe {
214            //trace!("dropping {:u32}", self.ptr as u32);
215            self.ptr.as_ptr().drop_in_place();
216            P::get().free(self.ptr);
217        };
218    }
219}
220
221unsafe impl<P: Pool> Send for Box<P> where P::Item: Send {}
222
223unsafe impl<P: Pool> Sync for Box<P> where P::Item: Sync {}
224
225unsafe impl<P: Pool> stable_deref_trait::StableDeref for Box<P> {}
226
227impl<P: Pool> as_slice_01::AsSlice for Box<P>
228where
229    P::Item: as_slice_01::AsSlice,
230{
231    type Element = <P::Item as as_slice_01::AsSlice>::Element;
232
233    fn as_slice(&self) -> &[Self::Element] {
234        self.deref().as_slice()
235    }
236}
237
238impl<P: Pool> as_slice_01::AsMutSlice for Box<P>
239where
240    P::Item: as_slice_01::AsMutSlice,
241{
242    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
243        self.deref_mut().as_mut_slice()
244    }
245}
246
247impl<P: Pool> as_slice_02::AsSlice for Box<P>
248where
249    P::Item: as_slice_02::AsSlice,
250{
251    type Element = <P::Item as as_slice_02::AsSlice>::Element;
252
253    fn as_slice(&self) -> &[Self::Element] {
254        self.deref().as_slice()
255    }
256}
257
258impl<P: Pool> as_slice_02::AsMutSlice for Box<P>
259where
260    P::Item: as_slice_02::AsMutSlice,
261{
262    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
263        self.deref_mut().as_mut_slice()
264    }
265}
266
267impl<P: Pool> Deref for Box<P> {
268    type Target = P::Item;
269
270    fn deref(&self) -> &P::Item {
271        unsafe { self.ptr.as_ref() }
272    }
273}
274
275impl<P: Pool> DerefMut for Box<P> {
276    fn deref_mut(&mut self) -> &mut P::Item {
277        unsafe { self.ptr.as_mut() }
278    }
279}
280
281impl<P: Pool> core::fmt::Debug for Box<P>
282where
283    P::Item: core::fmt::Debug,
284{
285    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
286        <P::Item as core::fmt::Debug>::fmt(self, f)
287    }
288}
289
290impl<P: Pool> core::fmt::Display for Box<P>
291where
292    P::Item: core::fmt::Display,
293{
294    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
295        <P::Item as core::fmt::Display>::fmt(self, f)
296    }
297}
298
299impl<P: Pool> PartialEq for Box<P>
300where
301    P::Item: PartialEq,
302{
303    fn eq(&self, rhs: &Box<P>) -> bool {
304        <P::Item as PartialEq>::eq(self, rhs)
305    }
306}
307
308impl<P: Pool> Eq for Box<P> where P::Item: Eq {}
309
310impl<P: Pool> PartialOrd for Box<P>
311where
312    P::Item: PartialOrd,
313{
314    fn partial_cmp(&self, rhs: &Box<P>) -> Option<cmp::Ordering> {
315        <P::Item as PartialOrd>::partial_cmp(self, rhs)
316    }
317}
318
319impl<P: Pool> Ord for Box<P>
320where
321    P::Item: Ord,
322{
323    fn cmp(&self, rhs: &Box<P>) -> cmp::Ordering {
324        <P::Item as Ord>::cmp(self, rhs)
325    }
326}
327
328impl<P: Pool> Hash for Box<P>
329where
330    P::Item: Hash,
331{
332    fn hash<H>(&self, state: &mut H)
333    where
334        H: Hasher,
335    {
336        <P::Item as Hash>::hash(self, state)
337    }
338}
339
340/// Create a item pool of a given type and size, as well as a waker pool of a given length.
341///
342/// The waker pool is used to wake up tasks waiting for an item to become available in the data pool.
343/// Its length should be at least the number of tasks that can be waiting for an item at the same time.
344/// Example:
345/// ```
346/// use async_pool::{pool, Box};
347///
348/// #[derive(Debug)]
349/// #[allow(dead_code)]
350/// struct Packet(u32);
351///
352/// pool!(PacketPool: [Packet; 4], 2); // Item pool of 4 Packet instances, waker pool of 2 wakers
353/// ```
354#[macro_export]
355macro_rules! pool {
356    ($vis:vis $name:ident: [$ty:ty; $n:expr], $wn:expr) => {
357        $vis struct $name { _uninhabited: ::core::convert::Infallible }
358        impl $crate::Pool for $name {
359            type Item = $ty;
360            type Storage = $crate::PoolStorageImpl<$ty, {$n}, {($n+31)/32}, {$wn}, {($wn+31)/32}>;
361            fn get() -> &'static Self::Storage {
362                static POOL: $crate::PoolStorageImpl<$ty, {$n}, {($n+31)/32}, {$wn}, {($wn+31)/32}> = $crate::PoolStorageImpl::new();
363                &POOL
364            }
365        }
366    };
367}
368
369#[cfg(test)]
370mod test {
371    use super::*;
372    use core::mem;
373
374    pool!(TestPool: [u32; 4], 0);
375    pool!(TestPool2: [u32; 4], 2);
376
377    #[test]
378    fn test_pool() {
379        let b1 = Box::<TestPool>::new(111).unwrap();
380        let b2 = Box::<TestPool>::new(222).unwrap();
381        let b3 = Box::<TestPool>::new(333).unwrap();
382        let b4 = Box::<TestPool>::new(444).unwrap();
383        assert!(Box::<TestPool>::new(555).is_none());
384        assert_eq!(*b1, 111);
385        assert_eq!(*b2, 222);
386        assert_eq!(*b3, 333);
387        assert_eq!(*b4, 444);
388        mem::drop(b3);
389        let b5 = Box::<TestPool>::new(555).unwrap();
390        assert!(Box::<TestPool>::new(666).is_none());
391        assert_eq!(*b1, 111);
392        assert_eq!(*b2, 222);
393        assert_eq!(*b4, 444);
394        assert_eq!(*b5, 555);
395    }
396
397    #[test]
398    fn test_async_sizes() {
399        let pool1 = <TestPool as Pool>::get();
400        let pool2 = <TestPool2 as Pool>::get();
401        assert!(mem::size_of_val(pool1) < mem::size_of_val(pool2));
402    }
403}
async_pool/lib.rs

async_pool/
lib.rs
