Skip to main content

dynomite/io/
mbuf.rs

1//! Pooled fixed-size byte buffers.
2//!
3//! Connection read and write paths use chained fixed-size byte
4//! chunks. Each chunk has a read cursor (`pos`), a write cursor
5//! (`last`), a writable end boundary (`end`), and an extra trailing
6//! region (`end_extra`) that downstream code uses for crypto padding
7//! and overflow guards. Chunks are recycled through a free list so
8//! steady-state traffic does not call into the allocator.
9//!
10//! The pieces:
11//!
12//! * [`Mbuf`] owns a [`Box<[u8]>`] of `chunk_size` bytes plus the four
13//!   cursors that track readable, writable, and reserved regions.
14//! * [`MbufQueue`] is a tail-queue of chunks, backed by a
15//!   [`VecDeque`].
16//! * [`MbufPool`] owns the recycled buffers behind a parking-lot
17//!   mutex and is shared across worker tasks via [`Arc`].
18//!
19//! The default chunk size is [`MBUF_SIZE`] (16 KiB), tunable through
20//! [`MbufPool::new`]. The configurable range is
21//! `[MBUF_MIN_SIZE, MBUF_MAX_SIZE]`. The trailing [`MBUF_ESIZE`] bytes
22//! of every chunk are reserved for the crypto MAC region; normal
23//! writes stop at [`Mbuf::capacity`].
24//!
25//! # Examples
26//!
27//! ```
28//! use dynomite::io::mbuf::{Mbuf, MbufPool, MbufQueue};
29//!
30//! let pool = MbufPool::default();
31//! let mut buf = pool.get();
32//! assert_eq!(buf.recv(b"hello"), 5);
33//! assert_eq!(buf.len(), 5);
34//!
35//! let mut out = [0u8; 5];
36//! assert_eq!(buf.send(&mut out), 5);
37//! assert_eq!(&out, b"hello");
38//!
39//! let mut q = MbufQueue::new();
40//! q.push_back(buf);
41//! assert_eq!(q.len(), 1);
42//! ```
43
44use std::collections::VecDeque;
45use std::sync::Arc;
46
47use parking_lot::Mutex;
48
49/// Default mbuf chunk size in bytes. Mirrors the reference `MBUF_SIZE`.
50pub const MBUF_SIZE: usize = 16384;
51
52/// Minimum permitted mbuf chunk size. Mirrors the reference
53/// `MBUF_MIN_SIZE`.
54pub const MBUF_MIN_SIZE: usize = 512;
55
56/// Maximum permitted mbuf chunk size. Mirrors the reference
57/// `MBUF_MAX_SIZE`.
58pub const MBUF_MAX_SIZE: usize = 512_000;
59
60/// Bytes reserved at the tail of every chunk for crypto padding and
61/// overflow guards. Mirrors the reference `MBUF_ESIZE`.
62pub const MBUF_ESIZE: usize = 16;
63
64/// Maximum number of free chunks the pool retains before dropping
65/// returned buffers on the floor.
66///
67/// An upper bound on the free list keeps a misbehaving caller from
68/// driving unbounded memory growth. The default tracks a
69/// connection-budget order-of-magnitude (4096 connections * 4 chunks
70/// per direction).
71pub const MBUF_POOL_MAX_FREE: usize = 16384;
72
73/// Connection identifier carried by an mbuf to mark the conn that owns
74/// it. The reactor stamps this when it hands a chunk to a connection
75/// state machine; the chunk pool itself never inspects the value.
76pub type OwnerId = u64;
77
78/// A single chunk of a connection's I/O buffer chain.
79///
80/// An `Mbuf` exposes three regions:
81///
82/// * `[0, pos)` - already-consumed bytes (drained by the parser).
83/// * `[pos, last)` - readable bytes (parser input or pre-flight write).
84/// * `[last, end)` - writable bytes (target of [`recv`](Self::recv)).
85///
86/// The slice `[end, end_extra)` of length [`MBUF_ESIZE`] is reserved
87/// for crypto MAC and overflow guards. It is not visible through the
88/// normal write API.
89pub struct Mbuf {
90    buf: Box<[u8]>,
91    pos: usize,
92    last: usize,
93    end: usize,
94    flags: u32,
95    owner: Option<OwnerId>,
96}
97
98/// Bit flag set when the buffer has been flipped from write to read
99/// orientation. Mirrors `MBUF_FLAGS_READ_FLIP`.
100pub const MBUF_FLAG_READ_FLIP: u32 = 0x0000_0001;
101
102/// Bit flag set when the contents have just been decrypted. Mirrors
103/// `MBUF_FLAGS_JUST_DECRYPTED`.
104pub const MBUF_FLAG_JUST_DECRYPTED: u32 = 0x0000_0002;
105
106impl Mbuf {
107    /// Allocate a fresh chunk of the given total size. Used by the
108    /// pool to fill its free list and by callers that need a one-off
109    /// non-pooled buffer.
110    ///
111    /// `chunk_size` must be in `[MBUF_MIN_SIZE, MBUF_MAX_SIZE]`.
112    ///
113    /// # Examples
114    ///
115    /// ```
116    /// use dynomite::io::mbuf::{Mbuf, MBUF_SIZE};
117    /// let buf = Mbuf::with_chunk_size(MBUF_SIZE);
118    /// assert_eq!(buf.chunk_size(), MBUF_SIZE);
119    /// assert!(buf.is_empty());
120    /// ```
121    pub fn with_chunk_size(chunk_size: usize) -> Self {
122        assert!(
123            (MBUF_MIN_SIZE..=MBUF_MAX_SIZE).contains(&chunk_size),
124            "mbuf chunk_size {chunk_size} outside [{MBUF_MIN_SIZE}, {MBUF_MAX_SIZE}]"
125        );
126        let buf = vec![0u8; chunk_size].into_boxed_slice();
127        let end = chunk_size - MBUF_ESIZE;
128        Self {
129            buf,
130            pos: 0,
131            last: 0,
132            end,
133            flags: 0,
134            owner: None,
135        }
136    }
137
138    /// Total byte length of the chunk allocation.
139    ///
140    /// # Examples
141    ///
142    /// ```
143    /// use dynomite::io::mbuf::{Mbuf, MBUF_SIZE};
144    /// let buf = Mbuf::with_chunk_size(MBUF_SIZE);
145    /// assert_eq!(buf.chunk_size(), MBUF_SIZE);
146    /// ```
147    pub fn chunk_size(&self) -> usize {
148        self.buf.len()
149    }
150
151    /// Number of bytes that can be written before [`is_full`](Self::is_full)
152    /// trips, equal to `chunk_size - MBUF_ESIZE`. Mirrors `mbuf_data_size`.
153    ///
154    /// # Examples
155    ///
156    /// ```
157    /// use dynomite::io::mbuf::{Mbuf, MBUF_SIZE, MBUF_ESIZE};
158    /// let buf = Mbuf::with_chunk_size(MBUF_SIZE);
159    /// assert_eq!(buf.data_size(), MBUF_SIZE - MBUF_ESIZE);
160    /// ```
161    pub fn data_size(&self) -> usize {
162        self.end
163    }
164
165    /// Total byte addressable region including the trailing extra
166    /// area, equal to `chunk_size`. This spans past the
167    /// [`Mbuf::capacity`] write boundary up to the `end_extra` edge.
168    ///
169    /// # Examples
170    ///
171    /// ```
172    /// use dynomite::io::mbuf::{Mbuf, MBUF_SIZE};
173    /// let buf = Mbuf::with_chunk_size(MBUF_SIZE);
174    /// assert_eq!(buf.usable_capacity(), MBUF_SIZE);
175    /// ```
176    pub fn usable_capacity(&self) -> usize {
177        self.buf.len()
178    }
179
180    /// Bytes currently available to read from the buffer
181    /// (`last - pos`). Mirrors `mbuf_length`.
182    ///
183    /// # Examples
184    ///
185    /// ```
186    /// use dynomite::io::mbuf::Mbuf;
187    /// let mut buf = Mbuf::with_chunk_size(1024);
188    /// buf.recv(b"abc");
189    /// assert_eq!(buf.len(), 3);
190    /// ```
191    pub fn len(&self) -> usize {
192        self.last - self.pos
193    }
194
195    /// Maximum number of bytes the buffer can ever hold (the writable
196    /// region size, equal to [`data_size`](Self::data_size)). Mirrors
197    /// `mbuf_size` in the reference's higher-level usage.
198    pub fn capacity(&self) -> usize {
199        self.end
200    }
201
202    /// Bytes still writable into the buffer (`end - last`). Mirrors
203    /// `mbuf_remaining_space`.
204    ///
205    /// # Examples
206    ///
207    /// ```
208    /// use dynomite::io::mbuf::{Mbuf, MBUF_SIZE, MBUF_ESIZE};
209    /// let buf = Mbuf::with_chunk_size(MBUF_SIZE);
210    /// assert_eq!(buf.remaining(), MBUF_SIZE - MBUF_ESIZE);
211    /// ```
212    pub fn remaining(&self) -> usize {
213        self.end - self.last
214    }
215
216    /// True when no readable bytes remain. Mirrors `mbuf_empty`.
217    ///
218    /// # Examples
219    ///
220    /// ```
221    /// use dynomite::io::mbuf::Mbuf;
222    /// let buf = Mbuf::with_chunk_size(1024);
223    /// assert!(buf.is_empty());
224    /// ```
225    pub fn is_empty(&self) -> bool {
226        self.pos == self.last
227    }
228
229    /// True when the writable region is full. Mirrors `mbuf_full`.
230    ///
231    /// # Examples
232    ///
233    /// ```
234    /// use dynomite::io::mbuf::Mbuf;
235    /// let mut buf = Mbuf::with_chunk_size(1024);
236    /// while buf.remaining() > 0 {
237    ///     buf.recv(&[0]);
238    /// }
239    /// assert!(buf.is_full());
240    /// ```
241    pub fn is_full(&self) -> bool {
242        self.last == self.end
243    }
244
245    /// Discard all buffered bytes and rewind both cursors to the
246    /// origin. Mirrors `mbuf_rewind`.
247    pub fn rewind(&mut self) {
248        self.pos = 0;
249        self.last = 0;
250    }
251
252    /// Read the buffer's flag bits.
253    pub fn flags(&self) -> u32 {
254        self.flags
255    }
256
257    /// Set or clear individual flag bits.
258    pub fn set_flag(&mut self, flag: u32, on: bool) {
259        if on {
260            self.flags |= flag;
261        } else {
262            self.flags &= !flag;
263        }
264    }
265
266    /// Return the optional owner connection id stamped on the buffer.
267    pub fn owner(&self) -> Option<OwnerId> {
268        self.owner
269    }
270
271    /// Stamp or clear the owner connection id.
272    pub fn set_owner(&mut self, owner: Option<OwnerId>) {
273        self.owner = owner;
274    }
275
276    /// Borrow the readable region (`[pos, last)`).
277    ///
278    /// # Examples
279    ///
280    /// ```
281    /// use dynomite::io::mbuf::Mbuf;
282    /// let mut buf = Mbuf::with_chunk_size(1024);
283    /// buf.recv(b"abc");
284    /// assert_eq!(buf.readable(), b"abc");
285    /// ```
286    pub fn readable(&self) -> &[u8] {
287        &self.buf[self.pos..self.last]
288    }
289
290    /// Borrow the writable region (`[last, end)`).
291    pub fn writable(&mut self) -> &mut [u8] {
292        &mut self.buf[self.last..self.end]
293    }
294
295    /// Copy bytes from `src` into the writable region. Stops when the
296    /// buffer fills. Returns the number of bytes copied. Mirrors
297    /// `mbuf_recv` / `mbuf_copy` for the inbound direction.
298    ///
299    /// # Examples
300    ///
301    /// ```
302    /// use dynomite::io::mbuf::Mbuf;
303    /// let mut buf = Mbuf::with_chunk_size(1024);
304    /// assert_eq!(buf.recv(b"hello"), 5);
305    /// assert_eq!(buf.readable(), b"hello");
306    /// ```
307    pub fn recv(&mut self, src: &[u8]) -> usize {
308        let n = src.len().min(self.remaining());
309        self.buf[self.last..self.last + n].copy_from_slice(&src[..n]);
310        self.last += n;
311        n
312    }
313
314    /// Drain bytes from the readable region into `dst`. Returns the
315    /// number of bytes copied. Mirrors `mbuf_send` for the outbound
316    /// direction.
317    ///
318    /// # Examples
319    ///
320    /// ```
321    /// use dynomite::io::mbuf::Mbuf;
322    /// let mut buf = Mbuf::with_chunk_size(1024);
323    /// buf.recv(b"hello");
324    /// let mut out = [0u8; 8];
325    /// let n = buf.send(&mut out);
326    /// assert_eq!(n, 5);
327    /// assert_eq!(&out[..n], b"hello");
328    /// ```
329    pub fn send(&mut self, dst: &mut [u8]) -> usize {
330        let n = dst.len().min(self.len());
331        dst[..n].copy_from_slice(&self.buf[self.pos..self.pos + n]);
332        self.pos += n;
333        n
334    }
335
336    /// Copy `n` bytes from `src` into the writable region without
337    /// bounds-checking against partial copy. Panics if `n` exceeds
338    /// [`remaining`](Self::remaining). Mirrors `mbuf_copy`.
339    ///
340    /// # Examples
341    ///
342    /// ```
343    /// use dynomite::io::mbuf::Mbuf;
344    /// let mut buf = Mbuf::with_chunk_size(1024);
345    /// buf.copy_from_slice(b"abc");
346    /// assert_eq!(buf.readable(), b"abc");
347    /// ```
348    pub fn copy_from_slice(&mut self, src: &[u8]) {
349        assert!(
350            src.len() <= self.remaining(),
351            "mbuf copy of {} bytes exceeds remaining {}",
352            src.len(),
353            self.remaining()
354        );
355        let end = self.last + src.len();
356        self.buf[self.last..end].copy_from_slice(src);
357        self.last = end;
358    }
359
360    /// Split the buffer at offset `at` (relative to `pos`). The bytes
361    /// in `[pos+at, last)` are moved into a new mbuf taken from `pool`
362    /// and returned. The original keeps `[pos, pos+at)`. Mirrors
363    /// `mbuf_split` (without the precopy callback - callers that
364    /// previously injected a header into the new buffer can do so on
365    /// the returned `Mbuf` before chaining it).
366    ///
367    /// Returns `None` if `at` is greater than [`len`](Self::len) or if
368    /// the moved tail would not fit in a fresh pool buffer.
369    ///
370    /// # Examples
371    ///
372    /// ```
373    /// use dynomite::io::mbuf::MbufPool;
374    /// let pool = MbufPool::default();
375    /// let mut head = pool.get();
376    /// head.recv(b"hello world");
377    /// let tail = head.split_off(5, &pool).unwrap();
378    /// assert_eq!(head.readable(), b"hello");
379    /// assert_eq!(tail.readable(), b" world");
380    /// ```
381    pub fn split_off(&mut self, at: usize, pool: &MbufPool) -> Option<Mbuf> {
382        if at > self.len() {
383            return None;
384        }
385        let mut tail = pool.get();
386        let cut = self.pos + at;
387        let moved = self.last - cut;
388        if moved > tail.remaining() {
389            pool.put(tail);
390            return None;
391        }
392        tail.copy_from_slice(&self.buf[cut..self.last]);
393        self.last = cut;
394        Some(tail)
395    }
396
397    /// Append the readable region of `other` into this buffer.
398    /// Returns the number of bytes appended (which may be less than
399    /// `other.len()` if this buffer fills first).
400    ///
401    /// # Examples
402    ///
403    /// ```
404    /// use dynomite::io::mbuf::Mbuf;
405    /// let mut a = Mbuf::with_chunk_size(1024);
406    /// let mut b = Mbuf::with_chunk_size(1024);
407    /// a.recv(b"foo");
408    /// b.recv(b"bar");
409    /// a.append(&b);
410    /// assert_eq!(a.readable(), b"foobar");
411    /// ```
412    pub fn append(&mut self, other: &Mbuf) -> usize {
413        self.recv(other.readable())
414    }
415
416    /// Mark `n` bytes of the readable region as consumed by advancing
417    /// `pos`. Useful when downstream code wrote directly into
418    /// [`writable`](Self::writable). Panics if `n` exceeds
419    /// [`len`](Self::len).
420    pub fn advance_pos(&mut self, n: usize) {
421        assert!(
422            n <= self.len(),
423            "advance_pos {n} exceeds len {}",
424            self.len()
425        );
426        self.pos += n;
427    }
428
429    /// Mark `n` bytes of the writable region as filled by advancing
430    /// `last`. Used after writing into [`writable`](Self::writable)
431    /// directly (for example through an `AsyncRead` call). Panics if
432    /// `n` exceeds [`remaining`](Self::remaining).
433    pub fn advance_last(&mut self, n: usize) {
434        assert!(
435            n <= self.remaining(),
436            "advance_last {n} exceeds remaining {}",
437            self.remaining()
438        );
439        self.last += n;
440    }
441}
442
443impl std::fmt::Debug for Mbuf {
444    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
445        f.debug_struct("Mbuf")
446            .field("chunk_size", &self.buf.len())
447            .field("pos", &self.pos)
448            .field("last", &self.last)
449            .field("end", &self.end)
450            .field("flags", &self.flags)
451            .field("owner", &self.owner)
452            .finish()
453    }
454}
455
456/// Tail-queue of mbufs.
457///
458/// The container is a [`VecDeque`]; insertion at either end is O(1).
459///
460/// # Examples
461///
462/// ```
463/// use dynomite::io::mbuf::{MbufPool, MbufQueue};
464/// let pool = MbufPool::default();
465/// let mut q = MbufQueue::new();
466/// q.push_back(pool.get());
467/// q.push_back(pool.get());
468/// assert_eq!(q.len(), 2);
469/// let _ = q.pop_front();
470/// assert_eq!(q.len(), 1);
471/// ```
472#[derive(Debug, Default)]
473pub struct MbufQueue {
474    inner: VecDeque<Mbuf>,
475}
476
477impl MbufQueue {
478    /// Create an empty queue.
479    pub fn new() -> Self {
480        Self {
481            inner: VecDeque::new(),
482        }
483    }
484
485    /// Number of buffers currently chained.
486    pub fn len(&self) -> usize {
487        self.inner.len()
488    }
489
490    /// True when the queue holds no buffers.
491    pub fn is_empty(&self) -> bool {
492        self.inner.is_empty()
493    }
494
495    /// Append `mbuf` at the tail. Mirrors `mbuf_insert`.
496    pub fn push_back(&mut self, mbuf: Mbuf) {
497        self.inner.push_back(mbuf);
498    }
499
500    /// Insert `mbuf` at the head. Mirrors `mbuf_insert_head`.
501    pub fn push_front(&mut self, mbuf: Mbuf) {
502        self.inner.push_front(mbuf);
503    }
504
505    /// Remove and return the head buffer.
506    pub fn pop_front(&mut self) -> Option<Mbuf> {
507        self.inner.pop_front()
508    }
509
510    /// Remove and return the tail buffer.
511    pub fn pop_back(&mut self) -> Option<Mbuf> {
512        self.inner.pop_back()
513    }
514
515    /// Borrow the tail buffer mutably without removing it. Used by
516    /// the `mbuf_split` consumers that want to operate on the latest
517    /// chunk.
518    pub fn back_mut(&mut self) -> Option<&mut Mbuf> {
519        self.inner.back_mut()
520    }
521
522    /// Borrow the head buffer mutably without removing it.
523    pub fn front_mut(&mut self) -> Option<&mut Mbuf> {
524        self.inner.front_mut()
525    }
526
527    /// Iterate the queued buffers in head-to-tail order.
528    pub fn iter(&self) -> std::collections::vec_deque::Iter<'_, Mbuf> {
529        self.inner.iter()
530    }
531}
532
533impl<'a> IntoIterator for &'a MbufQueue {
534    type Item = &'a Mbuf;
535    type IntoIter = std::collections::vec_deque::Iter<'a, Mbuf>;
536
537    fn into_iter(self) -> Self::IntoIter {
538        self.inner.iter()
539    }
540}
541
542impl MbufQueue {
543    /// Total readable bytes across the chain.
544    pub fn total_len(&self) -> usize {
545        self.inner.iter().map(Mbuf::len).sum()
546    }
547
548    /// Drain the queue into the pool, recycling every chunk.
549    ///
550    /// # Examples
551    ///
552    /// ```
553    /// use dynomite::io::mbuf::{MbufPool, MbufQueue};
554    /// let pool = MbufPool::default();
555    /// let mut q = MbufQueue::new();
556    /// q.push_back(pool.get());
557    /// q.push_back(pool.get());
558    /// q.recycle(&pool);
559    /// assert!(q.is_empty());
560    /// ```
561    pub fn recycle(&mut self, pool: &MbufPool) {
562        while let Some(buf) = self.inner.pop_front() {
563            pool.put(buf);
564        }
565    }
566}
567
568/// Free-list backed mbuf pool.
569///
570/// The pool keeps a parking-lot mutex around a stash of recyclable
571/// chunk allocations. New chunks are taken from the stash if available
572/// and freshly allocated otherwise. [`MbufPool::put`] returns chunks
573/// to the stash up to [`MBUF_POOL_MAX_FREE`]; chunks beyond that cap
574/// are dropped.
575///
576/// The pool tracks total live and free counts for diagnostics.
577#[derive(Clone)]
578pub struct MbufPool {
579    inner: Arc<MbufPoolInner>,
580}
581
582struct MbufPoolInner {
583    chunk_size: usize,
584    max_free: usize,
585    state: Mutex<MbufPoolState>,
586}
587
588struct MbufPoolState {
589    free: Vec<Box<[u8]>>,
590    total_allocated: u64,
591}
592
593impl Default for MbufPool {
594    fn default() -> Self {
595        Self::new(MBUF_SIZE, MBUF_POOL_MAX_FREE)
596    }
597}
598
599impl MbufPool {
600    /// Construct a new pool with `chunk_size` byte chunks and a free
601    /// list capped at `max_free`. Mirrors `mbuf_init` plus the Rust-
602    /// only free-list bound.
603    ///
604    /// # Examples
605    ///
606    /// ```
607    /// use dynomite::io::mbuf::{MbufPool, MBUF_SIZE, MBUF_POOL_MAX_FREE};
608    /// let pool = MbufPool::new(MBUF_SIZE, MBUF_POOL_MAX_FREE);
609    /// let buf = pool.get();
610    /// assert_eq!(buf.chunk_size(), MBUF_SIZE);
611    /// pool.put(buf);
612    /// ```
613    pub fn new(chunk_size: usize, max_free: usize) -> Self {
614        assert!(
615            (MBUF_MIN_SIZE..=MBUF_MAX_SIZE).contains(&chunk_size),
616            "mbuf chunk_size {chunk_size} outside [{MBUF_MIN_SIZE}, {MBUF_MAX_SIZE}]"
617        );
618        Self {
619            inner: Arc::new(MbufPoolInner {
620                chunk_size,
621                max_free,
622                state: Mutex::new(MbufPoolState {
623                    free: Vec::new(),
624                    total_allocated: 0,
625                }),
626            }),
627        }
628    }
629
630    /// Configured chunk size, in bytes.
631    pub fn chunk_size(&self) -> usize {
632        self.inner.chunk_size
633    }
634
635    /// Maximum number of chunks the free list retains.
636    pub fn max_free(&self) -> usize {
637        self.inner.max_free
638    }
639
640    /// Take a fresh or recycled chunk from the pool. Mirrors
641    /// `mbuf_get`.
642    ///
643    /// # Examples
644    ///
645    /// ```
646    /// use dynomite::io::mbuf::MbufPool;
647    /// let pool = MbufPool::default();
648    /// let buf = pool.get();
649    /// assert!(buf.is_empty());
650    /// ```
651    pub fn get(&self) -> Mbuf {
652        let buf = self.alloc_buffer();
653        let chunk_size = buf.len();
654        let end = chunk_size - MBUF_ESIZE;
655        Mbuf {
656            buf,
657            pos: 0,
658            last: 0,
659            end,
660            flags: 0,
661            owner: None,
662        }
663    }
664
665    /// Return a chunk to the free list. Resets cursors and flags
666    /// before storing. Mirrors `mbuf_put`. Chunks past
667    /// [`max_free`](Self::max_free) are dropped.
668    ///
669    /// # Examples
670    ///
671    /// ```
672    /// use dynomite::io::mbuf::MbufPool;
673    /// let pool = MbufPool::default();
674    /// let buf = pool.get();
675    /// pool.put(buf);
676    /// assert_eq!(pool.free_count(), 1);
677    /// ```
678    pub fn put(&self, mut mbuf: Mbuf) {
679        if mbuf.buf.len() != self.inner.chunk_size {
680            // Off-size buffers (e.g. from [`Mbuf::with_chunk_size`])
681            // are simply dropped rather than poisoning the free list.
682            return;
683        }
684        mbuf.rewind();
685        mbuf.flags = 0;
686        mbuf.owner = None;
687        let mut state = self.inner.state.lock();
688        if state.free.len() < self.inner.max_free {
689            state.free.push(mbuf.buf);
690        }
691    }
692
693    /// Number of chunks currently sitting in the free list.
694    pub fn free_count(&self) -> usize {
695        self.inner.state.lock().free.len()
696    }
697
698    /// Lifetime count of fresh allocations performed by the pool.
699    /// Useful for tests asserting
700    /// that recycling avoids the allocator path.
701    pub fn total_allocated(&self) -> u64 {
702        self.inner.state.lock().total_allocated
703    }
704
705    fn alloc_buffer(&self) -> Box<[u8]> {
706        let mut state = self.inner.state.lock();
707        if let Some(buf) = state.free.pop() {
708            return buf;
709        }
710        state.total_allocated += 1;
711        drop(state);
712        vec![0u8; self.inner.chunk_size].into_boxed_slice()
713    }
714}
715
716impl std::fmt::Debug for MbufPool {
717    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
718        let state = self.inner.state.lock();
719        f.debug_struct("MbufPool")
720            .field("chunk_size", &self.inner.chunk_size)
721            .field("max_free", &self.inner.max_free)
722            .field("free_count", &state.free.len())
723            .field("total_allocated", &state.total_allocated)
724            .finish()
725    }
726}
727
728#[cfg(test)]
729mod tests {
730    use super::*;
731
732    #[test]
733    fn mbuf_recv_send_round_trip() {
734        let mut buf = Mbuf::with_chunk_size(1024);
735        assert_eq!(buf.recv(b"abc"), 3);
736        let mut out = [0u8; 8];
737        assert_eq!(buf.send(&mut out), 3);
738        assert_eq!(&out[..3], b"abc");
739        assert!(buf.is_empty());
740    }
741
742    #[test]
743    fn mbuf_recv_truncates_to_remaining() {
744        let mut buf = Mbuf::with_chunk_size(MBUF_MIN_SIZE);
745        let payload = vec![7u8; MBUF_MIN_SIZE * 2];
746        let n = buf.recv(&payload);
747        assert_eq!(n, MBUF_MIN_SIZE - MBUF_ESIZE);
748        assert!(buf.is_full());
749    }
750
751    #[test]
752    fn mbuf_split_then_append_reconstructs() {
753        let pool = MbufPool::default();
754        let mut head = pool.get();
755        head.recv(b"hello world");
756        let tail = head.split_off(5, &pool).unwrap();
757        assert_eq!(head.readable(), b"hello");
758        assert_eq!(tail.readable(), b" world");
759        head.append(&tail);
760        assert_eq!(head.readable(), b"hello world");
761    }
762
763    #[test]
764    fn mbuf_split_off_out_of_range_returns_none() {
765        let pool = MbufPool::default();
766        let mut head = pool.get();
767        head.recv(b"abc");
768        assert!(head.split_off(99, &pool).is_none());
769    }
770
771    #[test]
772    fn mbuf_pool_recycles_buffers_without_reallocation() {
773        let pool = MbufPool::default();
774        let n = 4;
775        let mut taken = Vec::new();
776        for _ in 0..n {
777            taken.push(pool.get());
778        }
779        assert_eq!(pool.total_allocated(), n as u64);
780        for buf in taken.drain(..) {
781            pool.put(buf);
782        }
783        assert_eq!(pool.free_count(), n);
784        for _ in 0..n {
785            taken.push(pool.get());
786        }
787        // Reused; allocation count is unchanged.
788        assert_eq!(pool.total_allocated(), n as u64);
789    }
790
791    #[test]
792    fn mbuf_pool_drops_chunks_beyond_cap() {
793        let cap = 4;
794        let pool = MbufPool::new(MBUF_SIZE, cap);
795        let n = cap * 2;
796        let mut taken = Vec::with_capacity(n);
797        for _ in 0..n {
798            taken.push(pool.get());
799        }
800        for buf in taken.drain(..) {
801            pool.put(buf);
802        }
803        // The pool refuses to grow past `max_free`; the surplus
804        // buffers are dropped on return rather than poisoning the
805        // free list.
806        assert_eq!(pool.free_count(), cap);
807    }
808
809    #[test]
810    fn mbuf_queue_push_pop_fifo() {
811        let pool = MbufPool::default();
812        let mut q = MbufQueue::new();
813        for i in 0..3u8 {
814            let mut buf = pool.get();
815            buf.recv(&[i]);
816            q.push_back(buf);
817        }
818        for i in 0..3u8 {
819            let buf = q.pop_front().unwrap();
820            assert_eq!(buf.readable(), &[i]);
821        }
822        assert!(q.is_empty());
823    }
824
825    #[test]
826    fn mbuf_flags_round_trip() {
827        let mut buf = Mbuf::with_chunk_size(1024);
828        buf.set_flag(MBUF_FLAG_READ_FLIP, true);
829        assert_eq!(buf.flags() & MBUF_FLAG_READ_FLIP, MBUF_FLAG_READ_FLIP);
830        buf.set_flag(MBUF_FLAG_READ_FLIP, false);
831        assert_eq!(buf.flags(), 0);
832    }
833
834    #[test]
835    fn mbuf_owner_round_trip() {
836        let mut buf = Mbuf::with_chunk_size(1024);
837        assert!(buf.owner().is_none());
838        buf.set_owner(Some(42));
839        assert_eq!(buf.owner(), Some(42));
840    }
841
842    #[test]
843    #[should_panic(expected = "outside")]
844    fn mbuf_too_small_panics() {
845        let _ = Mbuf::with_chunk_size(MBUF_MIN_SIZE - 1);
846    }
847
848    #[test]
849    #[should_panic(expected = "outside")]
850    fn mbuf_too_large_panics() {
851        let _ = Mbuf::with_chunk_size(MBUF_MAX_SIZE + 1);
852    }
853
854    #[test]
855    fn mbuf_put_drops_offsize_chunks() {
856        let pool = MbufPool::default();
857        // Use Mbuf::with_chunk_size to construct a buffer of a
858        // different size and confirm the pool refuses to admit it.
859        let stray = Mbuf::with_chunk_size(MBUF_MIN_SIZE);
860        pool.put(stray);
861        assert_eq!(pool.free_count(), 0);
862    }
863}