rzmq 0.5.18

#![cfg(feature = "io-uring")]

use crate::ZmqError;
use bytes::Bytes;
use io_uring::IoUring;
use libc;
use parking_lot::Mutex;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tracing::{error, info, trace, warn};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct RegisteredSendBufferId(pub(crate) u16);

#[derive(Debug)]
struct SendBufferSlot {
  id: RegisteredSendBufferId,
  buffer: Vec<u8>,     // Owns the memory for the buffer
  in_kernel_use: bool, // True if SEND_ZC submitted, awaiting F_NOTIFY
}

impl SendBufferSlot {
  fn ptr(&self) -> *const u8 {
    self.buffer.as_ptr()
  }
  fn capacity(&self) -> usize {
    self.buffer.capacity()
  }
  fn as_mut_slice(&mut self) -> &mut [u8] {
    &mut self.buffer
  }
  fn iovec(&self) -> libc::iovec {
    libc::iovec {
      iov_base: self.buffer.as_ptr() as *mut libc::c_void,
      iov_len: self.buffer.len(), // For registration, use the full allocated length
    }
  }
}

#[derive(Debug)]
struct SendBufferPoolInner {
  pool: Vec<SendBufferSlot>,
  free_ids: VecDeque<RegisteredSendBufferId>,
}

#[derive(Debug)]
pub(crate) struct SendBufferPool {
  inner: Mutex<SendBufferPoolInner>,
}

/// A leased slot from the `SendBufferPool`, giving the holder exclusive write access to a
/// pre-registered kernel buffer. On drop, the slot is automatically recycled unless
/// `released_to_worker` has been set (indicating the worker took ownership).
pub(crate) struct SendBufferLease {
  pub id: RegisteredSendBufferId,
  /// Pointer into the registered buffer (stable for pool lifetime).
  pub ptr: *mut u8,
  pub capacity: usize,
  pub released_to_worker: Arc<AtomicBool>,
  pub pool: Arc<SendBufferPool>,
}

// SAFETY: ptr points into a pool buffer that is pinned for the worker's lifetime.
// Only this lease holder writes to it until released_to_worker is set.
unsafe impl Send for SendBufferLease {}
unsafe impl Sync for SendBufferLease {}

impl Drop for SendBufferLease {
  fn drop(&mut self) {
    if !self.released_to_worker.load(Ordering::Acquire) {
      trace!("SendBufferLease {:?}: dropped before reaching worker — recycling.", self.id);
      self.pool.release_buffer(self.id);
    }
  }
}

impl std::fmt::Debug for SendBufferLease {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.debug_struct("SendBufferLease")
      .field("id", &self.id)
      .field("capacity", &self.capacity)
      .field("released", &self.released_to_worker.load(Ordering::Relaxed))
      .finish_non_exhaustive()
  }
}

impl SendBufferPool {
  pub fn new(ring: &IoUring, count: usize, capacity_per_buffer: usize) -> Result<Self, ZmqError> {
    if count == 0 || capacity_per_buffer == 0 {
      warn!("SendBufferPool initialized with zero count or capacity. Zero-copy send will be effectively disabled.");
      return Ok(Self {
        inner: Mutex::new(SendBufferPoolInner {
          pool: Vec::new(),
          free_ids: VecDeque::new(),
        }),
      });
    }

    // TODO: Consider integrating RLIMIT_MEMLOCK check here or ensuring UringWorker does it.
    // If count * capacity_per_buffer > memlock_limit, adjust or error.
    // For now, assume parameters are valid.

    let mut slots = Vec::with_capacity(count);
    let mut free_ids = VecDeque::with_capacity(count);

    for i in 0..count {
      slots.push(SendBufferSlot {
        id: RegisteredSendBufferId(i as u16),
        buffer: vec![0u8; capacity_per_buffer],
        in_kernel_use: false,
      });
      free_ids.push_back(RegisteredSendBufferId(i as u16));
    }

    let iovecs_to_register: Vec<libc::iovec> = slots.iter().map(|slot| slot.iovec()).collect();

    unsafe {
      ring.submitter().register_buffers(&iovecs_to_register).map_err(|e| {
        error!("SendBufferPool: Failed to register_buffers: {}", e);
        ZmqError::Internal(format!("SendBufferPool: Failed to register_buffers: {}", e))
      })?;
    }

    info!(
      "SendBufferPool: Registered {} send buffers ({} bytes each).",
      count, capacity_per_buffer
    );
    Ok(Self {
      inner: Mutex::new(SendBufferPoolInner { pool: slots, free_ids }),
    })
  }

  /// Attempts to acquire a free buffer, copies `data_to_copy` into it,
  /// and marks it as in kernel use.
  /// Returns the buffer's ID, a pointer to its data, and the length of data copied.
  pub fn acquire_and_prep_buffer(&self, data_to_copy: &Bytes) -> Option<(RegisteredSendBufferId, *const u8, u32)> {
    if data_to_copy.is_empty() {
      trace!("SendBufferPool: acquire_and_prep_buffer called with empty data, skipping.");
      return None; // Cannot SEND_ZC empty data
    }

    let mut inner_guard = self.inner.lock();
    if inner_guard.pool.is_empty() {
      trace!("SendBufferPool: No buffers configured in the pool.");
      return None;
    }

    if let Some(buffer_id) = inner_guard.free_ids.pop_front() {
      let slot = &mut inner_guard.pool[buffer_id.0 as usize];

      if data_to_copy.len() > slot.capacity() {
        warn!(
          "Data ({} bytes) too large for send buffer slot {:?} (capacity {} bytes). Cannot use zero-copy for this send.",
          data_to_copy.len(),
          buffer_id,
          slot.capacity()
        );
        inner_guard.free_ids.push_front(buffer_id); // Put it back, it's still free
        return None;
      }

      // Copy data into the registered buffer
      slot.as_mut_slice()[..data_to_copy.len()].copy_from_slice(data_to_copy);
      slot.in_kernel_use = true; // Mark as given to kernel

      trace!(
        "SendBufferPool: Acquired buffer {:?} for {} bytes.",
        buffer_id,
        data_to_copy.len()
      );
      Some((buffer_id, slot.ptr(), data_to_copy.len() as u32))
    } else {
      trace!("SendBufferPool: No free send buffers available.");
      None // No free buffers
    }
  }

  /// Leases a free buffer slot, giving the caller exclusive write access without copying.
  /// The returned `SendBufferLease` auto-recycles on drop if not handed to the worker.
  pub fn acquire_lease(self: &Arc<Self>) -> Option<SendBufferLease> {
    let mut inner = self.inner.lock();
    if let Some(id) = inner.free_ids.pop_front() {
      let slot = &mut inner.pool[id.0 as usize];
      slot.in_kernel_use = true;
      Some(SendBufferLease {
        id,
        ptr: slot.as_mut_slice().as_mut_ptr(),
        capacity: slot.capacity(),
        released_to_worker: Arc::new(AtomicBool::new(false)),
        pool: self.clone(),
      })
    } else {
      None
    }
  }

  /// Releases a buffer, making it available for reuse.
  /// Called after the kernel signals it's done (e.g., via SEND_ZC F_NOTIFY CQE).
  pub fn release_buffer(&self, id: RegisteredSendBufferId) {
    let mut inner_guard = self.inner.lock();
    if let Some(slot_index) = inner_guard.pool.iter().position(|s| s.id == id) {
      let slot = &mut inner_guard.pool[slot_index];
      if slot.in_kernel_use {
        slot.in_kernel_use = false;
        // Check if it's already in free_ids to prevent duplicates, though ideally it shouldn't be.
        if !inner_guard.free_ids.contains(&id) {
          inner_guard.free_ids.push_back(id);
        } else {
          warn!("SendBufferPool: Buffer ID {:?} was already in free_ids when trying to release (in_kernel_use was true). State may be inconsistent.", id);
        }
        trace!("SendBufferPool: Released buffer {:?}.", id);
      } else {
        warn!(
          "SendBufferPool: Attempted to release buffer {:?} which was not marked as in_kernel_use.",
          id
        );
        // If it wasn't marked in_kernel_use, it should ideally already be in free_ids or never taken.
        // Adding it defensively if not present.
        if !inner_guard.free_ids.contains(&id) {
          inner_guard.free_ids.push_back(id);
        }
      }
    } else {
      error!("SendBufferPool: Attempted to release an unknown buffer ID: {:?}", id);
    }
  }

  /// Unregisters all buffers from io_uring. Called on UringWorker shutdown.
  pub unsafe fn unregister_all(&self, ring: &IoUring) -> Result<(), ZmqError> {
    let inner_guard = self.inner.lock(); // Ensure no modifications during unregister
    if inner_guard.pool.is_empty() {
      return Ok(());
    }
    // It's important that no operations are pending on these buffers when unregistering.
    // The UringWorker should ensure all SEND_ZC ops are complete (notifications received).
    drop(inner_guard); // Release lock

    info!("SendBufferPool: Unregistering all send buffers.");
    ring.submitter().unregister_buffers().map_err(|e| {
      error!("SendBufferPool: Failed to unregister_buffers: {}", e);
      ZmqError::Internal(format!("SendBufferPool: Failed to unregister_buffers: {}", e))
    })
  }
}

#[cfg(all(test, feature = "io-uring"))]
mod pool_tests {
  use super::*;
  use bytes::Bytes;
  use io_uring::IoUring;

  fn create_test_ring_and_pool(count: usize, cap: usize) -> (IoUring, SendBufferPool) {
    let ring = IoUring::new(16).expect("Failed to initialize test io_uring");
    let pool =
      SendBufferPool::new(&ring, count, cap).expect("Failed to initialize SendBufferPool");
    (ring, pool)
  }

  #[test]
  fn test_out_of_order_recycling() {
    let (_ring, pool) = create_test_ring_and_pool(5, 128);
    let mut acquired = Vec::new();

    for i in 0..5usize {
      let data = Bytes::from(format!("payload-data-for-buffer-{}", i));
      let lease = pool
        .acquire_and_prep_buffer(&data)
        .expect("Should acquire successfully");
      acquired.push(lease);
    }

    // Release in scrambled order.
    for idx in [3usize, 1, 4, 0, 2] {
      let (id, _, _) = acquired[idx];
      pool.release_buffer(id);
    }

    // All slots should be recyclable — no corruption.
    for i in 0..5usize {
      let data = Bytes::from(format!("new-payload-{}", i));
      let lease = pool.acquire_and_prep_buffer(&data);
      assert!(lease.is_some(), "Buffer {} should be recycled and acquirable", i);
    }
  }

  #[test]
  fn test_pool_exhaustion_and_recovery() {
    let (_ring, pool) = create_test_ring_and_pool(2, 64);

    let data = Bytes::from_static(b"test-bytes");
    let lease1 = pool.acquire_and_prep_buffer(&data);
    let lease2 = pool.acquire_and_prep_buffer(&data);

    assert!(lease1.is_some());
    assert!(lease2.is_some());

    // 3rd acquisition must fail — pool is exhausted.
    let lease3 = pool.acquire_and_prep_buffer(&data);
    assert!(lease3.is_none(), "Pool exhaustion should return None");

    // Release one slot, then verify recovery.
    let (id1, _, _) = lease1.unwrap();
    pool.release_buffer(id1);

    let lease_recovered = pool.acquire_and_prep_buffer(&data);
    assert!(lease_recovered.is_some(), "Should recover one buffer from pool");
    assert_eq!(
      lease_recovered.unwrap().0,
      id1,
      "Should recycle the exact returned buffer id"
    );
  }
}