trouble_host/
packet_pool.rs

1//! A packet pool for allocating and freeing packet buffers with quality of service policy.
2use core::cell::RefCell;
3
4use embassy_sync::blocking_mutex::raw::{CriticalSectionRawMutex, RawMutex};
5use embassy_sync::blocking_mutex::Mutex;
6
7use crate::{config, Packet, PacketPool};
8
9struct PacketBuf<const MTU: usize> {
10    buf: [u8; MTU],
11    free: bool,
12}
13
14impl<const MTU: usize> PacketBuf<MTU> {
15    const NEW: PacketBuf<MTU> = PacketBuf::new();
16
17    pub(crate) const fn new() -> Self {
18        Self {
19            buf: [0; MTU],
20            free: true,
21        }
22    }
23}
24
25struct State<const MTU: usize, const N: usize> {
26    packets: [PacketBuf<MTU>; N],
27}
28
29impl<const MTU: usize, const N: usize> State<MTU, N> {
30    pub(crate) const fn new() -> Self {
31        Self {
32            packets: [PacketBuf::NEW; N],
33        }
34    }
35
36    fn alloc(&mut self) -> Option<PacketRef<MTU>> {
37        for (idx, packet) in self.packets.iter_mut().enumerate() {
38            if packet.free {
39                // info!("[{}] alloc {}", id.0, idx);
40                packet.free = false;
41                packet.buf.iter_mut().for_each(|b| *b = 0);
42                return Some(PacketRef {
43                    idx,
44                    buf: packet.buf.as_mut_ptr(),
45                });
46            }
47        }
48        None
49    }
50
51    fn free(&mut self, p_ref: &PacketRef<MTU>) {
52        // info!("[{}] free {}", id.0, p_ref.idx);
53        self.packets[p_ref.idx].free = true;
54    }
55
56    fn available(&mut self) -> usize {
57        self.packets.iter().filter(|p| p.free).count()
58    }
59}
60
61/// A packet pool holds a pool of packet buffers that can be dynamically allocated
62/// and free'd.
63pub struct StaticPacketPool<M: RawMutex, const MTU: usize, const N: usize> {
64    state: Mutex<M, RefCell<State<MTU, N>>>,
65}
66
67impl<M: RawMutex, const MTU: usize, const N: usize> Default for StaticPacketPool<M, MTU, N> {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73impl<M: RawMutex, const MTU: usize, const N: usize> StaticPacketPool<M, MTU, N> {
74    /// Create a new packet pool with the given QoS policy
75    const fn new() -> Self {
76        Self {
77            state: Mutex::new(RefCell::new(State::new())),
78        }
79    }
80
81    fn alloc(&self) -> Option<PacketRef<MTU>> {
82        self.state.lock(|state| {
83            let mut state = state.borrow_mut();
84            state.alloc()
85        })
86    }
87
88    fn free(&self, p_ref: &PacketRef<MTU>) {
89        self.state.lock(|state| {
90            let mut state = state.borrow_mut();
91            state.free(p_ref);
92        });
93    }
94
95    fn available(&self) -> usize {
96        self.state.lock(|state| {
97            let mut state = state.borrow_mut();
98            state.available()
99        })
100    }
101}
102
103/// Represents a reference to a packet.
104#[repr(C)]
105pub struct PacketRef<const MTU: usize> {
106    idx: usize,
107    buf: *mut u8,
108}
109
110/// Global default packet pool.
111pub type DefaultPacketPool = StaticPacketPool<
112    CriticalSectionRawMutex,
113    { config::DEFAULT_PACKET_POOL_MTU },
114    { config::DEFAULT_PACKET_POOL_SIZE },
115>;
116
117static DEFAULT_POOL: StaticPacketPool<
118    CriticalSectionRawMutex,
119    { config::DEFAULT_PACKET_POOL_MTU },
120    { config::DEFAULT_PACKET_POOL_SIZE },
121> = StaticPacketPool::new();
122
123impl PacketPool for DefaultPacketPool {
124    type Packet = DefaultPacket;
125    const MTU: usize = { config::DEFAULT_PACKET_POOL_MTU };
126    fn capacity() -> usize {
127        config::DEFAULT_PACKET_POOL_SIZE
128    }
129
130    fn allocate() -> Option<DefaultPacket> {
131        DEFAULT_POOL.alloc().map(|p| DefaultPacket {
132            p_ref: p,
133            pool: &DEFAULT_POOL,
134        })
135    }
136}
137
138/// Type representing the packet from the default packet pool.
139pub struct DefaultPacket {
140    p_ref: PacketRef<{ config::DEFAULT_PACKET_POOL_MTU }>,
141    pool: &'static DefaultPacketPool,
142}
143
144impl Packet for DefaultPacket {}
145impl AsRef<[u8]> for DefaultPacket {
146    fn as_ref(&self) -> &[u8] {
147        unsafe { core::slice::from_raw_parts(self.p_ref.buf, config::DEFAULT_PACKET_POOL_MTU) }
148    }
149}
150
151impl AsMut<[u8]> for DefaultPacket {
152    fn as_mut(&mut self) -> &mut [u8] {
153        unsafe { core::slice::from_raw_parts_mut(self.p_ref.buf, config::DEFAULT_PACKET_POOL_MTU) }
154    }
155}
156
157impl Drop for DefaultPacket {
158    fn drop(&mut self) {
159        self.pool.free(&self.p_ref);
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use embassy_sync::blocking_mutex::raw::NoopRawMutex;
166
167    use super::*;
168
169    #[test]
170    fn test_none_qos() {
171        let pool: StaticPacketPool<NoopRawMutex, 27, 8> = StaticPacketPool::new();
172
173        let a1 = pool.alloc();
174        assert!(a1.is_some());
175        let a2 = pool.alloc();
176        assert!(a2.is_some());
177        let a3 = pool.alloc();
178        assert!(a3.is_some());
179        let a4 = pool.alloc();
180        assert!(a4.is_some());
181        let a5 = pool.alloc();
182        assert!(a5.is_some());
183        let a6 = pool.alloc();
184        assert!(a6.is_some());
185        let a7 = pool.alloc();
186        assert!(a7.is_some());
187
188        let b1 = pool.alloc();
189        assert!(b1.is_some());
190
191        let b2 = pool.alloc();
192        assert!(b2.is_none());
193    }
194}