rperf3-rs 0.6.1

//! Buffer pooling for efficient memory reuse
//!
//! This module provides a buffer pool to reduce allocation overhead during network tests.
//! Buffers are pre-allocated and reused across multiple send/receive operations, significantly
//! improving performance in high-throughput scenarios.
//!
//! # Performance Benefits
//!
//! Buffer pooling eliminates repeated allocations in hot paths, providing:
//! - **10-20% improvement** for UDP tests
//! - **5-10% improvement** for TCP tests
//! - Reduced garbage collection pressure
//! - Better cache locality through buffer reuse
//!
//! # Thread Safety
//!
//! The [`BufferPool`] is thread-safe and can be shared across multiple async tasks using
//! `Arc<BufferPool>`. Internal synchronization is handled automatically.
//!
//! # Examples
//!
//! ## Basic Usage
//!
//! ```
//! use rperf3::buffer_pool::BufferPool;
//!
//! // Create a pool with 1KB buffers, storing up to 10 buffers
//! let pool = BufferPool::new(1024, 10);
//!
//! // Get a buffer (allocates on first use)
//! let mut buffer = pool.get();
//! assert_eq!(buffer.len(), 1024);
//!
//! // Use the buffer for I/O operations
//! buffer[0] = 42;
//!
//! // Return it to the pool for reuse
//! pool.put(buffer);
//!
//! // Next get() will reuse the buffer
//! let buffer2 = pool.get();
//! assert_eq!(buffer2[0], 0); // Buffer is cleared on return
//! ```
//!
//! ## Pre-allocation for Performance
//!
//! ```
//! use rperf3::buffer_pool::BufferPool;
//!
//! let pool = BufferPool::new(8192, 5);
//!
//! // Pre-allocate all buffers before starting a test
//! pool.preallocate();
//! assert_eq!(pool.size(), 5);
//!
//! // Now all get() calls will reuse existing buffers
//! let buf = pool.get();
//! assert_eq!(buf.len(), 8192);
//! ```
//!
//! ## Shared Pool with Arc
//!
//! ```
//! use rperf3::buffer_pool::BufferPool;
//! use std::sync::Arc;
//!
//! let pool = Arc::new(BufferPool::new(4096, 8));
//!
//! // Clone the Arc to share across threads/tasks
//! let pool_clone = pool.clone();
//!
//! // Both references use the same pool
//! let buf1 = pool.get();
//! pool.put(buf1);
//!
//! let buf2 = pool_clone.get(); // Reuses the buffer
//! pool_clone.put(buf2);
//! ```

use std::sync::Mutex;

/// A thread-safe pool of reusable byte buffers
///
/// `BufferPool` manages a collection of pre-allocated buffers that can be reused across
/// multiple I/O operations, reducing allocation overhead and improving throughput.
///
/// # Thread Safety
///
/// All methods are thread-safe and can be called concurrently from multiple threads.
/// The pool uses a [`Mutex`] internally to synchronize access.
///
/// # Performance Characteristics
///
/// - **Get**: O(1) - pops from vector or allocates new buffer
/// - **Put**: O(1) - pushes to vector or drops if full
/// - **Size**: O(1) - returns vector length
/// - **Clear**: O(n) - drops all n buffers
///
/// # Examples
///
/// ```
/// use rperf3::buffer_pool::BufferPool;
///
/// // Create a pool for 64KB buffers (UDP max packet size)
/// let pool = BufferPool::new(65536, 10);
///
/// // Get a buffer for receiving data
/// let mut buffer = pool.get();
/// assert_eq!(buffer.len(), 65536);
///
/// // Simulate receiving data
/// buffer[0..5].copy_from_slice(b"hello");
///
/// // Return the buffer to the pool
/// pool.put(buffer);
///
/// // The next get will reuse the buffer (but it's cleared)
/// let buffer2 = pool.get();
/// assert_eq!(buffer2[0], 0); // Buffer was zeroed
/// ```
pub struct BufferPool {
    pool: Mutex<Vec<Vec<u8>>>,
    buffer_size: usize,
    max_pool_size: usize,
}

impl BufferPool {
    /// Creates a new buffer pool with the specified buffer size and capacity
    ///
    /// # Arguments
    ///
    /// * `buffer_size` - Size of each buffer in bytes. All buffers in the pool will be this size.
    /// * `max_pool_size` - Maximum number of buffers to keep in the pool. When the pool is full,
    ///   additional buffers returned via [`put`](Self::put) will be dropped.
    ///
    /// # Examples
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// // Create a pool for TCP with 128KB buffers, storing up to 20 buffers
    /// let tcp_pool = BufferPool::new(131072, 20);
    ///
    /// // Create a pool for UDP with 64KB buffers, storing up to 10 buffers
    /// let udp_pool = BufferPool::new(65536, 10);
    /// ```
    pub fn new(buffer_size: usize, max_pool_size: usize) -> Self {
        Self {
            pool: Mutex::new(Vec::with_capacity(max_pool_size)),
            buffer_size,
            max_pool_size,
        }
    }

    /// Gets a buffer from the pool or allocates a new one
    ///
    /// Returns a zeroed buffer of the configured size. If the pool contains available buffers,
    /// one is reused. Otherwise, a new buffer is allocated.
    ///
    /// # Returns
    ///
    /// A `Vec<u8>` of length `buffer_size`, filled with zeros.
    ///
    /// # Performance
    ///
    /// - **Pool hit**: O(1) - fast, just pops from internal vector
    /// - **Pool miss**: Allocates new buffer (slower, but only happens when pool is empty)
    ///
    /// # Examples
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(1024, 5);
    ///
    /// // First get allocates a new buffer
    /// let buf1 = pool.get();
    /// assert_eq!(buf1.len(), 1024);
    ///
    /// // Return it to the pool
    /// pool.put(buf1);
    ///
    /// // Second get reuses the buffer (faster!)
    /// let buf2 = pool.get();
    /// assert_eq!(buf2.len(), 1024);
    /// ```
    pub fn get(&self) -> Vec<u8> {
        let mut pool = self.pool.lock().unwrap();
        pool.pop().unwrap_or_else(|| vec![0u8; self.buffer_size])
    }

    /// Returns a buffer to the pool for reuse
    ///
    /// The buffer is cleared (all bytes set to 0) and added back to the pool if there's space.
    /// If the pool is at capacity, the buffer is dropped. Buffers with incorrect size are
    /// rejected and dropped immediately.
    ///
    /// # Arguments
    ///
    /// * `buffer` - The buffer to return to the pool. Must have length equal to `buffer_size`.
    ///
    /// # Behavior
    ///
    /// - **Size mismatch**: Buffer is dropped (not added to pool)
    /// - **Pool full**: Buffer is dropped (not added to pool)
    /// - **Pool has space**: Buffer is zeroed and added to pool
    ///
    /// # Examples
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(1024, 3);
    ///
    /// let mut buf = pool.get();
    /// buf[0] = 42; // Use the buffer
    ///
    /// // Return it - will be cleared and stored
    /// pool.put(buf);
    /// assert_eq!(pool.size(), 1);
    ///
    /// // Get it back - it's been zeroed
    /// let buf2 = pool.get();
    /// assert_eq!(buf2[0], 0);
    /// ```
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(1024, 2);
    ///
    /// // Fill the pool
    /// let buf1 = pool.get();
    /// let buf2 = pool.get();
    /// pool.put(buf1);
    /// pool.put(buf2);
    /// assert_eq!(pool.size(), 2);
    ///
    /// // Pool is full - this buffer will be dropped
    /// let buf3 = pool.get();
    /// pool.put(buf3);
    /// assert_eq!(pool.size(), 2); // Still 2
    /// ```
    pub fn put(&self, mut buffer: Vec<u8>) {
        // Only return buffers of the correct size
        if buffer.len() != self.buffer_size {
            return;
        }

        let mut pool = self.pool.lock().unwrap();
        if pool.len() < self.max_pool_size {
            // Clear the buffer for reuse
            buffer.fill(0);
            pool.push(buffer);
        }
        // Otherwise, let the buffer be dropped
    }

    /// Pre-allocates buffers to fill the pool to capacity
    ///
    /// This method allocates buffers up to `max_pool_size`, warming up the pool before
    /// starting performance-critical operations. This eliminates allocation overhead during
    /// the test run.
    ///
    /// If the pool already contains buffers, only the remaining slots are filled.
    ///
    /// # Examples
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(8192, 10);
    /// assert_eq!(pool.size(), 0);
    ///
    /// // Pre-allocate all buffers before a benchmark
    /// pool.preallocate();
    /// assert_eq!(pool.size(), 10);
    ///
    /// // All get() calls now reuse buffers (no allocation)
    /// let buf = pool.get();
    /// assert_eq!(buf.len(), 8192);
    /// ```
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(4096, 5);
    ///
    /// // Add some buffers manually
    /// let buf1 = pool.get();
    /// let buf2 = pool.get();
    /// pool.put(buf1);
    /// pool.put(buf2);
    /// assert_eq!(pool.size(), 2);
    ///
    /// // Preallocate fills the rest
    /// pool.preallocate();
    /// assert_eq!(pool.size(), 5);
    /// ```
    pub fn preallocate(&self) {
        let mut pool = self.pool.lock().unwrap();
        let current_size = pool.len();

        for _ in current_size..self.max_pool_size {
            pool.push(vec![0u8; self.buffer_size]);
        }
    }

    /// Returns the current number of buffers in the pool
    ///
    /// This count represents buffers that are available for reuse via [`get`](Self::get).
    /// Buffers that have been retrieved but not yet returned are not counted.
    ///
    /// # Returns
    ///
    /// The number of buffers currently stored in the pool (0 to `max_pool_size`).
    ///
    /// # Examples
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(1024, 5);
    /// assert_eq!(pool.size(), 0); // Empty pool
    ///
    /// // Add some buffers
    /// let buf1 = pool.get();
    /// let buf2 = pool.get();
    /// pool.put(buf1);
    /// pool.put(buf2);
    /// assert_eq!(pool.size(), 2);
    ///
    /// // Get one back
    /// let _buf = pool.get();
    /// assert_eq!(pool.size(), 1);
    /// ```
    pub fn size(&self) -> usize {
        self.pool.lock().unwrap().len()
    }

    /// Clears all buffers from the pool
    ///
    /// Removes and drops all buffers currently stored in the pool, freeing memory.
    /// After calling this method, [`size`](Self::size) will return 0.
    ///
    /// # Examples
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(1024, 10);
    /// pool.preallocate();
    /// assert_eq!(pool.size(), 10);
    ///
    /// // Clear all buffers
    /// pool.clear();
    /// assert_eq!(pool.size(), 0);
    /// ```
    ///
    /// ```
    /// use rperf3::buffer_pool::BufferPool;
    ///
    /// let pool = BufferPool::new(8192, 5);
    ///
    /// // Add buffers
    /// let mut buffers = Vec::new();
    /// for _ in 0..3 {
    ///     buffers.push(pool.get());
    /// }
    /// for buf in buffers {
    ///     pool.put(buf);
    /// }
    /// assert_eq!(pool.size(), 3);
    ///
    /// // Clear them
    /// pool.clear();
    /// assert_eq!(pool.size(), 0);
    ///
    /// // Pool can be reused after clearing
    /// let buf = pool.get();
    /// pool.put(buf);
    /// assert_eq!(pool.size(), 1);
    /// ```
    pub fn clear(&self) {
        self.pool.lock().unwrap().clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_buffer_pool_basic() {
        let pool = BufferPool::new(1024, 10);

        // Get a buffer
        let buf1 = pool.get();
        assert_eq!(buf1.len(), 1024);
        assert_eq!(pool.size(), 0);

        // Return it
        pool.put(buf1);
        assert_eq!(pool.size(), 1);

        // Get it again (should reuse)
        let buf2 = pool.get();
        assert_eq!(buf2.len(), 1024);
        assert_eq!(pool.size(), 0);
    }

    #[test]
    fn test_buffer_pool_max_size() {
        let pool = BufferPool::new(1024, 3);

        // Add 5 buffers, but only 3 should be kept
        let mut buffers = Vec::new();
        for _ in 0..5 {
            buffers.push(pool.get());
        }

        // Return all buffers
        for buf in buffers {
            pool.put(buf);
        }

        assert_eq!(pool.size(), 3);
    }

    #[test]
    fn test_buffer_pool_preallocate() {
        let pool = BufferPool::new(1024, 5);
        assert_eq!(pool.size(), 0);

        pool.preallocate();
        assert_eq!(pool.size(), 5);
    }

    #[test]
    fn test_buffer_pool_wrong_size() {
        let pool = BufferPool::new(1024, 3);

        // Try to return a buffer of wrong size
        let wrong_buf = vec![0u8; 2048];
        pool.put(wrong_buf);

        assert_eq!(pool.size(), 0);
    }

    #[test]
    fn test_buffer_pool_clear() {
        let pool = BufferPool::new(1024, 5);
        pool.preallocate();
        assert_eq!(pool.size(), 5);

        pool.clear();
        assert_eq!(pool.size(), 0);
    }
}