khata_rs/

reader.rs

//! Reader for consuming messages from a persistent queue

use std::sync::Arc;

use crate::error::{Error, Result};
use crate::queue::Queue;
use crate::store::Store;

/// State of a reader in the queue.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReaderState {
    /// Not yet initialized.
    Uninitialized,
    /// Positioned at a valid entry.
    Found,
    /// Reached the end of the current cycle.
    EndOfCycle,
    /// The requested cycle doesn't exist.
    CycleNotFound,
    /// Before the start of a cycle.
    BeforeStart,
    /// Entry not yet written.
    NotReached,
}

/// Reading direction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Direction {
    /// Read forward (default).
    #[default]
    Forward,
    /// Read backward.
    Backward,
    /// Stay at current position.
    None,
}

/// A reader for consuming messages from a [`Queue`].
///
/// Provides sequential and random access to messages, automatically handling
/// cycle boundaries.
///
/// # Thread Safety
///
/// A single `Reader` should only be used by one thread. Create separate
/// instances for concurrent readers.
///
/// # Example
///
/// ```no_run
/// use khata_rs::Queue;
///
/// let queue = Queue::new("/tmp/my-queue").build()?;
/// let mut reader = queue.reader()?;
///
/// reader.rewind()?;
/// while let Some(data) = reader.read()? {
///     println!("Read: {:?}", String::from_utf8_lossy(&data));
/// }
/// # Ok::<(), khata_rs::Error>(())
/// ```
pub struct Reader<'q> {
    queue: &'q Queue,
    cycle: i32,
    index: u64,
    sequence: u64,
    store: Option<Arc<Store>>,
    state: ReaderState,
    direction: Direction,
    /// Cached position of next message for O(1) sequential reads.
    /// Format: position in file (0 = unknown/invalid).
    next_position: u64,
}

impl<'q> Reader<'q> {
    /// Creates a new reader for the given queue.
    pub(crate) fn new(queue: &'q Queue) -> Result<Self> {
        Ok(Self {
            queue,
            cycle: 0,
            index: 0,
            sequence: 0,
            store: None,
            state: ReaderState::Uninitialized,
            direction: Direction::Forward,
            next_position: 0,
        })
    }

    /// Reads the next message (zero-copy).
    ///
    /// Returns a reference to the message data directly in the memory-mapped file.
    /// Returns `None` if there are no more messages. Automatically advances
    /// based on the configured direction.
    ///
    /// # Zero-Copy Performance
    ///
    /// This method returns a slice into the memory-mapped file without copying.
    /// Sequential reads are O(1) due to internal position caching.
    /// The reference is valid as long as the Reader is not dropped or moved to
    /// a different cycle.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use khata_rs::Queue;
    /// # let queue = Queue::new("/tmp/q").build().unwrap();
    /// let mut reader = queue.reader()?;
    /// reader.rewind()?;
    ///
    /// while let Some(data) = reader.read()? {
    ///     // Process data directly without allocation
    ///     println!("Length: {}", data.len());
    /// }
    /// # Ok::<(), khata_rs::Error>(())
    /// ```
    #[inline]
    pub fn read(&mut self) -> Result<Option<&[u8]>> {
        if self.state == ReaderState::Uninitialized {
            self.rewind()?;
        }

        loop {
            // Try fast path with cached position first
            if let Some((ptr, len, next_pos)) = self.read_fast_path()? {
                // Only cache next position if we're actually advancing
                if self.direction == Direction::Forward {
                    self.next_position = next_pos;
                }
                self.advance()?;
                // SAFETY: The slice data lives in the mmap which is owned by
                // self.store (Arc<Store>). advance() only modifies cursor state
                // (sequence/index numbers), not the mmap contents. The mmap
                // remains valid for the lifetime of &self since self holds the
                // Arc<Store>. The borrow checker ensures self is not dropped
                // while this reference is alive.
                let data = unsafe { std::slice::from_raw_parts(ptr, len) };
                return Ok(Some(data));
            }

            if !self.try_next_cycle()? {
                return Ok(None);
            }
        }
    }

    /// Reads the next message with a closure for zero-copy access.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use khata_rs::Queue;
    /// # let queue = Queue::new("/tmp/q").build().unwrap();
    /// let mut reader = queue.reader()?;
    /// reader.rewind()?;
    ///
    /// while let Some(len) = reader.read_with(|data| Ok(data.len()))? {
    ///     println!("Message length: {len}");
    /// }
    /// # Ok::<(), khata_rs::Error>(())
    /// ```
    pub fn read_with<F, T>(&mut self, f: F) -> Result<Option<T>>
    where
        F: FnOnce(&[u8]) -> Result<T>,
    {
        if self.state == ReaderState::Uninitialized {
            self.rewind()?;
        }

        loop {
            if let Some(store) = &self.store {
                if let Some(data) = store.read_ref(self.sequence)? {
                    let result = f(data)?;
                    self.advance()?;
                    return Ok(Some(result));
                }
            }

            if !self.try_next_cycle()? {
                return Ok(None);
            }
        }
    }

    /// Moves to the start of the queue.
    pub fn rewind(&mut self) -> Result<&mut Self> {
        self.next_position = 0; // Reset position cache
        if let Some(first_cycle) = self.queue.first_cycle() {
            self.move_to_cycle(first_cycle)?;
            self.sequence = 0;
            self.update_index();
            self.state = ReaderState::Found;
            // Set initial position to start of data (after file header)
            self.next_position = crate::store::FILE_HEADER_SIZE;
        } else {
            self.cycle = self.queue.current_cycle();
            self.sequence = 0;
            self.store = None;
            self.state = ReaderState::Uninitialized;
        }
        Ok(self)
    }

    /// Moves to the end of the queue.
    ///
    /// Subsequent reads will return new messages as they are written.
    pub fn seek_end(&mut self) -> Result<&mut Self> {
        self.next_position = 0; // Reset position cache
        if let Some(last_cycle) = self.queue.last_cycle() {
            self.move_to_cycle(last_cycle)?;

            if let Some(ref store) = self.store {
                self.sequence = store.message_count();
            }
            self.update_index();
            self.state = ReaderState::NotReached;
        } else {
            self.cycle = self.queue.current_cycle();
            self.sequence = 0;
            self.store = None;
            self.state = ReaderState::Uninitialized;
        }
        Ok(self)
    }

    /// Seeks to a specific index.
    ///
    /// Returns `true` if successful, `false` if the index doesn't exist.
    pub fn seek(&mut self, index: u64) -> Result<bool> {
        self.next_position = 0; // Reset position cache
        let cycle = self.queue.roll_cycle().to_cycle(index);
        let sequence = self.queue.roll_cycle().to_sequence(index);

        if let Err(Error::CycleNotFound(_)) = self.move_to_cycle(cycle) {
            self.state = ReaderState::CycleNotFound;
            return Ok(false);
        }

        if let Some(ref store) = self.store {
            if sequence >= store.message_count() {
                self.state = ReaderState::NotReached;
                return Ok(false);
            }
        }

        self.sequence = sequence;
        self.index = index;
        self.state = ReaderState::Found;
        Ok(true)
    }

    /// Moves to a specific cycle.
    pub fn move_to_cycle(&mut self, cycle: i32) -> Result<&mut Self> {
        if self.cycle == cycle && self.store.is_some() {
            return Ok(self);
        }

        self.next_position = 0; // Reset position cache on cycle change
        match self.queue.acquire_store(cycle) {
            Ok(store) => {
                self.store = Some(store);
                self.cycle = cycle;
                self.sequence = 0;
                self.update_index();
                self.state = ReaderState::Found;
                Ok(self)
            }
            Err(e) => {
                self.state = ReaderState::CycleNotFound;
                Err(e)
            }
        }
    }

    /// Returns the current index.
    #[inline]
    #[must_use]
    pub fn index(&self) -> u64 {
        self.index
    }

    /// Returns the current cycle number.
    #[inline]
    #[must_use]
    pub fn cycle(&self) -> i32 {
        self.cycle
    }

    /// Returns the current sequence within the cycle.
    #[inline]
    #[must_use]
    pub fn sequence(&self) -> u64 {
        self.sequence
    }

    /// Returns the current state.
    #[inline]
    #[must_use]
    pub fn state(&self) -> ReaderState {
        self.state
    }

    /// Returns the current direction.
    #[inline]
    #[must_use]
    pub fn direction(&self) -> Direction {
        self.direction
    }

    /// Sets the reading direction.
    pub fn set_direction(&mut self, direction: Direction) -> &mut Self {
        self.direction = direction;
        self
    }

    /// Peeks at the next message without advancing (zero-copy).
    ///
    /// Returns a reference to the message data directly in the memory-mapped file.
    pub fn peek(&self) -> Result<Option<&[u8]>> {
        if let Some(ref store) = self.store {
            return store.read_ref(self.sequence);
        }
        Ok(None)
    }

    /// Fast path for sequential reads with O(1) position lookup.
    ///
    /// Returns (ptr, len, next_position) where next_position is the file position
    /// of the next message for caching.
    #[inline(always)]
    fn read_fast_path(&self) -> Result<Option<(*const u8, usize, u64)>> {
        let store = match &self.store {
            Some(s) => s,
            None => return Ok(None),
        };

        // Use cached position if available
        if self.next_position > 0 {
            // Fast path: read directly at cached position
            return store.read_at_position_raw(self.next_position);
        }

        // Slow path: first read or after seek, use Store's read mechanism
        if let Some(data) = store.read_ref(self.sequence)? {
            return Ok(Some((data.as_ptr(), data.len(), 0)));
        }

        Ok(None)
    }

    /// Advances the sequence after a successful read.
    fn advance(&mut self) -> Result<()> {
        match self.direction {
            Direction::Forward => {
                self.sequence += 1;
                self.update_index();
            }
            Direction::Backward => {
                if self.sequence > 0 {
                    self.sequence -= 1;
                    self.update_index();
                }
            }
            Direction::None => {}
        }
        Ok(())
    }

    /// Tries to move to the next cycle.
    fn try_next_cycle(&mut self) -> Result<bool> {
        match self.direction {
            Direction::Forward => {
                let next_cycle = self.cycle + 1;

                if let Some(last_cycle) = self.queue.last_cycle() {
                    if next_cycle <= last_cycle && self.move_to_cycle(next_cycle).is_ok() {
                        self.sequence = 0;
                        self.update_index();
                        return Ok(true);
                    }
                }

                self.state = ReaderState::EndOfCycle;
                Ok(false)
            }
            Direction::Backward => {
                let prev_cycle = self.cycle - 1;

                if let Some(first_cycle) = self.queue.first_cycle() {
                    if prev_cycle >= first_cycle && self.move_to_cycle(prev_cycle).is_ok() {
                        if let Some(ref store) = self.store {
                            let count = store.message_count();
                            self.sequence = if count > 0 { count - 1 } else { 0 };
                        }
                        self.update_index();
                        return Ok(true);
                    }
                }

                self.state = ReaderState::BeforeStart;
                Ok(false)
            }
            Direction::None => Ok(false),
        }
    }

    /// Updates the index from cycle and sequence.
    fn update_index(&mut self) {
        self.index = self.queue.roll_cycle().to_index(self.cycle, self.sequence);
    }
}

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

    fn create_queue_with_messages(dir: &TempDir, messages: &[&[u8]]) -> Queue {
        let queue = Queue::new(dir.path()).build().unwrap();
        {
            let mut writer = queue.writer().unwrap();
            for msg in messages {
                writer.write(msg).unwrap();
            }
            writer.flush().unwrap();
        }
        queue
    }

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

        let reader = queue.reader().unwrap();
        assert_eq!(reader.state(), ReaderState::Uninitialized);
    }

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

        let mut reader = queue.reader().unwrap();
        let result = reader.read().unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_single_message() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"Hello"]);

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

        assert_eq!(reader.read().unwrap(), Some(&b"Hello"[..]));
        assert!(reader.read().unwrap().is_none());
    }

    #[test]
    fn test_multiple_messages() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"One", b"Two", b"Three"]);

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

        assert_eq!(reader.read().unwrap(), Some(&b"One"[..]));
        assert_eq!(reader.read().unwrap(), Some(&b"Two"[..]));
        assert_eq!(reader.read().unwrap(), Some(&b"Three"[..]));
        assert!(reader.read().unwrap().is_none());
    }

    #[test]
    fn test_rewind() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"First", b"Second"]);

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

        reader.rewind().unwrap();
        assert_eq!(reader.read().unwrap(), Some(&b"First"[..]));

        reader.rewind().unwrap();
        assert_eq!(reader.read().unwrap(), Some(&b"First"[..]));
    }

    #[test]
    fn test_seek_end() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"One", b"Two"]);

        let mut reader = queue.reader().unwrap();
        reader.seek_end().unwrap();

        assert!(reader.read().unwrap().is_none());
    }

    #[test]
    fn test_read_with() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"Hello, World!"]);

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

        let len = reader.read_with(|data| Ok(data.len())).unwrap().unwrap();
        assert_eq!(len, 13);
    }

    #[test]
    fn test_peek() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"Test"]);

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

        assert_eq!(reader.peek().unwrap(), Some(&b"Test"[..]));
        assert_eq!(reader.peek().unwrap(), Some(&b"Test"[..]));
        assert_eq!(reader.read().unwrap(), Some(&b"Test"[..]));
        assert!(reader.peek().unwrap().is_none());
    }

    #[test]
    fn test_direction_none() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"A", b"B"]);

        let mut reader = queue.reader().unwrap();
        reader.rewind().unwrap();
        reader.set_direction(Direction::None);

        assert_eq!(reader.read().unwrap(), Some(&b"A"[..]));
        assert_eq!(reader.read().unwrap(), Some(&b"A"[..]));
    }

    #[test]
    fn test_seek() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"A", b"B", b"C", b"D", b"E"]);

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

        reader.rewind().unwrap();
        reader.read().unwrap(); // A (seq 0)
        reader.read().unwrap(); // B (seq 1)
        let index_c = reader.index();

        reader.rewind().unwrap();
        assert_eq!(reader.read().unwrap(), Some(&b"A"[..]));

        assert!(reader.seek(index_c).unwrap());
        assert_eq!(reader.read().unwrap(), Some(&b"C"[..]));
        assert_eq!(reader.read().unwrap(), Some(&b"D"[..]));
    }

    #[test]
    fn test_sequence_tracking() {
        let temp_dir = TempDir::new().unwrap();
        let queue = create_queue_with_messages(&temp_dir, &[b"One", b"Two", b"Three"]);

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

        assert_eq!(reader.sequence(), 0);
        reader.read().unwrap();
        assert_eq!(reader.sequence(), 1);
        reader.read().unwrap();
        assert_eq!(reader.sequence(), 2);
        reader.read().unwrap();
        assert_eq!(reader.sequence(), 3);
    }
}