Skip to main content

prns_interfaces_embassy/bluetooth_auto/
connection_slots.rs

1use core::cell::Cell;
2
3use embassy_sync::blocking_mutex::raw::RawMutex;
4use embassy_sync::blocking_mutex::Mutex as BlockingMutex;
5use embassy_sync::semaphore::{FairSemaphore, Semaphore, SemaphoreReleaser};
6use embassy_sync::signal::Signal;
7use portable_atomic::{AtomicU8, Ordering};
8use prns_core::interfaces::bluetooth_auto::Origin;
9
10type Availability<M> = FairSemaphore<M, 1>;
11
12pub struct ConnectionSlotPool<M: RawMutex + 'static, const SLOTS: usize> {
13    availability: Availability<M>,
14    slots: [ConnectionSlotState<M>; SLOTS],
15}
16
17struct ConnectionSlotState<M: RawMutex + 'static> {
18    owners: AtomicU8,
19    index: AtomicU8,
20    availability: BlockingMutex<M, Cell<Option<&'static Availability<M>>>>,
21    closed: Signal<M, ()>,
22}
23
24impl<M: RawMutex + 'static> ConnectionSlotState<M> {
25    const fn new() -> Self {
26        Self {
27            owners: AtomicU8::new(0),
28            index: AtomicU8::new(0),
29            availability: BlockingMutex::new(Cell::new(None)),
30            closed: Signal::new(),
31        }
32    }
33
34    fn add_owner(&self) {
35        self.owners.fetch_add(1, Ordering::AcqRel);
36    }
37
38    fn request_close(&self) {
39        self.closed.signal(());
40    }
41
42    fn release(&self) {
43        self.request_close();
44        if self.owners.fetch_sub(1, Ordering::AcqRel) != 1 {
45            return;
46        }
47        if let Some(availability) = self.availability.lock(Cell::get) {
48            availability.release(1);
49        }
50    }
51}
52
53impl<M: RawMutex + 'static, const SLOTS: usize> ConnectionSlotPool<M, SLOTS> {
54    #[must_use]
55    pub const fn new() -> Self {
56        assert!(SLOTS <= u8::MAX as usize + 1);
57        Self {
58            availability: FairSemaphore::new(SLOTS),
59            slots: [const { ConnectionSlotState::new() }; SLOTS],
60        }
61    }
62
63    pub async fn acquire(
64        &'static self,
65    ) -> Result<ConnectionSlotLease<M>, ConnectionSlotAcquireError> {
66        let permit = self
67            .availability
68            .acquire(1)
69            .await
70            .map_err(|_| ConnectionSlotAcquireError::WaitQueueFull)?;
71        self.claim(permit)
72            .ok_or(ConnectionSlotAcquireError::PermitWithoutAvailableSlot)
73    }
74
75    pub fn try_acquire(
76        &'static self,
77    ) -> Result<Option<ConnectionSlotLease<M>>, ConnectionSlotAcquireError> {
78        let Some(permit) = self.availability.try_acquire(1) else {
79            return Ok(None);
80        };
81        self.claim(permit)
82            .map(Some)
83            .ok_or(ConnectionSlotAcquireError::PermitWithoutAvailableSlot)
84    }
85
86    pub fn request_close(&self, index: usize) {
87        if let Some(slot) = self.slots.get(index) {
88            slot.request_close();
89        }
90    }
91
92    fn claim(
93        &'static self,
94        permit: SemaphoreReleaser<'static, Availability<M>>,
95    ) -> Option<ConnectionSlotLease<M>> {
96        for index in 0..SLOTS {
97            let slot = &self.slots[index];
98            if slot
99                .owners
100                .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire)
101                .is_ok()
102            {
103                slot.index.store(index as u8, Ordering::Release);
104                slot.availability
105                    .lock(|availability| availability.set(Some(&self.availability)));
106                slot.closed.reset();
107                permit.disarm();
108                return Some(ConnectionSlotLease {
109                    owner: ConnectionSlotOwner { slot },
110                });
111            }
112        }
113        None
114    }
115}
116
117#[derive(Debug, PartialEq, Eq)]
118pub enum ConnectionSlotAcquireError {
119    WaitQueueFull,
120    PermitWithoutAvailableSlot,
121}
122
123impl<M: RawMutex + 'static, const SLOTS: usize> Default for ConnectionSlotPool<M, SLOTS> {
124    fn default() -> Self {
125        Self::new()
126    }
127}
128
129#[must_use]
130pub struct ConnectionSlotLease<M: RawMutex + 'static> {
131    owner: ConnectionSlotOwner<M>,
132}
133
134struct ConnectionSlotOwner<M: RawMutex + 'static> {
135    slot: &'static ConnectionSlotState<M>,
136}
137
138impl<M: RawMutex + 'static> ConnectionSlotOwner<M> {
139    fn index(&self) -> usize {
140        usize::from(self.slot.index.load(Ordering::Acquire))
141    }
142
143    fn wait_for_close(&self) -> impl core::future::Future<Output = ()> + '_ {
144        self.slot.closed.wait()
145    }
146
147    fn split(self) -> (Self, Self) {
148        self.slot.add_owner();
149        let slot = self.slot;
150        core::mem::forget(self);
151        (Self { slot }, Self { slot })
152    }
153}
154
155impl<M: RawMutex + 'static> Drop for ConnectionSlotOwner<M> {
156    fn drop(&mut self) {
157        self.slot.release();
158    }
159}
160
161impl<M: RawMutex + 'static> ConnectionSlotLease<M> {
162    #[must_use]
163    pub fn index(&self) -> usize {
164        self.owner.index()
165    }
166
167    pub fn activate(self) -> ConnectionSlotOwners<M> {
168        let (worker, link) = self.owner.split();
169        ConnectionSlotOwners {
170            worker: ConnectionSlotWorkerLease { owner: worker },
171            link: ConnectionSlotLinkLease { owner: link },
172        }
173    }
174}
175
176#[must_use]
177pub struct ConnectionSlotOwners<M: RawMutex + 'static> {
178    pub worker: ConnectionSlotWorkerLease<M>,
179    pub link: ConnectionSlotLinkLease<M>,
180}
181
182#[must_use]
183pub struct ConnectionSlotWorkerLease<M: RawMutex + 'static> {
184    owner: ConnectionSlotOwner<M>,
185}
186
187impl<M: RawMutex + 'static> ConnectionSlotWorkerLease<M> {
188    pub fn wait_for_close(&self) -> impl core::future::Future<Output = ()> + '_ {
189        self.owner.wait_for_close()
190    }
191
192    pub fn request_close(&self) {
193        self.owner.slot.request_close();
194    }
195}
196
197#[must_use]
198pub struct ConnectionSlotLinkLease<M: RawMutex + 'static> {
199    owner: ConnectionSlotOwner<M>,
200}
201
202impl<M: RawMutex + 'static> ConnectionSlotLinkLease<M> {
203    #[must_use]
204    pub fn index(&self) -> usize {
205        self.owner.index()
206    }
207
208    pub fn wait_for_close(&self) -> impl core::future::Future<Output = ()> + '_ {
209        self.owner.wait_for_close()
210    }
211
212    pub fn into_ready(self, origin: Origin) -> ReadyConnectionSlot<M> {
213        ReadyConnectionSlot { link: self, origin }
214    }
215
216    pub fn into_data(self) -> ConnectionSlotDataOwners<M> {
217        let (source, sink) = self.owner.split();
218        ConnectionSlotDataOwners {
219            source: ConnectionSlotSourceLease { owner: source },
220            sink: ConnectionSlotSinkLease { owner: sink },
221        }
222    }
223}
224
225#[must_use]
226pub struct ConnectionSlotDataOwners<M: RawMutex + 'static> {
227    pub source: ConnectionSlotSourceLease<M>,
228    pub sink: ConnectionSlotSinkLease<M>,
229}
230
231#[must_use]
232pub struct ConnectionSlotSourceLease<M: RawMutex + 'static> {
233    owner: ConnectionSlotOwner<M>,
234}
235
236impl<M: RawMutex + 'static> ConnectionSlotSourceLease<M> {
237    pub fn wait_for_close(&self) -> impl core::future::Future<Output = ()> + '_ {
238        self.owner.wait_for_close()
239    }
240}
241
242#[must_use]
243pub struct ConnectionSlotSinkLease<M: RawMutex + 'static> {
244    owner: ConnectionSlotOwner<M>,
245}
246
247impl<M: RawMutex + 'static> ConnectionSlotSinkLease<M> {
248    pub fn wait_for_close(&self) -> impl core::future::Future<Output = ()> + '_ {
249        self.owner.wait_for_close()
250    }
251}
252
253#[must_use]
254pub struct ReadyConnectionSlot<M: RawMutex + 'static> {
255    link: ConnectionSlotLinkLease<M>,
256    origin: Origin,
257}
258
259impl<M: RawMutex + 'static> ReadyConnectionSlot<M> {
260    pub fn into_parts(self) -> ReadyConnectionSlotParts<M> {
261        ReadyConnectionSlotParts {
262            origin: self.origin,
263            link: self.link,
264        }
265    }
266}
267
268#[must_use]
269pub struct ReadyConnectionSlotParts<M: RawMutex + 'static> {
270    pub origin: Origin,
271    pub link: ConnectionSlotLinkLease<M>,
272}
273
274#[cfg(test)]
275mod tests {
276    use core::future::ready;
277
278    use embassy_futures::block_on;
279    use embassy_futures::select::{select, Either};
280    use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
281    use prns_core::interfaces::bluetooth_auto::Origin;
282
283    use super::{
284        ConnectionSlotAcquireError, ConnectionSlotDataOwners, ConnectionSlotOwners,
285        ConnectionSlotPool, ReadyConnectionSlotParts,
286    };
287
288    #[test]
289    fn leases_reserve_unique_slots_until_drop() {
290        static POOL: ConnectionSlotPool<CriticalSectionRawMutex, 3> = ConnectionSlotPool::new();
291
292        let first = POOL.try_acquire().ok().flatten();
293        let second = POOL.try_acquire().ok().flatten();
294        let third = POOL.try_acquire().ok().flatten();
295
296        assert_eq!(first.as_ref().map(|lease| lease.index()), Some(0));
297        assert_eq!(second.as_ref().map(|lease| lease.index()), Some(1));
298        assert_eq!(third.as_ref().map(|lease| lease.index()), Some(2));
299        assert_eq!(POOL.try_acquire().map(|lease| lease.is_none()), Ok(true));
300
301        drop(second);
302        let replacement = POOL.try_acquire().ok().flatten();
303        assert_eq!(replacement.as_ref().map(|lease| lease.index()), Some(1));
304    }
305
306    #[test]
307    fn split_lease_releases_after_both_owners_drop() {
308        static POOL: ConnectionSlotPool<CriticalSectionRawMutex, 1> = ConnectionSlotPool::new();
309
310        let lease = POOL.try_acquire().ok().flatten();
311        assert!(lease.is_some());
312        if let Some(lease) = lease {
313            let ConnectionSlotOwners { worker, link } = lease.activate();
314            let ready = link.into_ready(Origin::Dialed);
315            let ReadyConnectionSlotParts { origin, link } = ready.into_parts();
316            assert_eq!(origin, Origin::Dialed);
317
318            drop(worker);
319            assert_eq!(POOL.try_acquire().map(|lease| lease.is_none()), Ok(true));
320            let ConnectionSlotDataOwners { source, sink } = link.into_data();
321            drop(source);
322            assert_eq!(POOL.try_acquire().map(|lease| lease.is_none()), Ok(true));
323            drop(sink);
324            let replacement = POOL.try_acquire().ok().flatten();
325            assert!(replacement.is_some());
326            assert_eq!(POOL.try_acquire().map(|lease| lease.is_none()), Ok(true));
327        }
328    }
329
330    #[test]
331    fn cancelled_acquisition_leaves_capacity_available() {
332        static POOL: ConnectionSlotPool<CriticalSectionRawMutex, 1> = ConnectionSlotPool::new();
333
334        let held = POOL.try_acquire().ok().flatten();
335        assert!(held.is_some());
336        block_on(async {
337            assert!(matches!(
338                select(POOL.acquire(), ready(())).await,
339                Either::Second(())
340            ));
341        });
342        drop(held);
343        assert_eq!(POOL.try_acquire().map(|lease| lease.is_some()), Ok(true));
344    }
345
346    #[test]
347    fn wait_queue_saturation_is_explicit() {
348        static POOL: ConnectionSlotPool<CriticalSectionRawMutex, 1> = ConnectionSlotPool::new();
349
350        let held = POOL.try_acquire().ok().flatten();
351        assert!(held.is_some());
352        block_on(async {
353            assert!(matches!(
354                select(POOL.acquire(), POOL.acquire()).await,
355                Either::Second(Err(ConnectionSlotAcquireError::WaitQueueFull))
356            ));
357        });
358        drop(held);
359        assert_eq!(POOL.try_acquire().map(|lease| lease.is_some()), Ok(true));
360    }
361
362    #[test]
363    fn owner_drop_closes_peer_and_reuse_resets_close_signal() {
364        static POOL: ConnectionSlotPool<CriticalSectionRawMutex, 1> = ConnectionSlotPool::new();
365
366        let lease = POOL.try_acquire().ok().flatten();
367        assert!(lease.is_some());
368        if let Some(lease) = lease {
369            let ConnectionSlotOwners { worker, link } = lease.activate();
370            drop(link);
371            block_on(async {
372                assert!(matches!(
373                    select(worker.wait_for_close(), ready(())).await,
374                    Either::First(())
375                ));
376            });
377            assert_eq!(POOL.try_acquire().map(|lease| lease.is_none()), Ok(true));
378            drop(worker);
379
380            let reused = POOL.try_acquire().ok().flatten();
381            assert!(reused.is_some());
382            if let Some(reused) = reused {
383                let ConnectionSlotOwners { worker, link } = reused.activate();
384                block_on(async {
385                    assert!(matches!(
386                        select(worker.wait_for_close(), ready(())).await,
387                        Either::Second(())
388                    ));
389                });
390                drop(link);
391                drop(worker);
392            }
393        }
394    }
395
396    #[test]
397    fn explicit_close_requests_wake_workers_without_releasing_capacity() {
398        static POOL: ConnectionSlotPool<CriticalSectionRawMutex, 1> = ConnectionSlotPool::new();
399
400        let lease = POOL.try_acquire().ok().flatten();
401        assert!(lease.is_some());
402        if let Some(lease) = lease {
403            let ConnectionSlotOwners { worker, link } = lease.activate();
404            POOL.request_close(0);
405            block_on(async {
406                assert!(matches!(
407                    select(worker.wait_for_close(), ready(())).await,
408                    Either::First(())
409                ));
410            });
411            assert_eq!(POOL.try_acquire().map(|lease| lease.is_none()), Ok(true));
412            drop(link);
413            drop(worker);
414
415            let reused = POOL.try_acquire().ok().flatten();
416            assert!(reused.is_some());
417            if let Some(reused) = reused {
418                let ConnectionSlotOwners { worker, link } = reused.activate();
419                worker.request_close();
420                block_on(async {
421                    assert!(matches!(
422                        select(worker.wait_for_close(), ready(())).await,
423                        Either::First(())
424                    ));
425                });
426                drop(link);
427                drop(worker);
428            }
429        }
430    }
431}