rustfs_io_core/

pool.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Tiered buffer pool for zero-copy buffer management.
//!
//! Migrated from rustfs-ecstore to provide unified buffer pooling
//! across rustfs and rustfs-ecstore without cyclic dependencies.

use bytes::BytesMut;
use std::mem::ManuallyDrop;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};

// Tier size thresholds
const SMALL_MAX: usize = 64 * 1024;
const MEDIUM_MAX: usize = 512 * 1024;
const LARGE_MAX: usize = 4 * 1024 * 1024;

/// Tiered buffer pool for zero-copy buffer management.
///
/// This pool provides 4 tiers of buffers for different size ranges:
/// - Small: 4KB - 64KB
/// - Medium: 64KB - 512KB
/// - Large: 512KB - 4MB
/// - XLarge: > 4MB
///
/// Buffers are automatically reused when returned to the pool.
///
/// # Example
///
/// ```ignore
/// let pool = BytesPool::new_tiered();
///
/// // Acquire a buffer (automatically selects tier based on size)
/// let mut buffer = pool.acquire_buffer(8192).await;
///
/// // Use the buffer...
/// buffer.put_slice(b"hello world");
///
/// // Return to pool (automatic when dropped)
/// drop(buffer);
///
/// // Next acquisition will reuse the buffer
/// let mut buffer2 = pool.acquire_buffer(8192).await;
/// assert!(pool.hit_rate() > 0.0); // Buffer was reused!
/// ```
#[derive(Clone)]
pub struct BytesPool {
    /// Small object pool (4KB - 64KB)
    small_pool: Arc<PoolTier>,
    /// Medium object pool (64KB - 512KB)
    medium_pool: Arc<PoolTier>,
    /// Large object pool (512KB - 4MB)
    large_pool: Arc<PoolTier>,
    /// Extra large pool (> 4MB)
    xlarge_pool: Arc<PoolTier>,
    /// Pool metrics
    metrics: Arc<BytesPoolMetrics>,
}

/// Single pool tier with concurrent access control and buffer reuse.
struct PoolTier {
    /// Buffer size for this tier
    buffer_size: usize,
    /// Maximum concurrent buffers
    max_buffers: usize,
    /// Semaphore for concurrency control
    semaphore: Arc<Semaphore>,
    /// Pool name for metrics
    name: &'static str,
    /// Queue of available buffers for reuse
    available_buffers: Mutex<Vec<BytesMut>>,
    /// Metrics for tracking this tier
    metrics: Mutex<Option<Arc<BytesPoolMetrics>>>,
    /// Total acquisitions for this tier
    tier_total_acquires: AtomicU64,
    /// Total hits for this tier
    tier_pool_hits: AtomicU64,
    /// Current allocated bytes for this tier
    tier_current_allocated_bytes: AtomicU64,
}

/// Pool metrics for monitoring and optimization.
///
/// Tracks acquisition patterns and memory usage.
#[derive(Debug, Default)]
pub struct BytesPoolMetrics {
    /// Total buffer acquisitions
    pub total_acquires: AtomicU64,
    /// Pool hits (buffer reused)
    pub pool_hits: AtomicU64,
    /// Pool misses (new allocation)
    pub pool_misses: AtomicU64,
    /// Total bytes allocated
    pub total_bytes_allocated: AtomicU64,
    /// Current allocated bytes
    pub current_allocated_bytes: AtomicU64,
    /// Current available buffers in pool
    pub available_buffers: AtomicU64,
}

/// A buffer managed by the BytesPool.
///
/// When dropped, the buffer is automatically returned to the pool for reuse.
pub struct PooledBuffer {
    /// The underlying buffer (ManuallyDrop to allow taking on drop)
    pub buffer: ManuallyDrop<BytesMut>,
    /// Reference to pool tier for returning buffer
    tier: Option<Arc<PoolTier>>,
    /// The semaphore permit (must be dropped last to release slot)
    _permit: Option<OwnedSemaphorePermit>,
}

/// BytesPool configuration.
///
/// Allows customization of buffer sizes and limits for each tier.
pub struct BytesPoolConfig {
    pub small_size: usize,
    pub small_max: usize,
    pub medium_size: usize,
    pub medium_max: usize,
    pub large_size: usize,
    pub large_max: usize,
    pub xlarge_size: usize,
    pub xlarge_max: usize,
}

impl Default for BytesPoolConfig {
    fn default() -> Self {
        Self {
            small_size: 4 * 1024,
            small_max: 1000,
            medium_size: 64 * 1024,
            medium_max: 500,
            large_size: 512 * 1024,
            large_max: 100,
            xlarge_size: 4 * 1024 * 1024,
            xlarge_max: 25,
        }
    }
}

impl BytesPool {
    /// Create new tiered pool with default configuration.
    ///
    /// # Tier Configuration
    ///
    /// - Small: 4KB buffers, max 1000 concurrent
    /// - Medium: 64KB buffers, max 500 concurrent
    /// - Large: 512KB buffers, max 100 concurrent
    /// - XLarge: 4MB buffers, max 25 concurrent
    ///
    /// # Example
    ///
    /// ```ignore
    /// let pool = BytesPool::new_tiered();
    /// ```
    pub fn new_tiered() -> Self {
        Self::with_config(BytesPoolConfig::default())
    }

    /// Create pool with custom configuration.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let config = BytesPoolConfig {
    ///     small_size: 8 * 1024,  // 8KB small buffers
    ///     small_max: 2000,
    ///     ..Default::default()
    /// };
    /// let pool = BytesPool::with_config(config);
    /// ```
    pub fn with_config(config: BytesPoolConfig) -> Self {
        let metrics = Arc::new(BytesPoolMetrics::default());
        let small_pool = Arc::new(PoolTier::new(config.small_size, config.small_max, "small"));
        let medium_pool = Arc::new(PoolTier::new(config.medium_size, config.medium_max, "medium"));
        let large_pool = Arc::new(PoolTier::new(config.large_size, config.large_max, "large"));
        let xlarge_pool = Arc::new(PoolTier::new(config.xlarge_size, config.xlarge_max, "xlarge"));

        // Set metrics reference in all tiers
        small_pool.set_metrics(Arc::clone(&metrics));
        medium_pool.set_metrics(Arc::clone(&metrics));
        large_pool.set_metrics(Arc::clone(&metrics));
        xlarge_pool.set_metrics(Arc::clone(&metrics));

        Self {
            small_pool,
            medium_pool,
            large_pool,
            xlarge_pool,
            metrics,
        }
    }

    /// Acquire buffer with automatic tier selection.
    ///
    /// Selects the appropriate tier based on requested size and blocks
    /// until a buffer is available. Reuses returned buffers when available.
    ///
    /// # Arguments
    ///
    /// * `size` - Minimum capacity for the buffer
    ///
    /// # Returns
    ///
    /// A PooledBuffer that releases the permit and returns buffer to pool when dropped.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let mut buffer = pool.acquire_buffer(8192).await;
    /// ```
    pub async fn acquire_buffer(&self, size: usize) -> PooledBuffer {
        let tier = self.select_tier(size);
        let mut buffer = tier.acquire_buffer(size, &self.metrics).await;
        if buffer._permit.is_some() {
            buffer.tier = Some(Arc::clone(tier));
        }
        buffer
    }

    /// Try to acquire buffer without blocking.
    ///
    /// # Arguments
    ///
    /// * `size` - Minimum capacity for the buffer
    ///
    /// # Returns
    ///
    /// * `Some(buffer)` - If a buffer was available
    /// * `None` - If the pool is at capacity
    ///
    /// # Example
    ///
    /// ```ignore
    /// if let Some(mut buffer) = pool.try_acquire_buffer(8192) {
    ///     // Use buffer...
    /// }
    /// ```
    pub fn try_acquire_buffer(&self, size: usize) -> Option<PooledBuffer> {
        let tier = self.select_tier(size);
        let mut buffer = tier.try_acquire_buffer(size, &self.metrics)?;
        // Set tier reference for return on drop
        buffer.tier = Some(Arc::clone(tier));
        Some(buffer)
    }

    /// Select appropriate tier based on size.
    fn select_tier(&self, size: usize) -> &Arc<PoolTier> {
        if size <= SMALL_MAX {
            &self.small_pool
        } else if size <= MEDIUM_MAX {
            &self.medium_pool
        } else if size <= LARGE_MAX {
            &self.large_pool
        } else {
            &self.xlarge_pool
        }
    }

    /// Get pool metrics.
    pub fn metrics(&self) -> &BytesPoolMetrics {
        &self.metrics
    }

    /// Get pool hit rate (0.0 - 1.0).
    pub fn hit_rate(&self) -> f64 {
        let hits = self.metrics.pool_hits.load(Ordering::Relaxed);
        let total = self.metrics.total_acquires.load(Ordering::Relaxed);
        if total == 0 { 0.0 } else { hits as f64 / total as f64 }
    }

    /// Get the number of available buffers in the pool.
    pub fn available_buffers(&self) -> u64 {
        self.metrics.available_buffers.load(Ordering::Relaxed)
    }
}

impl PoolTier {
    fn new(buffer_size: usize, max_buffers: usize, name: &'static str) -> Self {
        Self {
            buffer_size,
            max_buffers,
            semaphore: Arc::new(Semaphore::new(max_buffers)),
            name,
            available_buffers: Mutex::new(Vec::new()),
            metrics: Mutex::new(None),
            tier_total_acquires: AtomicU64::new(0),
            tier_pool_hits: AtomicU64::new(0),
            tier_current_allocated_bytes: AtomicU64::new(0),
        }
    }

    fn set_metrics(&self, metrics: Arc<BytesPoolMetrics>) {
        *self.metrics.lock().unwrap_or_else(|e| e.into_inner()) = Some(metrics);
    }

    fn take_or_allocate_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> (BytesMut, bool) {
        let buffer_opt = {
            let mut available = self.available_buffers.lock().unwrap_or_else(|e| e.into_inner());
            available.pop()
        };
        let was_reused = buffer_opt.is_some();
        if was_reused {
            // A reused buffer leaves the available pool. `return_buffer` does a
            // matching `fetch_add(1)` on push, so without this the
            // `available_buffers` gauge only ever grows and never reflects the real
            // pool size (backlog#806).
            pool_metrics.available_buffers.fetch_sub(1, Ordering::Relaxed);
        }

        let buffer = if let Some(mut buf) = buffer_opt {
            let previous_capacity = buf.capacity();
            buf.clear();
            if previous_capacity < size {
                buf.reserve(size - previous_capacity);
            }
            let current_capacity = buf.capacity();
            if current_capacity > previous_capacity {
                let delta = (current_capacity - previous_capacity) as u64;
                pool_metrics.total_bytes_allocated.fetch_add(delta, Ordering::Relaxed);
                pool_metrics.current_allocated_bytes.fetch_add(delta, Ordering::Relaxed);
                self.tier_current_allocated_bytes.fetch_add(delta, Ordering::Relaxed);
            }
            buf
        } else {
            let buf = BytesMut::with_capacity(size.max(self.buffer_size));
            let allocated_bytes = buf.capacity() as u64;
            pool_metrics
                .total_bytes_allocated
                .fetch_add(allocated_bytes, Ordering::Relaxed);
            pool_metrics
                .current_allocated_bytes
                .fetch_add(allocated_bytes, Ordering::Relaxed);
            self.tier_current_allocated_bytes
                .fetch_add(allocated_bytes, Ordering::Relaxed);
            buf
        };

        (buffer, was_reused)
    }

    fn record_acquire_metrics(&self, pool_metrics: &BytesPoolMetrics, buffer_capacity: usize, was_reused: bool) {
        rustfs_io_metrics::record_bytes_pool_acquire(self.name, buffer_capacity, was_reused);

        if was_reused {
            pool_metrics.pool_hits.fetch_add(1, Ordering::Relaxed);
            self.tier_pool_hits.fetch_add(1, Ordering::Relaxed);
        } else {
            pool_metrics.pool_misses.fetch_add(1, Ordering::Relaxed);
        }

        let tier_total_acquires = self.tier_total_acquires.load(Ordering::Relaxed);
        let tier_pool_hits = self.tier_pool_hits.load(Ordering::Relaxed);
        let tier_hit_rate = if tier_total_acquires == 0 {
            0.0
        } else {
            tier_pool_hits as f64 / tier_total_acquires as f64
        };
        rustfs_io_metrics::record_bytes_pool_hit_rate(self.name, tier_hit_rate);
        rustfs_io_metrics::record_bytes_pool_allocated(self.name, self.tier_current_allocated_bytes.load(Ordering::Relaxed));
    }

    async fn acquire_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> PooledBuffer {
        // Acquire semaphore permit (owned for storage in PooledBuffer)
        let permit = match Arc::clone(&self.semaphore).acquire_owned().await {
            Ok(p) => p,
            Err(_) => {
                let buffer = BytesMut::with_capacity(size.max(self.buffer_size));
                return PooledBuffer {
                    buffer: ManuallyDrop::new(buffer),
                    tier: None,
                    _permit: None,
                };
            }
        };

        // Use the pool's shared metrics for recording
        let _metrics_lock = self.metrics.lock().unwrap_or_else(|e| e.into_inner());
        let _metrics = _metrics_lock.as_ref().expect("operation should succeed");

        // Record acquisition
        pool_metrics.total_acquires.fetch_add(1, Ordering::Relaxed);
        self.tier_total_acquires.fetch_add(1, Ordering::Relaxed);

        let (buffer, was_reused) = self.take_or_allocate_buffer(size, pool_metrics);
        let buffer_capacity = buffer.capacity();
        self.record_acquire_metrics(pool_metrics, buffer_capacity, was_reused);

        PooledBuffer {
            buffer: ManuallyDrop::new(buffer),
            tier: None, // Will be set after creating Arc<PoolTier>
            _permit: Some(permit),
        }
    }

    fn try_acquire_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> Option<PooledBuffer> {
        // Try to acquire permit without blocking
        let permit = Arc::clone(&self.semaphore).try_acquire_owned().ok()?;

        // Use the pool's shared metrics for recording
        let _metrics_lock = self.metrics.lock().unwrap_or_else(|e| e.into_inner());
        let _metrics = _metrics_lock.as_ref().expect("operation should succeed");

        // Record acquisition
        pool_metrics.total_acquires.fetch_add(1, Ordering::Relaxed);
        self.tier_total_acquires.fetch_add(1, Ordering::Relaxed);

        let (buffer, was_reused) = self.take_or_allocate_buffer(size, pool_metrics);
        let buffer_capacity = buffer.capacity();
        self.record_acquire_metrics(pool_metrics, buffer_capacity, was_reused);

        Some(PooledBuffer {
            buffer: ManuallyDrop::new(buffer),
            tier: None,
            _permit: Some(permit),
        })
    }

    /// Return a buffer to the pool for reuse without ever blocking the caller.
    fn return_buffer(&self, buffer: BytesMut) {
        let mut buffer = Some(buffer);

        if let Ok(mut available) = self.available_buffers.try_lock()
            && available.len() < self.max_buffers
        {
            available.push(buffer.take().expect("buffer should be present until returned"));
            if let Ok(metrics) = self.metrics.try_lock()
                && let Some(metrics) = metrics.as_ref()
            {
                metrics.available_buffers.fetch_add(1, Ordering::Relaxed);
            }
        }

        if let Some(buffer) = buffer {
            let released_bytes = buffer.capacity() as u64;
            self.tier_current_allocated_bytes
                .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
                    Some(current.saturating_sub(released_bytes))
                })
                .ok();
            if let Ok(metrics) = self.metrics.try_lock()
                && let Some(metrics) = metrics.as_ref()
            {
                metrics
                    .current_allocated_bytes
                    .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
                        Some(current.saturating_sub(released_bytes))
                    })
                    .ok();
            }
        }
        rustfs_io_metrics::record_bytes_pool_allocated(self.name, self.tier_current_allocated_bytes.load(Ordering::Relaxed));
    }
}

impl Drop for PooledBuffer {
    // SAFETY: Drop has exclusive access to `self`; taking the `ManuallyDrop`
    // buffer moves it exactly once into the pool when a tier still owns it.
    #[allow(unsafe_code)]
    fn drop(&mut self) {
        let buffer = unsafe { ManuallyDrop::take(&mut self.buffer) };
        if let Some(ref tier) = self.tier {
            tier.return_buffer(buffer);
        }
        // The permit is automatically dropped here, releasing the semaphore slot
    }
}

impl AsRef<[u8]> for PooledBuffer {
    fn as_ref(&self) -> &[u8] {
        self.buffer.as_ref()
    }
}

impl AsMut<[u8]> for PooledBuffer {
    fn as_mut(&mut self) -> &mut [u8] {
        self.buffer.as_mut()
    }
}

impl std::ops::Deref for PooledBuffer {
    type Target = BytesMut;

    fn deref(&self) -> &Self::Target {
        &self.buffer
    }
}

impl std::ops::DerefMut for PooledBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.buffer
    }
}

impl std::fmt::Debug for BytesPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BytesPool")
            .field("small_pool", &self.small_pool)
            .field("medium_pool", &self.medium_pool)
            .field("large_pool", &self.large_pool)
            .field("xlarge_pool", &self.xlarge_pool)
            .field("metrics", &self.metrics)
            .finish()
    }
}

impl std::fmt::Debug for PoolTier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PoolTier")
            .field("name", &self.name)
            .field("buffer_size", &self.buffer_size)
            .field("max_buffers", &self.max_buffers)
            .field("available_permits", &self.semaphore.available_permits())
            .field(
                "available_buffers",
                &self.available_buffers.lock().unwrap_or_else(|e| e.into_inner()).len(),
            )
            .finish()
    }
}

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

    #[tokio::test]
    async fn test_new_tiered() {
        let pool = BytesPool::new_tiered();
        assert_eq!(pool.small_pool.buffer_size, 4 * 1024);
        assert_eq!(pool.small_pool.max_buffers, 1000);
    }

    #[tokio::test]
    async fn test_acquire_buffer() {
        let pool = BytesPool::new_tiered();
        let buffer = pool.acquire_buffer(2048).await;
        assert!(buffer.capacity() >= 2048);
    }

    #[tokio::test]
    async fn available_buffers_gauge_decrements_on_reuse() {
        // Regression (backlog#806): acquiring a pooled (reused) buffer must
        // decrement the available_buffers gauge, mirroring return_buffer's
        // increment, so the gauge reflects the real pool size instead of only
        // growing.
        let pool = BytesPool::new_tiered();
        let buf = pool.acquire_buffer(2048).await;
        drop(buf); // returns to the pool -> gauge += 1
        assert_eq!(pool.available_buffers(), 1, "returned buffer should be counted");
        let buf2 = pool.acquire_buffer(2048).await; // reuses the pooled buffer -> gauge -= 1
        assert_eq!(pool.available_buffers(), 0, "reused buffer must decrement the gauge");
        drop(buf2);
        assert_eq!(pool.available_buffers(), 1, "gauge tracks the real pool size");
    }

    #[tokio::test]
    async fn test_acquire_buffer_after_shutdown_is_unpooled() {
        let pool = BytesPool::new_tiered();
        pool.small_pool.semaphore.close();

        let buffer = pool.acquire_buffer(2048).await;

        assert!(buffer.tier.is_none());
        assert!(buffer._permit.is_none());
        assert!(buffer.capacity() >= pool.small_pool.buffer_size);
        drop(buffer);
        assert_eq!(pool.available_buffers(), 0);
    }

    #[tokio::test]
    async fn test_tier_selection() {
        let pool = BytesPool::new_tiered();

        // Small buffer (4KB - 64KB)
        let buf1 = pool.acquire_buffer(1024).await;
        assert_eq!(buf1.capacity(), 4 * 1024);

        // Medium buffer (64KB - 512KB) - capacity is max(requested, tier_size)
        let buf2 = pool.acquire_buffer(100 * 1024).await;
        assert_eq!(buf2.capacity(), 100 * 1024); // Requested size

        // Large buffer (512KB - 4MB)
        let buf3 = pool.acquire_buffer(1024 * 1024).await;
        assert_eq!(buf3.capacity(), 1024 * 1024); // Requested size

        // XLarge buffer (> 4MB)
        let buf4 = pool.acquire_buffer(8 * 1024 * 1024).await;
        assert_eq!(buf4.capacity(), 8 * 1024 * 1024); // Requested size
    }

    #[tokio::test]
    async fn test_try_acquire_buffer() {
        let pool = BytesPool::with_config(BytesPoolConfig {
            small_size: 1024,
            small_max: 1,
            ..Default::default()
        });

        // First acquisition should succeed
        let buffer1 = pool.try_acquire_buffer(512);
        assert!(buffer1.is_some());

        // Second should fail (pool at capacity)
        let buffer2 = pool.try_acquire_buffer(512);
        assert!(buffer2.is_none());
    }

    #[tokio::test]
    async fn test_metrics() {
        let pool = BytesPool::new_tiered();
        let _buffer = pool.acquire_buffer(1024).await;
        drop(_buffer);

        let metrics = pool.metrics();
        assert!(metrics.total_acquires.load(Ordering::Relaxed) > 0);
    }

    #[tokio::test]
    async fn test_hit_rate() {
        let pool = BytesPool::new_tiered();
        assert_eq!(pool.hit_rate(), 0.0); // No acquisitions yet

        let _buffer = pool.acquire_buffer(1024).await;
        drop(_buffer);

        // First acquire is a miss (no buffers available yet)
        assert_eq!(pool.hit_rate(), 0.0);
    }

    #[tokio::test]
    async fn test_available_buffers() {
        let pool = BytesPool::new_tiered();
        assert_eq!(pool.available_buffers(), 0);

        let _buffer = pool.acquire_buffer(1024).await;
        drop(_buffer);

        // After drop, buffer should be returned to pool
        assert_eq!(pool.available_buffers(), 1);
    }

    #[tokio::test]
    async fn test_buffer_reuse() {
        // This test verifies that buffers are reused when returned to the pool
        let pool = BytesPool::with_config(BytesPoolConfig {
            small_size: 1024,
            small_max: 2,
            ..Default::default()
        });

        // Record initial state
        let initial_acquires = pool.metrics().total_acquires.load(Ordering::Relaxed);
        let initial_hits = pool.metrics().pool_hits.load(Ordering::Relaxed);
        assert_eq!(initial_acquires, 0);

        // First acquisition - should allocate new (miss)
        let buffer1 = pool.acquire_buffer(512).await;
        let initial_bytes_allocated = pool.metrics().total_bytes_allocated.load(Ordering::Relaxed);
        assert!(initial_bytes_allocated >= 1024);

        // Return buffer (by dropping)
        drop(buffer1);

        // Second acquisition - should reuse (hit)
        let _buffer2 = pool.acquire_buffer(512).await;
        let bytes_after_reuse = pool.metrics().total_bytes_allocated.load(Ordering::Relaxed);

        // Bytes allocated should be the same (buffer was reused)
        assert_eq!(initial_bytes_allocated, bytes_after_reuse);

        // Total acquires should be 2
        let total_acquires = pool.metrics().total_acquires.load(Ordering::Relaxed) - initial_acquires;
        assert_eq!(total_acquires, 2);

        // Pool hits should be 1
        let delta_hits = pool.metrics().pool_hits.load(Ordering::Relaxed) - initial_hits;
        assert_eq!(delta_hits, 1);
    }

    #[tokio::test]
    async fn test_tier_allocated_bytes_tracks_real_allocations() {
        let pool = BytesPool::with_config(BytesPoolConfig {
            small_size: 1024,
            small_max: 2,
            ..Default::default()
        });

        // First acquire allocates one small-tier buffer.
        let buf1 = pool.acquire_buffer(512).await;
        assert_eq!(pool.small_pool.tier_current_allocated_bytes.load(Ordering::Relaxed), 1024);

        // Return and reuse should not increase allocated bytes.
        drop(buf1);
        let buf2 = pool.acquire_buffer(512).await;
        assert_eq!(pool.small_pool.tier_current_allocated_bytes.load(Ordering::Relaxed), 1024);

        // A second in-flight buffer forces one more allocation.
        let _buf3 = pool.acquire_buffer(512).await;
        assert_eq!(pool.small_pool.tier_current_allocated_bytes.load(Ordering::Relaxed), 2048);

        drop(buf2);
    }

    #[tokio::test]
    async fn test_tier_hit_rate_counters_track_reuse() {
        let pool = BytesPool::with_config(BytesPoolConfig {
            small_size: 1024,
            small_max: 2,
            ..Default::default()
        });

        // First acquire is miss.
        let buf1 = pool.acquire_buffer(512).await;
        drop(buf1);

        // Second acquire reuses previous buffer and counts as hit.
        let _buf2 = pool.acquire_buffer(512).await;

        assert_eq!(pool.small_pool.tier_total_acquires.load(Ordering::Relaxed), 2);
        assert_eq!(pool.small_pool.tier_pool_hits.load(Ordering::Relaxed), 1);
    }
}