grafos_collections/

queue.rs

//! A bounded SPSC ring buffer stored in leased fabric memory.
//!
//! [`FabricQueue<T>`] implements a single-producer single-consumer FIFO queue
//! backed by a leased memory region. Elements are serialized with postcard
//! into fixed-size slots arranged in a circular buffer.
//!
//! # Memory layout
//!
//! ```text
//! Offset 0:  [header]  head: u64 (8) | tail: u64 (8) | capacity: u64 (8) | stride: u64 (8)
//! Offset 32: [slots]   slot 0 (stride bytes) | slot 1 | ... | slot (capacity-1)
//! ```
//!
//! `head` is the next index to read from, `tail` is the next index to write
//! to. The queue is empty when `head == tail` and full when
//! `(tail + 1) % capacity == head`. One slot is always reserved for the
//! empty/full distinction, so the usable capacity is `capacity - 1`.
//!
//! # Example
//!
//! ```rust
//! use grafos_collections::queue::FabricQueue;
//! use grafos_std::mem::MemBuilder;
//!
//! # grafos_std::host::reset_mock();
//! # grafos_std::host::mock_set_fbmu_arena_size(65536);
//! let lease = MemBuilder::new().min_bytes(4096).acquire()?;
//! let mut q: FabricQueue<u32> = FabricQueue::new(lease, 16, 16)?;
//!
//! q.push(&1)?;
//! q.push(&2)?;
//! assert_eq!(q.pop()?, Some(1)); // FIFO order
//! # Ok::<(), grafos_std::FabricError>(())
//! ```

extern crate alloc;
use alloc::vec;

use grafos_std::error::{FabricError, Result};
use grafos_std::mem::{MemBuilder, MemLease};

use serde::{de::DeserializeOwned, Serialize};

const HEADER_SIZE: u64 = 32;

/// A bounded SPSC ring buffer stored in leased fabric memory.
///
/// The queue owns its [`MemLease`] and releases it on drop. Elements are
/// stored in FIFO order: `push()` appends to the tail, `pop()` reads from
/// the head.
///
/// # Non-blocking semantics
///
/// - `push()` returns `Ok(false)` when the queue is full.
/// - `pop()` returns `Ok(None)` when the queue is empty.
///
/// Neither operation blocks or retries.
///
/// # Future direction: MPMC
///
/// This queue is currently single-producer single-consumer. A future MPMC
/// variant would replace the plain `head`/`tail` fields with
/// compare-and-swap (CAS) atomic operations on the remote memory, or
/// alternatively use a lock-based protocol with a lease-scoped mutex in
/// the header region. MPMC support is not yet implemented.
///
/// # Example
///
/// ```rust
/// use grafos_collections::queue::FabricQueue;
///
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// let mut q: FabricQueue<u32> = FabricQueue::with_capacity(8, 16)?;
/// q.push(&10)?;
/// q.push(&20)?;
/// assert_eq!(q.pop()?, Some(10));
/// assert_eq!(q.pop()?, Some(20));
/// assert_eq!(q.pop()?, None);
/// # Ok::<(), grafos_std::FabricError>(())
/// ```
pub struct FabricQueue<T> {
    lease: MemLease,
    head: u64,
    tail: u64,
    capacity: u64,
    stride: u64,
    _marker: core::marker::PhantomData<T>,
}

impl<T: Serialize + DeserializeOwned> FabricQueue<T> {
    /// Create a new queue over the given lease.
    ///
    /// `capacity` is the total number of slots. `stride` is the slot width
    /// in bytes (must accommodate 4-byte length prefix plus serialized
    /// element). One slot is reserved for the empty/full distinction, so
    /// the usable capacity is `capacity - 1`.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the arena is too small
    /// to hold the header (32 bytes) plus `capacity * stride` bytes.
    pub fn new(lease: MemLease, capacity: usize, stride: usize) -> Result<Self> {
        let arena = lease.mem().arena_size()?;
        let needed = HEADER_SIZE + (capacity as u64) * (stride as u64);
        if arena < needed {
            return Err(FabricError::CapacityExceeded);
        }
        let q = FabricQueue {
            lease,
            head: 0,
            tail: 0,
            capacity: capacity as u64,
            stride: stride as u64,
            _marker: core::marker::PhantomData,
        };
        q.write_header()?;
        Ok(q)
    }

    /// Create a new queue by acquiring a lease sized for `capacity` elements
    /// of `stride` bytes each.
    ///
    /// Convenience constructor that acquires a lease via `MemBuilder` and
    /// then calls [`FabricQueue::new`].
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the host cannot provide
    /// a large enough arena.
    pub fn with_capacity(capacity: usize, stride: usize) -> Result<Self> {
        let needed = HEADER_SIZE + (capacity as u64) * (stride as u64);
        let lease = MemBuilder::new().min_bytes(needed).acquire()?;
        Self::new(lease, capacity, stride)
    }

    /// Push an element onto the back of the queue.
    ///
    /// Returns `Ok(true)` if the element was enqueued, `Ok(false)` if the
    /// queue is full. Does not block.
    ///
    /// # Errors
    ///
    /// - [`FabricError::CapacityExceeded`] if the serialized element
    ///   (plus 4-byte length prefix) exceeds `stride`.
    /// - [`FabricError::IoError`] if serialization or remote write fails.
    pub fn push(&mut self, item: &T) -> Result<bool> {
        if self.is_full() {
            return Ok(false);
        }
        let bytes = postcard::to_allocvec(item).map_err(|_| FabricError::IoError(-1))?;
        if bytes.len() as u64 + 4 > self.stride {
            return Err(FabricError::CapacityExceeded);
        }
        let offset = self.slot_offset(self.tail);
        let mut slot = vec![0u8; self.stride as usize];
        let len_bytes = (bytes.len() as u32).to_le_bytes();
        slot[..4].copy_from_slice(&len_bytes);
        slot[4..4 + bytes.len()].copy_from_slice(&bytes);
        self.lease.mem().write(offset, &slot)?;
        self.tail = (self.tail + 1) % self.capacity;
        self.write_header()?;
        Ok(true)
    }

    /// Pop an element from the front of the queue.
    ///
    /// Returns `Ok(Some(item))` if an element was dequeued, `Ok(None)` if
    /// the queue is empty. Does not block.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::IoError`] if the remote read or
    /// deserialization fails.
    pub fn pop(&mut self) -> Result<Option<T>> {
        if self.is_empty() {
            return Ok(None);
        }
        let offset = self.slot_offset(self.head);
        let raw = self.lease.mem().read(offset, self.stride as u32)?;
        let ser_len = u32::from_le_bytes([raw[0], raw[1], raw[2], raw[3]]) as usize;
        let item: T =
            postcard::from_bytes(&raw[4..4 + ser_len]).map_err(|_| FabricError::IoError(-1))?;
        self.head = (self.head + 1) % self.capacity;
        self.write_header()?;
        Ok(Some(item))
    }

    /// Returns the number of elements currently in the queue.
    pub fn len(&self) -> usize {
        if self.tail >= self.head {
            (self.tail - self.head) as usize
        } else {
            (self.capacity - self.head + self.tail) as usize
        }
    }

    /// Returns `true` if the queue contains no elements.
    pub fn is_empty(&self) -> bool {
        self.head == self.tail
    }

    /// Returns `true` if the queue is full.
    pub fn is_full(&self) -> bool {
        (self.tail + 1) % self.capacity == self.head
    }

    /// Returns the lease ID of the memory lease for external renewal
    /// management (e.g. via [`grafos_leasekit::RenewalManager`]).
    pub fn lease_id(&self) -> u128 {
        self.lease.lease_id()
    }

    /// Returns the expiry time (unix seconds) of the memory lease for
    /// external renewal management.
    pub fn expires_at_unix_secs(&self) -> u64 {
        self.lease.expires_at_unix_secs()
    }

    /// Returns the maximum number of elements the queue can hold.
    ///
    /// This is `capacity - 1` since one slot is reserved for the
    /// empty/full distinction.
    pub fn max_len(&self) -> usize {
        (self.capacity - 1) as usize
    }

    fn slot_offset(&self, index: u64) -> u64 {
        HEADER_SIZE + index * self.stride
    }

    fn write_header(&self) -> Result<()> {
        let mut hdr = [0u8; HEADER_SIZE as usize];
        hdr[0..8].copy_from_slice(&self.head.to_le_bytes());
        hdr[8..16].copy_from_slice(&self.tail.to_le_bytes());
        hdr[16..24].copy_from_slice(&self.capacity.to_le_bytes());
        hdr[24..32].copy_from_slice(&self.stride.to_le_bytes());
        self.lease.mem().write(0, &hdr)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use grafos_std::host;
    use grafos_std::mem::MemBuilder;

    fn setup(arena_size: u64) -> MemLease {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(arena_size);
        MemBuilder::new().acquire().expect("acquire")
    }

    #[test]
    fn push_pop_fifo() {
        let lease = setup(4096);
        let mut q: FabricQueue<u32> = FabricQueue::new(lease, 16, 16).expect("new");

        assert!(q.is_empty());
        assert!(!q.is_full());
        assert_eq!(q.len(), 0);

        q.push(&1).expect("push");
        q.push(&2).expect("push");
        q.push(&3).expect("push");
        assert_eq!(q.len(), 3);

        assert_eq!(q.pop().expect("pop"), Some(1));
        assert_eq!(q.pop().expect("pop"), Some(2));
        assert_eq!(q.pop().expect("pop"), Some(3));
        assert_eq!(q.pop().expect("pop"), None);
    }

    #[test]
    fn full_returns_false() {
        // capacity=4, stride=16, so 3 usable slots
        let lease = setup(HEADER_SIZE + 4 * 16);
        let mut q: FabricQueue<u32> = FabricQueue::new(lease, 4, 16).expect("new");

        assert!(q.push(&1).expect("push 1"));
        assert!(q.push(&2).expect("push 2"));
        assert!(q.push(&3).expect("push 3"));
        assert!(q.is_full());
        assert!(!q.push(&4).expect("push 4 (full)"));
    }

    #[test]
    fn wraparound() {
        let lease = setup(HEADER_SIZE + 4 * 16);
        let mut q: FabricQueue<u32> = FabricQueue::new(lease, 4, 16).expect("new");

        // Fill and drain a few times to exercise wrap-around
        for round in 0..3u32 {
            let base = round * 100;
            q.push(&(base + 1)).expect("push");
            q.push(&(base + 2)).expect("push");
            q.push(&(base + 3)).expect("push");
            assert!(q.is_full());

            assert_eq!(q.pop().expect("pop"), Some(base + 1));
            assert_eq!(q.pop().expect("pop"), Some(base + 2));
            assert_eq!(q.pop().expect("pop"), Some(base + 3));
            assert!(q.is_empty());
        }
    }

    #[test]
    fn with_capacity_works() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);
        let q: FabricQueue<u64> = FabricQueue::with_capacity(32, 16).expect("with_capacity");
        assert_eq!(q.max_len(), 31);
    }

    #[test]
    fn interleaved_push_pop() {
        let lease = setup(4096);
        let mut q: FabricQueue<u32> = FabricQueue::new(lease, 8, 16).expect("new");

        q.push(&10).expect("push");
        q.push(&20).expect("push");
        assert_eq!(q.pop().expect("pop"), Some(10));
        q.push(&30).expect("push");
        assert_eq!(q.pop().expect("pop"), Some(20));
        assert_eq!(q.pop().expect("pop"), Some(30));
        assert_eq!(q.pop().expect("pop"), None);
    }
}