arcbox_datapath/
frame_buf.rs

1//! Owned frame buffer backed by [`PacketPool`] or a heap `Vec<u8>`.
2//!
3//! [`FrameBuf`] is the datapath's unit of frame ownership. It replaces raw
4//! `Vec<u8>` on the hot path with a pool-backed allocation that avoids the
5//! system allocator entirely (O(1) atomic CAS vs ~50-100 ns malloc).
6//!
7//! Code that produces frames from external sources (handshake synthesizer,
8//! DNS replies, ARP construction) can still use `FrameBuf::from(vec)` to
9//! wrap an existing `Vec<u8>` without copying. The datapath treats both
10//! variants uniformly via `Deref<Target = [u8]>`.
11
12use std::ops::Deref;
13use std::sync::Arc;
14
15use super::pool::PacketPool;
16
17/// A frame buffer backed by either the [`PacketPool`] or a heap `Vec`.
18///
19/// - `Pooled`: data lives in a pre-allocated [`PacketPool`] slot. On drop,
20///   the slot is returned to the pool (lock-free). This is the fast path.
21/// - `Heap`: data lives in a regular `Vec<u8>`. Used for frames that
22///   originate outside the pool (handshake replies, manually constructed
23///   frames).
24pub enum FrameBuf {
25    /// Pool-backed frame. Automatically freed on drop.
26    Pooled {
27        pool: Arc<PacketPool>,
28        index: u32,
29        len: u32,
30    },
31    /// Heap-backed frame (fallback for non-pool sources).
32    Heap(Vec<u8>),
33}
34
35impl FrameBuf {
36    /// Allocates a frame from the pool and copies `data` into it.
37    ///
38    /// Falls back to `Heap` if the pool is exhausted.
39    pub fn from_pool(pool: &Arc<PacketPool>, data: &[u8]) -> Self {
40        Self::from_pool_with_prefix(pool, &[], data)
41    }
42
43    /// Copies a header and payload into one owned frame without an intermediate buffer.
44    /// Falls back to `Heap` if the pool is exhausted or the combined frame is too large.
45    pub fn from_pool_with_prefix(pool: &Arc<PacketPool>, prefix: &[u8], data: &[u8]) -> Self {
46        let len = prefix.len() + data.len();
47        if let Some(mut pkt) = pool.alloc()
48            && let Some(frame) = pkt.as_full_mut_slice().get_mut(..len)
49        {
50            let (header, payload) = frame.split_at_mut(prefix.len());
51            header.copy_from_slice(prefix);
52            payload.copy_from_slice(data);
53            pkt.set_len(len);
54            let index = pkt.into_index();
55            return Self::Pooled {
56                pool: Arc::clone(pool),
57                index,
58                len: len as u32,
59            };
60        }
61        // Pool exhausted or frame too large — fall back to heap.
62        let mut frame = Vec::with_capacity(len);
63        frame.extend_from_slice(prefix);
64        frame.extend_from_slice(data);
65        Self::Heap(frame)
66    }
67
68    /// Returns the frame data length.
69    #[inline]
70    pub fn len(&self) -> usize {
71        match self {
72            Self::Pooled { len, .. } => *len as usize,
73            Self::Heap(v) => v.len(),
74        }
75    }
76
77    /// Returns true if the frame is empty.
78    #[inline]
79    pub fn is_empty(&self) -> bool {
80        self.len() == 0
81    }
82
83    /// Returns true if this frame is pool-backed.
84    #[inline]
85    pub fn is_pooled(&self) -> bool {
86        matches!(self, Self::Pooled { .. })
87    }
88}
89
90impl Deref for FrameBuf {
91    type Target = [u8];
92
93    #[inline]
94    fn deref(&self) -> &[u8] {
95        match self {
96            Self::Pooled { pool, index, len } => {
97                // SAFETY: The buffer at `index` is allocated (we hold ownership)
98                // and no other code has a mutable reference.
99                let buf = unsafe { pool.get(*index) };
100                &buf.as_full_slice()[..*len as usize]
101            }
102            Self::Heap(v) => v,
103        }
104    }
105}
106
107impl AsRef<[u8]> for FrameBuf {
108    #[inline]
109    fn as_ref(&self) -> &[u8] {
110        self
111    }
112}
113
114impl Drop for FrameBuf {
115    fn drop(&mut self) {
116        if let Self::Pooled { pool, index, .. } = self {
117            // SAFETY: We own this buffer slot exclusively. Returning it to the
118            // pool makes it available for reuse.
119            unsafe { pool.free_by_index(*index) };
120        }
121        // Heap variant: Vec<u8> drops normally.
122    }
123}
124
125impl From<Vec<u8>> for FrameBuf {
126    #[inline]
127    fn from(v: Vec<u8>) -> Self {
128        Self::Heap(v)
129    }
130}
131
132impl std::fmt::Debug for FrameBuf {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        match self {
135            Self::Pooled { index, len, .. } => f
136                .debug_struct("FrameBuf::Pooled")
137                .field("index", index)
138                .field("len", len)
139                .finish(),
140            Self::Heap(v) => f
141                .debug_struct("FrameBuf::Heap")
142                .field("len", &v.len())
143                .finish(),
144        }
145    }
146}
147
148// Compile-time assertion: FrameBuf must be Send (passed through channels).
149// Arc<PacketPool> is Send (PacketPool is Send+Sync), Vec<u8> is Send,
150// so FrameBuf derives Send automatically — no unsafe impl needed.
151const _: () = {
152    const fn assert_send<T: Send>() {}
153    assert_send::<FrameBuf>();
154};
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    fn make_pool() -> Arc<PacketPool> {
161        Arc::new(PacketPool::new(16).unwrap())
162    }
163
164    #[test]
165    fn test_from_pool() {
166        let pool = make_pool();
167        let data = [0xAB; 100];
168        let buf = FrameBuf::from_pool(&pool, &data);
169
170        assert!(buf.is_pooled());
171        assert_eq!(buf.len(), 100);
172        assert_eq!(&buf[..], &data[..]);
173        assert_eq!(pool.free_count(), 15);
174
175        drop(buf);
176        assert_eq!(pool.free_count(), 16);
177    }
178
179    #[test]
180    fn test_from_vec() {
181        let buf = FrameBuf::from(vec![1, 2, 3]);
182        assert!(!buf.is_pooled());
183        assert_eq!(buf.len(), 3);
184        assert_eq!(&buf[..], &[1, 2, 3]);
185    }
186
187    #[test]
188    fn test_pool_exhaustion_fallback() {
189        let pool = Arc::new(PacketPool::new(1).unwrap());
190
191        let _held = FrameBuf::from_pool(&pool, &[0; 10]);
192        assert_eq!(pool.free_count(), 0);
193
194        // Pool exhausted — should fall back to heap.
195        let fallback = FrameBuf::from_pool(&pool, &[0xFF; 20]);
196        assert!(!fallback.is_pooled());
197        assert_eq!(fallback.len(), 20);
198    }
199
200    #[test]
201    fn prefixed_frame_preserves_bytes_across_pool_exhaustion_and_reuse() {
202        let pool = Arc::new(PacketPool::new(1).unwrap());
203        let pooled = FrameBuf::from_pool_with_prefix(&pool, &[1, 2], &[3, 4, 5]);
204        let heap = FrameBuf::from_pool_with_prefix(&pool, &[6], &[7, 8]);
205        assert!(pooled.is_pooled());
206        assert!(!heap.is_pooled());
207        assert_eq!(&*pooled, &[1, 2, 3, 4, 5]);
208        assert_eq!(&*heap, &[6, 7, 8]);
209        drop(pooled);
210        let reused = FrameBuf::from_pool_with_prefix(&pool, &[9], &[]);
211        assert!(reused.is_pooled());
212        assert_eq!(&*reused, &[9]);
213        drop(reused);
214        assert_eq!(pool.free_count(), 1);
215    }
216
217    #[test]
218    fn prefixed_frame_exceeding_pool_slot_retains_the_complete_payload() {
219        let pool = Arc::new(PacketPool::new(1).unwrap());
220        let payload = vec![0xAB; crate::MAX_PACKET_SIZE];
221        let frame = FrameBuf::from_pool_with_prefix(&pool, &[1, 2], &payload);
222        assert!(!frame.is_pooled());
223        assert_eq!(&frame[..2], &[1, 2]);
224        assert_eq!(&frame[2..], payload);
225        assert_eq!(pool.free_count(), 1);
226    }
227
228    #[test]
229    fn test_deref_as_slice() {
230        let pool = make_pool();
231        let buf = FrameBuf::from_pool(&pool, &[10, 20, 30]);
232        let slice: &[u8] = &buf;
233        assert_eq!(slice, &[10, 20, 30]);
234    }
235
236    #[test]
237    fn test_as_ref() {
238        let buf = FrameBuf::from(vec![4, 5, 6]);
239        let r: &[u8] = buf.as_ref();
240        assert_eq!(r, &[4, 5, 6]);
241    }
242
243    #[test]
244    fn test_empty() {
245        let buf = FrameBuf::from(Vec::new());
246        assert!(buf.is_empty());
247        assert_eq!(buf.len(), 0);
248    }
249
250    #[test]
251    fn test_send() {
252        fn assert_send<T: Send>() {}
253        assert_send::<FrameBuf>();
254    }
255
256    #[test]
257    fn test_concurrent_alloc_drop() {
258        let pool = Arc::new(PacketPool::new(64).unwrap());
259        let handles: Vec<_> = (0..4)
260            .map(|_| {
261                let pool = Arc::clone(&pool);
262                std::thread::spawn(move || {
263                    for i in 0..500 {
264                        let data = vec![(i % 256) as u8; 100];
265                        let buf = FrameBuf::from_pool(&pool, &data);
266                        assert_eq!(buf.len(), 100);
267                        assert_eq!(buf[0], (i % 256) as u8);
268                        // buf dropped here — returned to pool
269                    }
270                })
271            })
272            .collect();
273
274        for h in handles {
275            h.join().unwrap();
276        }
277        assert_eq!(pool.free_count(), 64);
278    }
279}
arcbox_datapath/frame_buf.rs

arcbox_datapath/
frame_buf.rs
