grafos_mq/

consumer.rs

//! Consumer: poll messages from assigned partitions, commit offsets, seek.

extern crate alloc;
use alloc::string::String;
use alloc::vec::Vec;

use crate::message::Message;
use crate::offset::OffsetStore;
use crate::topic::TopicManager;
use grafos_std::error::{FabricError, Result};

/// Seek policy for initializing consumer position.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SeekPolicy {
    /// Start from the earliest available message.
    Earliest,
    /// Start from the latest message (only new messages).
    Latest,
    /// Start from the last committed offset (or earliest if none).
    Committed,
}

/// A consumer that reads from specific partitions of a topic.
pub struct Consumer {
    topic_name: String,
    group: String,
    /// Assigned partition indices.
    assigned_partitions: Vec<u32>,
    /// Current read offset per partition (by partition index in assigned_partitions).
    read_offsets: Vec<u64>,
    /// Whether offsets have been initialized.
    initialized: bool,
}

impl Consumer {
    /// Create a new consumer for the named topic and consumer group.
    pub fn new(topic_name: &str, group: &str) -> Self {
        Consumer {
            topic_name: String::from(topic_name),
            group: String::from(group),
            assigned_partitions: Vec::new(),
            read_offsets: Vec::new(),
            initialized: false,
        }
    }

    /// Assign specific partitions to this consumer.
    pub fn assign(&mut self, partitions: &[u32]) {
        self.assigned_partitions = partitions.to_vec();
        self.read_offsets = alloc::vec![0u64; partitions.len()];
        self.initialized = false;
    }

    /// Initialize read offsets according to the seek policy.
    pub fn seek(
        &mut self,
        mgr: &TopicManager,
        offset_store: &dyn OffsetStore,
        policy: SeekPolicy,
    ) -> Result<()> {
        let topic = mgr
            .open(&self.topic_name)
            .ok_or(FabricError::IoError(-11))?;

        for (i, &part_idx) in self.assigned_partitions.iter().enumerate() {
            let partition = topic.partition(part_idx).ok_or(FabricError::IoError(-12))?;
            let offset = match &policy {
                SeekPolicy::Earliest => partition.oldest_offset(),
                SeekPolicy::Latest => partition.next_offset(),
                SeekPolicy::Committed => {
                    if let Some(committed) =
                        offset_store.fetch(&self.topic_name, part_idx, &self.group)
                    {
                        // Resume from after the committed offset
                        committed + 1
                    } else {
                        partition.oldest_offset()
                    }
                }
            };
            self.read_offsets[i] = offset;
        }
        self.initialized = true;
        Ok(())
    }

    /// Poll for the next batch of messages across all assigned partitions.
    ///
    /// Returns up to `max_messages` messages. Messages are interleaved across
    /// partitions in round-robin order.
    pub fn poll(&mut self, mgr: &TopicManager, max_messages: usize) -> Result<Vec<(u32, Message)>> {
        if !self.initialized {
            return Err(FabricError::IoError(-13));
        }

        let topic = mgr
            .open(&self.topic_name)
            .ok_or(FabricError::IoError(-11))?;
        let mut results = Vec::new();

        for (i, &part_idx) in self.assigned_partitions.iter().enumerate() {
            if results.len() >= max_messages {
                break;
            }
            let partition = topic.partition(part_idx).ok_or(FabricError::IoError(-12))?;
            while results.len() < max_messages && self.read_offsets[i] < partition.next_offset() {
                if let Some(msg) = partition.read_at(self.read_offsets[i])? {
                    #[cfg(feature = "observe")]
                    crate::observe_hooks::on_message_consumed(&self.topic_name, &self.group);
                    results.push((part_idx, msg));
                }
                self.read_offsets[i] += 1;
            }
        }

        Ok(results)
    }

    /// Commit the current read offsets to the offset store.
    pub fn commit(&self, offset_store: &mut dyn OffsetStore) {
        for (i, &part_idx) in self.assigned_partitions.iter().enumerate() {
            if self.read_offsets[i] > 0 {
                offset_store.commit(
                    &self.topic_name,
                    part_idx,
                    &self.group,
                    self.read_offsets[i] - 1,
                );
            }
        }
    }

    /// Returns the assigned partitions.
    pub fn assigned_partitions(&self) -> &[u32] {
        &self.assigned_partitions
    }

    /// Returns the consumer group name.
    pub fn group(&self) -> &str {
        &self.group
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::offset::MemOffsetStore;
    use crate::producer::Producer;
    use crate::topic::{TopicConfig, TopicManager};
    use grafos_std::host;

    fn setup() -> TopicManager {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(1 << 20);
        let mut mgr = TopicManager::new();
        mgr.create(
            "test",
            TopicConfig {
                num_partitions: 2,
                partition_capacity: 32,
                partition_stride: 256,
            },
        )
        .expect("create");
        mgr
    }

    #[test]
    fn produce_consume_roundtrip() {
        let mut mgr = setup();
        let mut store = MemOffsetStore::new();

        // Produce to partition 0
        let mut prod = Producer::new("test");
        prod.send_to(&mut mgr, 0, None, b"msg0").expect("send");
        prod.send_to(&mut mgr, 0, None, b"msg1").expect("send");
        prod.send_to(&mut mgr, 1, None, b"msg2").expect("send");

        // Consume
        let mut consumer = Consumer::new("test", "group-a");
        consumer.assign(&[0, 1]);
        consumer
            .seek(&mgr, &store, SeekPolicy::Earliest)
            .expect("seek");

        let msgs = consumer.poll(&mgr, 100).expect("poll");
        assert_eq!(msgs.len(), 3);
        assert_eq!(msgs[0].1.value, b"msg0");
        assert_eq!(msgs[1].1.value, b"msg1");
        assert_eq!(msgs[2].1.value, b"msg2");

        // Commit
        consumer.commit(&mut store);
        assert_eq!(store.fetch("test", 0, "group-a"), Some(1));
        assert_eq!(store.fetch("test", 1, "group-a"), Some(0));
    }

    #[test]
    fn seek_committed_resumes() {
        let mut mgr = setup();
        let mut store = MemOffsetStore::new();

        let mut prod = Producer::new("test");
        prod.send_to(&mut mgr, 0, None, b"a").expect("send");
        prod.send_to(&mut mgr, 0, None, b"b").expect("send");
        prod.send_to(&mut mgr, 0, None, b"c").expect("send");

        // First consumer reads and commits
        let mut c1 = Consumer::new("test", "g1");
        c1.assign(&[0]);
        c1.seek(&mgr, &store, SeekPolicy::Earliest).expect("seek");
        let msgs = c1.poll(&mgr, 2).expect("poll");
        assert_eq!(msgs.len(), 2);
        c1.commit(&mut store);

        // Second consumer with same group resumes from committed offset
        let mut c2 = Consumer::new("test", "g1");
        c2.assign(&[0]);
        c2.seek(&mgr, &store, SeekPolicy::Committed).expect("seek");
        let msgs = c2.poll(&mgr, 100).expect("poll");
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].1.value, b"c");
    }

    #[test]
    fn seek_latest_skips_existing() {
        let mut mgr = setup();
        let store = MemOffsetStore::new();

        let mut prod = Producer::new("test");
        prod.send_to(&mut mgr, 0, None, b"old").expect("send");

        let mut consumer = Consumer::new("test", "g1");
        consumer.assign(&[0]);
        consumer
            .seek(&mgr, &store, SeekPolicy::Latest)
            .expect("seek");

        // No messages should be available (already past existing)
        let msgs = consumer.poll(&mgr, 100).expect("poll");
        assert!(msgs.is_empty());

        // New messages should be visible
        let mut prod2 = Producer::new("test");
        prod2.send_to(&mut mgr, 0, None, b"new").expect("send");
        let msgs = consumer.poll(&mgr, 100).expect("poll");
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].1.value, b"new");
    }

    #[test]
    fn poll_without_seek_fails() {
        let mgr = setup();
        let mut consumer = Consumer::new("test", "g1");
        consumer.assign(&[0]);
        let result = consumer.poll(&mgr, 100);
        assert!(result.is_err());
    }
}