cotton_usb_host/
async_pool.rs

1use crate::bitset::BitSet;
2use core::cell::{Cell, RefCell};
3use core::future::Future;
4use core::pin::Pin;
5use core::task::{Context, Poll, Waker};
6#[cfg(feature = "std")]
7use std::fmt::{self, Display};
8
9/// Managing access to N equivalent resources
10///
11/// Callers who wish to access a resource (but don't care which one of the N)
12/// can call the async function [`Pool::alloc`] which will return (awakening
13/// the task) as soon as a resource is available.
14///
15/// Once the returned resource (represented as a [`Pooled`]) is
16/// finished with, the caller can then drop it (i.e., let it go out of
17/// scope) and it will be returned to the pool for another user.
18///
19/// Not quite the same as async_semaphore, because callers need to know
20/// *which* of the N devices they were allotted.
21///
22/// # Example
23/// ```rust
24/// use cotton_usb_host::async_pool::Pool;
25/// let mut pool = Pool::new(2); // this pool has two resources
26/// let res = pool.try_alloc().unwrap(); // obtain a resource
27/// println!("I got resource {}", res.which());
28/// {
29///     let res2 = pool.try_alloc().unwrap(); // obtain a resource
30///     println!("I got resource {}", res.which());
31///     let res3 = pool.try_alloc();
32///     assert!(res3.is_none());   // oh dear, no resources available
33///     // But now res2 goes out of scope (i.e., back into the pool)
34/// }
35/// let res4 = pool.try_alloc().unwrap();
36/// println!("I got resource {}", res.which()); // success!
37/// ```
38///
39/// For a larger example, see how the RP2040 USB host-controller driver
40/// shares out its USB endpoints.
41pub struct Pool {
42    total: u8,
43    allocated: Cell<BitSet>,
44    waker: RefCell<Option<Waker>>,
45}
46
47/// Representing ownership of one of the resources in a [`Pool`]
48pub struct Pooled<'a> {
49    n: u8,
50    pool: &'a Pool,
51}
52
53impl Pooled<'_> {
54    /// Returns which one of the [`Pool`]'s N resources is owned by this `Pooled`
55    pub fn which(&self) -> u8 {
56        self.n
57    }
58}
59
60#[cfg(feature = "std")]
61impl Display for Pooled<'_> {
62    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
63        write!(f, "Pooled({})", self.n)
64    }
65}
66
67#[cfg(feature = "defmt")]
68impl defmt::Format for Pooled<'_> {
69    fn format(&self, f: defmt::Formatter) {
70        defmt::write!(f, "Pooled({})", self.n);
71    }
72}
73
74impl Drop for Pooled<'_> {
75    fn drop(&mut self) {
76        self.pool.dealloc_internal(self.n);
77    }
78}
79
80struct PoolFuture<'a> {
81    pool: &'a Pool,
82}
83
84impl<'a> Future for PoolFuture<'a> {
85    type Output = Pooled<'a>;
86
87    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
88        self.pool.waker.replace(Some(cx.waker().clone()));
89
90        if let Some(n) = self.pool.alloc_internal() {
91            Poll::Ready(Pooled { n, pool: self.pool })
92        } else {
93            Poll::Pending
94        }
95    }
96}
97
98impl Pool {
99    /// Create a new Pool, sharing out a number of equivalent resources
100    ///
101    /// # Parameters
102    /// - `total`: The number of resources (must be 0-32)
103    ///
104    /// # Panics
105    /// Will panic if `total`>32.
106    pub const fn new(total: u8) -> Self {
107        assert!(total <= 32);
108        Self {
109            total,
110            allocated: Cell::new(BitSet::new()),
111            waker: RefCell::new(None),
112        }
113    }
114
115    fn alloc_internal(&self) -> Option<u8> {
116        let mut bits = self.allocated.get();
117        let n = bits.set_any()?;
118        if n >= self.total {
119            None
120        } else {
121            self.allocated.replace(bits);
122            Some(n)
123        }
124    }
125
126    fn dealloc_internal(&self, n: u8) {
127        let mut bits = self.allocated.get();
128        debug_assert!(bits.contains(n));
129        bits.clear(n);
130        self.allocated.replace(bits);
131
132        if let Some(w) = self.waker.take() {
133            w.wake();
134        }
135    }
136
137    /// Obtain one of the resources
138    ///
139    /// This asynchronous function will return immediately if any of
140    /// the resources is currently idle (unused). Otherwise, it will
141    /// wait until a resource is available. Once it has returned, the
142    /// caller has ownership of the resource (represented by ownership
143    /// of the [`Pooled`] object) until the `Pooled` is dropped -- for instance,
144    /// at the end of a scope.
145    ///
146    /// It is not unsafe or unsound (in the Rust sense) to keep hold
147    /// of a `Pooled` indefinitely, nor to `mem::forget` it -- but it is
148    /// inadvisable, as this constitutes a denial-of-service against
149    /// other potential resource users.
150    ///
151    /// # See also
152    /// [`Pool::try_alloc()`] for a synchronous version
153    pub async fn alloc(&self) -> Pooled<'_> {
154        let fut = PoolFuture { pool: self };
155        fut.await
156    }
157
158    /// Obtain a resource if one is immediately available
159    ///
160    /// Returns `Some` if any of the resources is currently idle (unused).
161    /// Otherwise, returns `None`.
162    ///
163    /// # See also
164    /// [`Pool::alloc()`] for an asynchronous version
165    pub fn try_alloc(&self) -> Option<Pooled<'_>> {
166        Some(Pooled {
167            n: self.alloc_internal()?,
168            pool: self,
169        })
170    }
171}
172
173#[cfg(all(test, feature = "std"))]
174#[path = "tests/async_pool.rs"]
175mod tests;
cotton_usb_host/async_pool.rs

cotton_usb_host/
async_pool.rs
