fxcp_core/

buffer.rs

// SPDX-License-Identifier: GPL-2.0-or-later
// Copyright (C) 2025 Joel Wirāmu Pauling <aenertia@aenertia.net>
//
// fxcp-core/src/buffer.rs  --  Aligned buffer pool for io_uring registered I/O

//! Memory-aligned buffer pool for io_uring fixed-buffer operations.
//! Provides pinned, aligned allocations that can be registered with the kernel.

use std::alloc::{alloc, dealloc, Layout};
use std::{ops::{Deref, DerefMut}, slice};
use tracing::{debug, error, trace, info};
use crate::metrics::{GLOBAL_BUFFER_COUNT, GLOBAL_BUFFER_LIMIT, GLOBAL_MEMORY_USAGE_BYTES};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::cell::UnsafeCell;
use crossbeam::queue::ArrayQueue;
use crate::error::{FxcpError, Result};
use crate::constants;
use std::ptr;

pub(crate) struct AlignedBuffer {
    ptr: *mut u8,
    layout: Layout,
    capacity: usize,
    len: AtomicUsize
}
// SAFETY: AlignedBuffer owns its heap allocation exclusively. The raw pointer
// is never aliased, and all length mutations use atomic operations, so it is
// safe to send the buffer to another thread.
unsafe impl Send for AlignedBuffer {}
// SAFETY: All mutable state (len) is behind AtomicUsize, so concurrent reads
// from multiple threads are data-race-free. The raw pointer is only written
// through &mut self (DerefMut), which the borrow checker serializes.
unsafe impl Sync for AlignedBuffer {}

impl AlignedBuffer {
    /// NUMA-aware allocation stub for enterprise server configurations.
    /// Falls back to standard allocation if node < 0 or mbind unavailable.
    #[allow(dead_code, unused_variables)]
    pub fn try_new_numa(capacity: usize, alignment: usize, numa_node: i32) -> Result<Self> {
        // TODO: On dual-socket servers, use mmap + mbind(MPOL_BIND) to pin
        // buffers to the same NUMA node as the NVMe controller to avoid
        // QPI cross-talk. Discover node via /sys/class/block/*/device/numa_node.
        // For now, fall back to standard allocation.
        Self::try_new(capacity, alignment)
    }

    /// Allocate an aligned buffer with the given capacity and alignment.
    pub fn try_new(capacity: usize, alignment: usize) -> Result<Self> {
        // Fix: Limit is already in bytes, don't multiply by 1024*1024 again
        let limit_bytes = GLOBAL_BUFFER_LIMIT.get() as u64;
        let align = alignment.max(constants::MINIMUM_ALIGNMENT_BYTES);
        let actual_capacity = capacity.max(align);
        
        let _prev = GLOBAL_BUFFER_COUNT.try_update(
            Ordering::SeqCst,
            Ordering::SeqCst,
            |current| {
                if current + actual_capacity as u64 > limit_bytes {
                    None
                } else {
                    Some(current + actual_capacity as u64)
                }
            }
        ).map_err(|_| FxcpError::MemoryExhausted(format!(
            "Global memory limit exceeded. Refusing allocation of {} bytes.", capacity
        )))?;

        // Update Prometheus Gauge
        GLOBAL_MEMORY_USAGE_BYTES.add(actual_capacity as f64);

        let layout = Layout::from_size_align(actual_capacity, align)
            .map_err(|_e| {
                GLOBAL_BUFFER_COUNT.fetch_sub(actual_capacity as u64, Ordering::SeqCst);
                GLOBAL_MEMORY_USAGE_BYTES.sub(actual_capacity as f64);
                FxcpError::System(nix::Error::from(nix::errno::Errno::EINVAL))
            })?;
            
        // SAFETY: layout is valid (from_size_align succeeded above) with non-zero
        // size and power-of-two alignment. alloc returns a valid pointer or null.
        let ptr = unsafe { alloc(layout) };
        if ptr.is_null() {
            GLOBAL_BUFFER_COUNT.fetch_sub(actual_capacity as u64, Ordering::SeqCst);
            GLOBAL_MEMORY_USAGE_BYTES.sub(actual_capacity as f64);
            return Err(FxcpError::MemoryExhausted("Physical memory allocation failed".to_string()));
        }
        // SAFETY: ptr is non-null (checked above), and actual_capacity matches
        // the allocation size, so the write stays within the allocated region.
        unsafe { ptr::write_bytes(ptr, 0, actual_capacity); }

        // Pre-fault pages for large buffers to avoid page faults during io_uring I/O.
        // MADV_POPULATE_WRITE (Linux 5.14+) is an optimization; silently ignore EINVAL
        // on older kernels.
        if actual_capacity >= 1_048_576
            && let Some(nn) = std::ptr::NonNull::new(ptr) {
                // SAFETY: ptr is a valid allocation of actual_capacity bytes from alloc().
                // MADV_POPULATE_WRITE pre-faults pages for write access.
                let _ = unsafe {
                    nix::sys::mman::madvise(nn.cast(), actual_capacity, nix::sys::mman::MmapAdvise::MADV_POPULATE_WRITE)
                };
            }
        trace!("AlignedBuffer: Allocated {} bytes.", actual_capacity);
        Ok(Self {
            ptr,
            layout,
            capacity: actual_capacity,
            len: AtomicUsize::new(0)
        })
    }
    #[inline(always)]
    pub fn clear(&self) {
        self.len.store(0, Ordering::Release);
    }
    #[inline(always)]
    pub fn capacity(&self) -> usize { self.capacity }
    #[inline(always)]
    #[allow(dead_code)]
    pub fn set_full_len(&self) {
        self.len.store(self.capacity, Ordering::Release);
    }
    #[inline(always)]
    pub fn set_len(&self, len: usize) -> Result<()> {
        if len > self.capacity {
            return Err(FxcpError::Config(format!(
                "AlignedBuffer::set_len: {} exceeds capacity {}", len, self.capacity
            )));
        }
        self.len.store(len, Ordering::Release);
        Ok(())
    }
    #[inline(always)]
    pub fn get_len(&self) -> usize {
        self.len.load(Ordering::Acquire)
    }
    #[inline(always)]
    pub fn ptr(&self) -> *mut u8 { self.ptr }
    #[inline(always)]
    #[allow(dead_code)]
    pub fn alignment(&self) -> usize { self.layout.align() }
}

impl Drop for AlignedBuffer {
    fn drop(&mut self) {
        let _prev = GLOBAL_BUFFER_COUNT.fetch_sub(self.capacity as u64, Ordering::SeqCst);
        GLOBAL_MEMORY_USAGE_BYTES.sub(self.capacity as f64);
        // SAFETY: self.ptr was allocated with self.layout in try_new, and this
        // is the only deallocation (Drop runs exactly once).
        unsafe { dealloc(self.ptr, self.layout); }
    }
}

impl Deref for AlignedBuffer {
    type Target = [u8];
    fn deref(&self) -> &Self::Target {
        // SAFETY: self.ptr is valid for self.capacity bytes (allocated in try_new),
        // and get_len() is always <= capacity (enforced by set_len bounds check).
        unsafe { slice::from_raw_parts(self.ptr, self.get_len()) }
    }
}

impl DerefMut for AlignedBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: same as Deref, plus we have &mut self so no aliasing is possible.
        unsafe { slice::from_raw_parts_mut(self.ptr, self.get_len()) }
    }
}

enum QueueStrategy {
    Shared(ArrayQueue<u16>),
    Local(UnsafeCell<Vec<u16>>),
}
// SAFETY: QueueStrategy::Shared wraps ArrayQueue (already Send+Sync).
// QueueStrategy::Local wraps UnsafeCell<Vec<u16>>, which is only accessed
// from a single thread in local_mode (thread-local pool pattern).
unsafe impl Send for QueueStrategy {}
unsafe impl Sync for QueueStrategy {}

struct BufferPoolInner {
    buffers: Vec<UnsafeCell<AlignedBuffer>>,
    free_indices: QueueStrategy,
    capacity: usize,
    chunk_size: usize,
    alignment: usize,
}
// SAFETY: BufferPoolInner's Vec<UnsafeCell<AlignedBuffer>> is accessed through
// index-based acquire/release discipline. Each buffer index is held by at most
// one caller at a time (enforced by the free_indices queue), preventing aliasing.
unsafe impl Sync for BufferPoolInner {}
unsafe impl Send for BufferPoolInner {}

/// Thread-safe pool of pre-allocated aligned buffers for io_uring fixed-buffer I/O.
#[derive(Clone)]
pub struct BufferPool {
    inner: Arc<BufferPoolInner>,
}

impl BufferPool {
    /// Create a shared buffer pool with cross-thread access.
    pub fn new(requested_buffers: usize, requested_chunk_size: usize, alignment: usize) -> Result<Self> {
        Self::create_pool(requested_buffers, requested_chunk_size, alignment, false)
    }
    /// Create a thread-local buffer pool (no cross-thread synchronization).
    pub fn new_local(requested_buffers: usize, requested_chunk_size: usize, alignment: usize) -> Result<Self> {
        Self::create_pool(requested_buffers, requested_chunk_size, alignment, true)
    }
    fn create_pool(requested_buffers: usize, requested_chunk_size: usize, alignment: usize, local_mode: bool) -> Result<Self> {
        if requested_buffers > constants::BUFFER_POOL_MAX_REGISTERED_BUFFERS as usize {
            return Err(FxcpError::Config("BufferPool: max 65535 buffers allowed".to_string()));
        }
        let align = alignment.max(constants::MINIMUM_ALIGNMENT_BYTES);
        let min_chunk_size = constants::BUFFER_POOL_MIN_CHUNK_SIZE;
        let strategies = [
            (requested_buffers, requested_chunk_size),
            (requested_buffers / 2, requested_chunk_size),
            (requested_buffers / 4, requested_chunk_size),
            (requested_buffers, requested_chunk_size / 2),
            (requested_buffers / 2, requested_chunk_size / 2),
            (8, min_chunk_size),
        ];
        
        for (count, size) in strategies {
            if count < 2 || size < min_chunk_size { continue; }
            let mut buffers = Vec::with_capacity(count);
            let mut allocated_successfully = true;
            for i in 0..count {
                match AlignedBuffer::try_new(size, align) {
                    Ok(buf) => buffers.push(UnsafeCell::new(buf)),
                    Err(e) => {
                        debug!("BufferPool: Allocation failed at buffer {} of {}: {:?}", i, count, e);
                        allocated_successfully = false;
                        break;
                    }
                }
            }
            if !allocated_successfully { continue; }
            
            let actual_capacity = buffers.len();
            let queue = if local_mode {
                let mut v = Vec::with_capacity(actual_capacity);
                for i in 0..actual_capacity {
                    v.push(i as u16);
                }
                QueueStrategy::Local(UnsafeCell::new(v))
            } else {
                let q = ArrayQueue::new(actual_capacity);
                for i in 0..actual_capacity {
                    let _ = q.push(i as u16);
                }
                QueueStrategy::Shared(q)
            };
            
            let total_mb = (actual_capacity as u64 * size as u64) / 1024 / 1024;
            info!("BufferPool: Allocated {} x {}KB buffers ({}MB total, Mode: {})",
                  actual_capacity, size / 1024, total_mb, if local_mode { "Thread-Local" } else { "Shared" });
            crate::metrics::BUFFER_POOL_CAPACITY.set(actual_capacity as f64);
            crate::metrics::BUFFER_POOL_CHUNK_SIZE.set(size as f64);
            crate::metrics::BUFFER_POOL_TOTAL_BYTES.set((actual_capacity as u64 * size as u64) as f64);
            
            return Ok(Self {
                inner: Arc::new(BufferPoolInner {
                    buffers,
                    free_indices: queue,
                    capacity: actual_capacity,
                    chunk_size: size,
                    alignment: align,
                })
            });
        }
        Err(FxcpError::MemoryExhausted("BufferPool: All allocation strategies failed.".to_string()))
    }
    
    /// Returns the total number of buffers in the pool.
    #[inline(always)]
    pub fn capacity(&self) -> usize { self.inner.capacity }
    /// Returns the size in bytes of each buffer chunk.
    #[inline(always)]
    pub fn chunk_size(&self) -> usize { self.inner.chunk_size }
    /// Returns the memory alignment of each buffer in bytes.
    #[inline(always)]
    pub fn alignment(&self) -> usize { self.inner.alignment }
    /// Returns the raw pointer for the buffer at the given index.
    #[inline(always)]
    pub fn get_ptr(&self, index: u16) -> Option<*mut u8> {
        self.inner.buffers.get(index as usize).map(|cell| {
            // Build io_uring iovec descriptors from the pool's buffers.
            // SAFETY: index was obtained from acquire(), so no other caller
            // holds this index. We only read the pointer, not mutate the buffer.
            unsafe { (*cell.get()).ptr() }
        })
    }
    /// Returns the current data length of the buffer at the given index.
    #[inline(always)]
    pub fn get_len(&self, index: u16) -> usize {
        if let Some(cell) = self.inner.buffers.get(index as usize) {
            // SAFETY: index bounds-checked by get(). Caller holds exclusive
            // logical ownership of this index via acquire/release protocol.
            unsafe { (*cell.get()).get_len() }
        } else {
            0
        }
    }
    /// Sets the data length of the buffer at the given index.
    #[inline(always)]
    pub fn set_len(&self, index: u16, len: usize) -> Result<()> {
        if let Some(cell) = self.inner.buffers.get(index as usize) {
            // SAFETY: caller holds exclusive logical ownership of this index.
            // set_len only stores to an AtomicUsize, so no data race.
            unsafe {
                let buf = &*cell.get();
                buf.set_len(len)?;
            }
        }
        Ok(())
    }
    /// Build io_uring iovec descriptors from the pool's buffers.
    pub fn as_io_vecs(&mut self) -> Vec<libc::iovec> {
        self.inner.buffers.iter().map(|cell| {
            // SAFETY: &mut self guarantees exclusive access to the pool.
            // No buffers are acquired (caller must ensure this before registering).
            unsafe {
                let buf = &*cell.get();
                libc::iovec { iov_base: buf.ptr() as _, iov_len: buf.capacity() }
            }
        }).collect()
    }
    /// Acquires a free buffer index from the pool, returning `None` if exhausted.
    #[inline]
    pub fn acquire(&self) -> Option<u16> {
        match &self.inner.free_indices {
            QueueStrategy::Shared(q) => q.pop(),
            QueueStrategy::Local(cell) => {
                // SAFETY: Local mode is single-threaded; only one caller
                // accesses this UnsafeCell at a time.
                let vec = unsafe { &mut *cell.get() };
                vec.pop()
            }
        }
    }
    /// Returns a buffer to the pool, clearing its data length.
    #[inline]
    pub fn release(&self, index: u16) {
        if let Some(cell) = self.inner.buffers.get(index as usize) {
            // SAFETY: caller is returning exclusive ownership of this index.
            // clear() only stores to an AtomicUsize.
            unsafe { (*cell.get()).clear(); }
            match &self.inner.free_indices {
                QueueStrategy::Shared(q) => {
                    let _ = q.push(index);
                },
                QueueStrategy::Local(c) => {
                    // SAFETY: Local mode is single-threaded.
                    let vec = unsafe { &mut *c.get() };
                    vec.push(index);
                }
            }
        } else {
            error!("BufferPool: Attempted to release invalid index {}", index);
        }
    }
    /// Returns the number of currently available buffers in the pool.
    #[inline]
    pub fn free_count(&self) -> usize {
        match &self.inner.free_indices {
            QueueStrategy::Shared(q) => q.len(),
            // SAFETY: Local mode is single-threaded.
            QueueStrategy::Local(c) => unsafe { (*c.get()).len() },
        }
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
    use super::*;
    use std::sync::Mutex;

    /// Tests that mutate `GLOBAL_BUFFER_LIMIT` must hold this lock to prevent
    /// parallel test threads from racing on the shared Prometheus gauge.
    static BUFFER_LIMIT_LOCK: Mutex<()> = Mutex::new(());

    #[test]
    fn test_buffer_set_len_overflow_returns_error() {
        let _guard = BUFFER_LIMIT_LOCK.lock().unwrap();
        let saved = GLOBAL_BUFFER_LIMIT.get();
        GLOBAL_BUFFER_LIMIT.set(1024.0 * 1024.0);

        let buf = AlignedBuffer::try_new(16, 4096).unwrap();
        let cap = buf.capacity();
        assert!(buf.set_len(cap + 1).is_err(), "set_len beyond capacity should return Err");
        assert!(buf.set_len(cap).is_ok(), "set_len at exact capacity should succeed");
        assert!(buf.set_len(0).is_ok(), "set_len to zero should succeed");

        GLOBAL_BUFFER_LIMIT.set(saved);
    }

    #[test]
    fn test_large_buffer_madv_populate() {
        let _guard = BUFFER_LIMIT_LOCK.lock().unwrap();
        let saved = GLOBAL_BUFFER_LIMIT.get();
        GLOBAL_BUFFER_LIMIT.set(64.0 * 1024.0 * 1024.0);

        let buf = AlignedBuffer::try_new(1_048_576, 4096);
        assert!(buf.is_ok(), "1MB buffer allocation should succeed");
        let buf = buf.unwrap();
        assert!(buf.capacity() >= 1_048_576);

        GLOBAL_BUFFER_LIMIT.set(saved);
    }

    #[test]
    fn test_small_buffer_no_populate() {
        let _guard = BUFFER_LIMIT_LOCK.lock().unwrap();
        let saved = GLOBAL_BUFFER_LIMIT.get();
        GLOBAL_BUFFER_LIMIT.set(64.0 * 1024.0 * 1024.0);

        let buf = AlignedBuffer::try_new(4096, 4096);
        assert!(buf.is_ok(), "4KB buffer allocation should succeed");
        assert!(buf.unwrap().capacity() >= 4096);

        GLOBAL_BUFFER_LIMIT.set(saved);
    }
}