khata_rs/

queue.rs

//! khata-rs Queue implementation.

use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use parking_lot::RwLock;

use crate::cycle::{FastDaily, RollCycle};
use crate::error::{Error, Result};
use crate::reader::Reader;
use crate::store::{Store, StoreConfig};
use crate::writer::Writer;

/// File extension for Chronicle Queue data files.
pub const FILE_EXTENSION: &str = "cq4";

/// Builder for creating a [`Queue`].
///
/// # Example
///
/// ```no_run
/// use khata_rs::{QueueBuilder, cycle::FastHourly};
///
/// let queue = QueueBuilder::new("/tmp/my-queue")
///     .roll_cycle(FastHourly)
///     .block_size(32 * 1024 * 1024)
///     .build()?;
/// # Ok::<(), khata_rs::Error>(())
/// ```
pub struct QueueBuilder {
    path: PathBuf,
    roll_cycle: Arc<dyn RollCycle>,
    block_size: usize,
    index_spacing: u32,
    index_count: u32,
    epoch: u64,
    read_only: bool,
    enable_checksums: bool,
}

impl QueueBuilder {
    /// Creates a new builder for a queue at the given path.
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: path.into(),
            roll_cycle: Arc::new(FastDaily),
            block_size: 64 * 1024 * 1024,
            index_spacing: 256,
            index_count: 4096,
            epoch: 0,
            read_only: false,
            enable_checksums: false,
        }
    }

    /// Sets the roll cycle.
    pub fn roll_cycle<R: RollCycle + 'static>(mut self, cycle: R) -> Self {
        self.roll_cycle = Arc::new(cycle);
        self
    }

    /// Sets the block size for memory-mapped files (default: 64 MB).
    pub fn block_size(mut self, size: usize) -> Self {
        self.block_size = size;
        self
    }

    /// Sets the index spacing (default: 256)
    /// Has to be a power of two for bitwise masks
    pub fn index_spacing(mut self, spacing: u32) -> Self {
        self.index_spacing = spacing.next_power_of_two();
        self
    }

    /// Sets the index count per array (default: 4096).
    /// Has to be a power of two for bitwise masks.
    pub fn index_count(mut self, count: u32) -> Self {
        self.index_count = count.next_power_of_two();
        self
    }

    /// Sets the epoch offset in milliseconds (default: 0).
    pub fn epoch(mut self, epoch: u64) -> Self {
        self.epoch = epoch;
        self
    }

    /// Opens in read-only mode.
    pub fn read_only(mut self, read_only: bool) -> Self {
        self.read_only = read_only;
        self
    }

    /// Enables or disables CRC-16 checksums for data integrity (default: true).
    ///
    /// When enabled, each message is written with a CRC-16 checksum that can
    /// be verified on read. Disabling checksums reduces latency but removes
    /// data integrity verification.
    pub fn checksums(mut self, enable: bool) -> Self {
        self.enable_checksums = enable;
        self
    }

    /// Builds the queue.
    pub fn build(self) -> Result<Queue> {
        Queue::open(self)
    }
}

/// A persisted, low-latency message queue.
///
/// Uses memory-mapped files organized into cycle files based on time.
///
/// # Thread Safety
///
/// - One [`Writer`](crate::Writer) per thread
/// - Multiple [`Reader`]s can read concurrently
/// - The queue itself is `Send + Sync`
///
/// # Example
///
/// ```no_run
/// use khata_rs::Queue;
///
/// let queue = Queue::new("/tmp/my-queue").build()?;
///
/// let mut writer = queue.writer()?;
/// writer.write(b"Hello, World!")?;
/// # Ok::<(), khata_rs::Error>(())
/// ```
pub struct Queue {
    path: PathBuf,
    roll_cycle: Arc<dyn RollCycle>,
    block_size: usize,
    index_spacing: u32,
    index_count: u32,
    epoch: u64,
    read_only: bool,
    enable_checksums: bool,
    store_cache: RwLock<HashMap<i32, Arc<Store>>>,
    directory_listing: RwLock<Vec<i32>>,
}

impl Queue {
    /// Creates a new builder for a queue at the given path.
    #[must_use]
    pub fn new(path: impl Into<PathBuf>) -> QueueBuilder {
        QueueBuilder::new(path)
    }

    /// Opens or creates a queue from builder config.
    fn open(builder: QueueBuilder) -> Result<Self> {
        if !builder.path.exists() {
            if builder.read_only {
                return Err(Error::StoreFileMissing(builder.path));
            }
            fs::create_dir_all(&builder.path)
                .map_err(|_| Error::DirectoryCreationFailed(builder.path.clone()))?;
        }

        let queue = Self {
            path: builder.path,
            roll_cycle: builder.roll_cycle,
            block_size: builder.block_size,
            index_spacing: builder.index_spacing,
            index_count: builder.index_count,
            epoch: builder.epoch,
            read_only: builder.read_only,
            enable_checksums: builder.enable_checksums,
            store_cache: RwLock::new(HashMap::new()),
            directory_listing: RwLock::new(Vec::new()),
        };

        queue.refresh_directory_listing()?;
        Ok(queue)
    }

    /// Creates a writer for appending messages.
    ///
    /// # Errors
    ///
    /// Returns `Error::ReadOnly` if opened in read-only mode.
    pub fn writer(&self) -> Result<Writer<'_>> {
        if self.read_only {
            return Err(Error::ReadOnly);
        }
        Writer::new(self)
    }

    /// Creates a reader for consuming messages.
    pub fn reader(&self) -> Result<Reader<'_>> {
        Reader::new(self)
    }

    /// Returns the queue directory path.
    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Returns the roll cycle configuration.
    #[must_use]
    pub fn roll_cycle(&self) -> &dyn RollCycle {
        self.roll_cycle.as_ref()
    }

    /// Returns the block size.
    #[must_use]
    pub fn block_size(&self) -> usize {
        self.block_size
    }

    /// Returns the epoch offset.
    #[must_use]
    pub fn epoch(&self) -> u64 {
        self.epoch
    }

    /// Returns whether the queue is read-only.
    #[must_use]
    pub fn is_read_only(&self) -> bool {
        self.read_only
    }

    /// Returns whether CRC-16 checksums are enabled.
    #[must_use]
    pub fn checksums_enabled(&self) -> bool {
        self.enable_checksums
    }

    /// Returns the index spacing.
    #[must_use]
    pub fn index_spacing(&self) -> u32 {
        self.index_spacing
    }

    /// Returns the index count per array.
    #[must_use]
    pub fn index_count(&self) -> u32 {
        self.index_count
    }

    /// Returns the first (oldest) cycle, if any.
    #[must_use]
    pub fn first_cycle(&self) -> Option<i32> {
        self.directory_listing.read().first().copied()
    }

    /// Returns the last (newest) cycle, if any.
    #[must_use]
    pub fn last_cycle(&self) -> Option<i32> {
        self.directory_listing.read().last().copied()
    }

    /// Returns the current cycle based on current time.
    #[must_use]
    pub fn current_cycle(&self) -> i32 {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_millis() as u64)
            .unwrap_or(0);
        self.roll_cycle.current_cycle(now, self.epoch)
    }

    /// Acquires a store for the given cycle.
    pub(crate) fn acquire_store(&self, cycle: i32) -> Result<Arc<Store>> {
        {
            let cache = self.store_cache.read();
            if let Some(store) = cache.get(&cycle) {
                return Ok(Arc::clone(store));
            }
        }

        let store = Arc::new(self.create_store(cycle)?);

        {
            let mut cache = self.store_cache.write();
            cache.insert(cycle, Arc::clone(&store));
        }

        {
            let mut listing = self.directory_listing.write();
            if !listing.contains(&cycle) {
                listing.push(cycle);
                listing.sort_unstable();
            }
        }

        Ok(store)
    }

    /// Creates a new store for the given cycle.
    fn create_store(&self, cycle: i32) -> Result<Store> {
        let filename = format!(
            "{}.{}",
            self.roll_cycle.format_cycle(cycle, self.epoch),
            FILE_EXTENSION
        );
        let config = StoreConfig {
            block_size: self.block_size,
            enable_checksums: self.enable_checksums,
            ..StoreConfig::default()
        };
        Store::open(self.path.join(filename), cycle, config, self.read_only)
    }

    /// Generates the filename for a given cycle.
    #[must_use]
    pub fn filename_for_cycle(&self, cycle: i32) -> String {
        format!(
            "{}.{}",
            self.roll_cycle.format_cycle(cycle, self.epoch),
            FILE_EXTENSION
        )
    }

    /// Refreshes the directory listing.
    fn refresh_directory_listing(&self) -> Result<()> {
        let mut cycles = Vec::new();

        if let Ok(entries) = fs::read_dir(&self.path) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.extension().is_some_and(|ext| ext == FILE_EXTENSION) {
                    if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
                        if let Some(cycle) = self.roll_cycle.parse_cycle(stem, self.epoch) {
                            cycles.push(cycle);
                        }
                    }
                }
            }
        }

        cycles.sort_unstable();

        *self.directory_listing.write() = cycles;
        Ok(())
    }
}

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

    #[test]
    fn test_builder_defaults() {
        let builder = QueueBuilder::new("/tmp/test");
        assert_eq!(builder.block_size, 64 * 1024 * 1024);
        assert_eq!(builder.index_spacing, 256);
        assert_eq!(builder.index_count, 4096);
        assert_eq!(builder.epoch, 0);
        assert!(!builder.read_only);
        assert!(!builder.enable_checksums);
    }

    #[test]
    fn test_checksums_enabled() {
        let temp_dir = TempDir::new().unwrap();
        let queue = Queue::new(temp_dir.path()).checksums(true).build().unwrap();

        assert!(queue.checksums_enabled());
    }

    #[test]
    fn test_creation() {
        let temp_dir = TempDir::new().unwrap();
        let queue = Queue::new(temp_dir.path()).build().unwrap();

        assert!(queue.path().exists());
        assert!(!queue.is_read_only());
    }

    #[test]
    fn test_filename_for_cycle() {
        let temp_dir = TempDir::new().unwrap();
        let queue = Queue::new(temp_dir.path()).build().unwrap();

        assert_eq!(queue.filename_for_cycle(0), "19700101.cq4");
    }

    #[test]
    fn test_single_producer_multiple_consumer() {
        let temp_dir = TempDir::new().unwrap();
        let queue = Queue::new(temp_dir.path()).build().unwrap();

        let mut writer = queue.writer().unwrap();

        // Write multiple messages
        let messages: Vec<Vec<u8>> = (0..10)
            .map(|i| format!("message-{}", i).into_bytes())
            .collect();

        for msg in &messages {
            writer.write(msg).unwrap();
        }

        // Create two independent readers
        let mut r1 = queue.reader().unwrap();
        r1.rewind().unwrap();

        let mut r2 = queue.reader().unwrap();
        r2.rewind().unwrap();

        // Both readers should independently read all messages
        for expected in &messages {
            let read1 = r1.read().unwrap();
            assert!(read1.is_some(), "r1 should have data");
            assert_eq!(read1.unwrap().as_ref(), expected.as_slice());

            let read2 = r2.read().unwrap();
            assert!(read2.is_some(), "r2 should have data");
            assert_eq!(read2.unwrap().as_ref(), expected.as_slice());
        }

        // Both readers should now be at the end
        assert!(r1.read().unwrap().is_none());
        assert!(r2.read().unwrap().is_none());
    }
}