grafos_mq/

partition.rs

//! Partition: ring buffer in leased memory with slot-based indexing.
//!
//! Each partition is a bounded ring buffer backed by a [`MemLease`]. Messages
//! are serialized into fixed-size slots. When the buffer wraps, old messages
//! are overwritten (log compaction by offset).
//!
//! # Memory layout
//!
//! ```text
//! Offset 0:  [header]  write_offset: u64 (8) | capacity: u64 (8) | stride: u64 (8) | msg_count: u64 (8)
//! Offset 32: [slots]   slot 0 (stride bytes) | slot 1 | ... | slot (capacity-1)
//! ```

extern crate alloc;
use alloc::vec;

use crate::message::Message;
use grafos_std::error::{FabricError, Result};
use grafos_std::host;
use grafos_std::mem::{MemBuilder, MemLease};

const HEADER_SIZE: u64 = 32;

/// A single partition backed by a memory lease.
pub struct Partition {
    lease: MemLease,
    /// Next offset to write at (monotonically increasing, mod capacity for slot index).
    write_offset: u64,
    /// Number of slots in the ring buffer.
    capacity: u64,
    /// Size of each slot in bytes.
    stride: u64,
    /// Total number of messages ever written (for offset tracking).
    msg_count: u64,
}

impl Partition {
    /// Create a new partition with the given capacity and slot stride.
    ///
    /// Acquires a memory lease large enough for the header plus all slots.
    pub fn new(capacity: usize, stride: usize) -> Result<Self> {
        let needed = HEADER_SIZE + (capacity as u64) * (stride as u64);
        let lease = MemBuilder::new().min_bytes(needed).acquire()?;
        let p = Partition {
            lease,
            write_offset: 0,
            capacity: capacity as u64,
            stride: stride as u64,
            msg_count: 0,
        };
        p.write_header()?;
        Ok(p)
    }

    /// Append a message to the partition.
    ///
    /// The message's `offset` and `timestamp` fields are set automatically.
    /// Returns the assigned offset, or `CapacityExceeded` if the serialized
    /// message exceeds the slot stride.
    pub fn append(&mut self, key: Option<&[u8]>, value: &[u8]) -> Result<u64> {
        let now = host::unix_time_secs();
        let offset = self.msg_count;
        let msg = Message {
            offset,
            timestamp: now,
            key: key.map(|k| k.to_vec()),
            value: value.to_vec(),
            headers: alloc::vec::Vec::new(),
        };

        let bytes = postcard::to_allocvec(&msg).map_err(|_| FabricError::IoError(-1))?;
        if bytes.len() as u64 + 4 > self.stride {
            return Err(FabricError::CapacityExceeded);
        }

        let slot_idx = self.write_offset % self.capacity;
        let slot_offset = HEADER_SIZE + slot_idx * self.stride;

        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(slot_offset, &slot)?;
        self.write_offset += 1;
        self.msg_count += 1;
        self.write_header()?;

        Ok(offset)
    }

    /// Read the message at the given absolute offset.
    ///
    /// Returns `None` if the offset has been overwritten (wrapped) or is
    /// beyond the current write position.
    pub fn read_at(&self, offset: u64) -> Result<Option<Message>> {
        if offset >= self.msg_count {
            return Ok(None);
        }
        // Check if offset has been overwritten
        let oldest_available = self.msg_count.saturating_sub(self.capacity);
        if offset < oldest_available {
            return Ok(None);
        }

        let slot_idx = offset % self.capacity;
        let slot_offset = HEADER_SIZE + slot_idx * self.stride;

        let raw = self.lease.mem().read(slot_offset, self.stride as u32)?;
        let ser_len = u32::from_le_bytes([raw[0], raw[1], raw[2], raw[3]]) as usize;
        if ser_len == 0 || ser_len + 4 > raw.len() {
            return Ok(None);
        }
        let msg: Message =
            postcard::from_bytes(&raw[4..4 + ser_len]).map_err(|_| FabricError::IoError(-1))?;
        Ok(Some(msg))
    }

    /// Returns the next offset that will be assigned.
    pub fn next_offset(&self) -> u64 {
        self.msg_count
    }

    /// Returns the oldest available offset (messages before this have been overwritten).
    pub fn oldest_offset(&self) -> u64 {
        self.msg_count.saturating_sub(self.capacity)
    }

    /// Returns the number of messages currently retained in the ring buffer.
    pub fn len(&self) -> usize {
        let available = self.msg_count - self.oldest_offset();
        available as usize
    }

    /// Returns `true` if no messages are in the partition.
    pub fn is_empty(&self) -> bool {
        self.msg_count == 0
    }

    /// Returns the slot capacity of this partition.
    pub fn capacity(&self) -> usize {
        self.capacity as usize
    }

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

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

    fn setup() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);
    }

    #[test]
    fn append_and_read() {
        setup();
        let mut part = Partition::new(16, 256).expect("new");
        assert!(part.is_empty());

        let off = part.append(None, b"hello").expect("append");
        assert_eq!(off, 0);
        assert_eq!(part.len(), 1);

        let msg = part.read_at(0).expect("read").expect("some");
        assert_eq!(msg.offset, 0);
        assert_eq!(msg.value, b"hello");
        assert_eq!(msg.key, None);
    }

    #[test]
    fn keyed_append() {
        setup();
        let mut part = Partition::new(16, 256).expect("new");
        part.append(Some(b"k1"), b"v1").expect("append");

        let msg = part.read_at(0).expect("read").expect("some");
        assert_eq!(msg.key, Some(b"k1".to_vec()));
        assert_eq!(msg.value, b"v1");
    }

    #[test]
    fn wraparound() {
        setup();
        let mut part = Partition::new(4, 256).expect("new");

        for i in 0..6u64 {
            part.append(None, &[i as u8; 4]).expect("append");
        }

        // Offsets 0 and 1 should be overwritten
        assert_eq!(part.oldest_offset(), 2);
        assert!(part.read_at(0).expect("read").is_none());
        assert!(part.read_at(1).expect("read").is_none());

        // Offsets 2-5 should be available
        for i in 2..6u64 {
            let msg = part.read_at(i).expect("read").expect("some");
            assert_eq!(msg.offset, i);
        }
    }

    #[test]
    fn read_beyond_write() {
        setup();
        let mut part = Partition::new(8, 256).expect("new");
        part.append(None, b"one").expect("append");
        assert!(part.read_at(1).expect("read").is_none());
        assert!(part.read_at(100).expect("read").is_none());
    }

    #[test]
    fn capacity_exceeded_for_large_message() {
        setup();
        let mut part = Partition::new(4, 32).expect("new");
        // A message with 30+ bytes of payload won't fit in 32 - 4 = 28 bytes of data
        let big = vec![0u8; 200];
        let result = part.append(None, &big);
        assert_eq!(result.unwrap_err(), FabricError::CapacityExceeded);
    }
}