Skip to main content

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        if let Some(mut pkt) = pool.alloc()
41            && pkt.copy_from_slice(data).is_ok()
42        {
43            let index = pkt.into_index();
44            return Self::Pooled {
45                pool: Arc::clone(pool),
46                index,
47                len: data.len() as u32,
48            };
49        }
50        // Pool exhausted or frame too large — fall back to heap.
51        Self::Heap(data.to_vec())
52    }
53
54    /// Returns the frame data length.
55    #[inline]
56    pub fn len(&self) -> usize {
57        match self {
58            Self::Pooled { len, .. } => *len as usize,
59            Self::Heap(v) => v.len(),
60        }
61    }
62
63    /// Returns true if the frame is empty.
64    #[inline]
65    pub fn is_empty(&self) -> bool {
66        self.len() == 0
67    }
68
69    /// Returns true if this frame is pool-backed.
70    #[inline]
71    pub fn is_pooled(&self) -> bool {
72        matches!(self, Self::Pooled { .. })
73    }
74}
75
76impl Deref for FrameBuf {
77    type Target = [u8];
78
79    #[inline]
80    fn deref(&self) -> &[u8] {
81        match self {
82            Self::Pooled { pool, index, len } => {
83                // SAFETY: The buffer at `index` is allocated (we hold ownership)
84                // and no other code has a mutable reference.
85                let buf = unsafe { pool.get(*index) };
86                &buf.as_full_slice()[..*len as usize]
87            }
88            Self::Heap(v) => v,
89        }
90    }
91}
92
93impl AsRef<[u8]> for FrameBuf {
94    #[inline]
95    fn as_ref(&self) -> &[u8] {
96        self
97    }
98}
99
100impl Drop for FrameBuf {
101    fn drop(&mut self) {
102        if let Self::Pooled { pool, index, .. } = self {
103            // SAFETY: We own this buffer slot exclusively. Returning it to the
104            // pool makes it available for reuse.
105            unsafe { pool.free_by_index(*index) };
106        }
107        // Heap variant: Vec<u8> drops normally.
108    }
109}
110
111impl From<Vec<u8>> for FrameBuf {
112    #[inline]
113    fn from(v: Vec<u8>) -> Self {
114        Self::Heap(v)
115    }
116}
117
118impl std::fmt::Debug for FrameBuf {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        match self {
121            Self::Pooled { index, len, .. } => f
122                .debug_struct("FrameBuf::Pooled")
123                .field("index", index)
124                .field("len", len)
125                .finish(),
126            Self::Heap(v) => f
127                .debug_struct("FrameBuf::Heap")
128                .field("len", &v.len())
129                .finish(),
130        }
131    }
132}
133
134// Compile-time assertion: FrameBuf must be Send (passed through channels).
135// Arc<PacketPool> is Send (PacketPool is Send+Sync), Vec<u8> is Send,
136// so FrameBuf derives Send automatically — no unsafe impl needed.
137const _: () = {
138    const fn assert_send<T: Send>() {}
139    assert_send::<FrameBuf>();
140};
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    fn make_pool() -> Arc<PacketPool> {
147        Arc::new(PacketPool::new(16).unwrap())
148    }
149
150    #[test]
151    fn test_from_pool() {
152        let pool = make_pool();
153        let data = [0xAB; 100];
154        let buf = FrameBuf::from_pool(&pool, &data);
155
156        assert!(buf.is_pooled());
157        assert_eq!(buf.len(), 100);
158        assert_eq!(&buf[..], &data[..]);
159        assert_eq!(pool.free_count(), 15);
160
161        drop(buf);
162        assert_eq!(pool.free_count(), 16);
163    }
164
165    #[test]
166    fn test_from_vec() {
167        let buf = FrameBuf::from(vec![1, 2, 3]);
168        assert!(!buf.is_pooled());
169        assert_eq!(buf.len(), 3);
170        assert_eq!(&buf[..], &[1, 2, 3]);
171    }
172
173    #[test]
174    fn test_pool_exhaustion_fallback() {
175        let pool = Arc::new(PacketPool::new(1).unwrap());
176
177        let _held = FrameBuf::from_pool(&pool, &[0; 10]);
178        assert_eq!(pool.free_count(), 0);
179
180        // Pool exhausted — should fall back to heap.
181        let fallback = FrameBuf::from_pool(&pool, &[0xFF; 20]);
182        assert!(!fallback.is_pooled());
183        assert_eq!(fallback.len(), 20);
184    }
185
186    #[test]
187    fn test_deref_as_slice() {
188        let pool = make_pool();
189        let buf = FrameBuf::from_pool(&pool, &[10, 20, 30]);
190        let slice: &[u8] = &buf;
191        assert_eq!(slice, &[10, 20, 30]);
192    }
193
194    #[test]
195    fn test_as_ref() {
196        let buf = FrameBuf::from(vec![4, 5, 6]);
197        let r: &[u8] = buf.as_ref();
198        assert_eq!(r, &[4, 5, 6]);
199    }
200
201    #[test]
202    fn test_empty() {
203        let buf = FrameBuf::from(Vec::new());
204        assert!(buf.is_empty());
205        assert_eq!(buf.len(), 0);
206    }
207
208    #[test]
209    fn test_send() {
210        fn assert_send<T: Send>() {}
211        assert_send::<FrameBuf>();
212    }
213
214    #[test]
215    fn test_concurrent_alloc_drop() {
216        let pool = Arc::new(PacketPool::new(64).unwrap());
217        let handles: Vec<_> = (0..4)
218            .map(|_| {
219                let pool = Arc::clone(&pool);
220                std::thread::spawn(move || {
221                    for i in 0..500 {
222                        let data = vec![(i % 256) as u8; 100];
223                        let buf = FrameBuf::from_pool(&pool, &data);
224                        assert_eq!(buf.len(), 100);
225                        assert_eq!(buf[0], (i % 256) as u8);
226                        // buf dropped here — returned to pool
227                    }
228                })
229            })
230            .collect();
231
232        for h in handles {
233            h.join().unwrap();
234        }
235        assert_eq!(pool.free_count(), 64);
236    }
237}