grafos_mq/

producer.rs

//! Producer: send messages to topic partitions.
//!
//! Supports round-robin (default) and key-hash (FNV-1a) partitioning strategies.

extern crate alloc;

use crate::topic::TopicManager;
use grafos_std::error::{FabricError, Result};

/// Partitioning strategy for the producer.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Partitioner {
    /// Round-robin across partitions.
    RoundRobin,
    /// Hash the message key using FNV-1a and select partition by modulus.
    KeyHash,
}

/// A producer that sends messages to a named topic.
pub struct Producer {
    topic_name: alloc::string::String,
    partitioner: Partitioner,
    round_robin_counter: u32,
}

impl Producer {
    /// Create a new producer for the named topic.
    pub fn new(topic_name: &str) -> Self {
        Producer {
            topic_name: alloc::string::String::from(topic_name),
            partitioner: Partitioner::RoundRobin,
            round_robin_counter: 0,
        }
    }

    /// Set the partitioning strategy.
    pub fn with_partitioner(mut self, p: Partitioner) -> Self {
        self.partitioner = p;
        self
    }

    /// Send a message without a key (uses round-robin partitioning).
    pub fn send(&mut self, mgr: &mut TopicManager, value: &[u8]) -> Result<u64> {
        let topic = mgr
            .open_mut(&self.topic_name)
            .ok_or(FabricError::IoError(-11))?;
        let np = topic.num_partitions();
        let idx = self.round_robin_counter % np;
        self.round_robin_counter = self.round_robin_counter.wrapping_add(1);
        let part = topic.partition_mut(idx).ok_or(FabricError::IoError(-12))?;
        part.append(None, value)
    }

    /// Send a keyed message.
    ///
    /// With [`Partitioner::KeyHash`], the key determines the partition.
    /// With [`Partitioner::RoundRobin`], the key is stored but partition
    /// selection still uses round-robin.
    pub fn send_keyed(&mut self, mgr: &mut TopicManager, key: &[u8], value: &[u8]) -> Result<u64> {
        let topic = mgr
            .open_mut(&self.topic_name)
            .ok_or(FabricError::IoError(-11))?;
        let np = topic.num_partitions();
        let idx = match self.partitioner {
            Partitioner::RoundRobin => {
                let i = self.round_robin_counter % np;
                self.round_robin_counter = self.round_robin_counter.wrapping_add(1);
                i
            }
            Partitioner::KeyHash => fnv1a_hash(key) % np,
        };
        let part = topic.partition_mut(idx).ok_or(FabricError::IoError(-12))?;
        part.append(Some(key), value)
    }

    /// Send a message to a specific partition.
    pub fn send_to(
        &mut self,
        mgr: &mut TopicManager,
        partition: u32,
        key: Option<&[u8]>,
        value: &[u8],
    ) -> Result<u64> {
        let topic = mgr
            .open_mut(&self.topic_name)
            .ok_or(FabricError::IoError(-11))?;
        let part = topic
            .partition_mut(partition)
            .ok_or(FabricError::IoError(-12))?;
        let offset = part.append(key, value)?;
        #[cfg(feature = "observe")]
        crate::observe_hooks::on_message_published(&self.topic_name, value.len() as u64);
        Ok(offset)
    }
}

/// FNV-1a hash over a byte slice, returning a u32-range partition index.
fn fnv1a_hash(data: &[u8]) -> u32 {
    const FNV_OFFSET: u32 = 2166136261;
    const FNV_PRIME: u32 = 16777619;
    let mut hash = FNV_OFFSET;
    for &byte in data {
        hash ^= byte as u32;
        hash = hash.wrapping_mul(FNV_PRIME);
    }
    hash
}

#[cfg(test)]
mod tests {
    use super::*;
    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: 4,
                partition_capacity: 32,
                partition_stride: 256,
            },
        )
        .expect("create");
        mgr
    }

    #[test]
    fn round_robin_distributes() {
        let mut mgr = setup();
        let mut prod = Producer::new("test");

        for _ in 0..8 {
            prod.send(&mut mgr, b"msg").expect("send");
        }

        let topic = mgr.open("test").expect("open");
        // Each partition should have 2 messages
        for i in 0..4 {
            assert_eq!(topic.partition(i).unwrap().next_offset(), 2);
        }
    }

    #[test]
    fn key_hash_same_key_same_partition() {
        let mut mgr = setup();
        let mut prod = Producer::new("test").with_partitioner(Partitioner::KeyHash);

        for _ in 0..5 {
            prod.send_keyed(&mut mgr, b"same-key", b"val")
                .expect("send");
        }

        // All messages should land in the same partition
        let topic = mgr.open("test").expect("open");
        let mut counts = [0u64; 4];
        for i in 0..4 {
            counts[i as usize] = topic.partition(i).unwrap().next_offset();
        }
        let non_zero: Vec<_> = counts.iter().filter(|&&c| c > 0).collect();
        assert_eq!(non_zero.len(), 1);
        assert_eq!(*non_zero[0], 5);
    }

    #[test]
    fn send_to_explicit_partition() {
        let mut mgr = setup();
        let mut prod = Producer::new("test");

        prod.send_to(&mut mgr, 2, Some(b"k"), b"v")
            .expect("send_to");
        let topic = mgr.open("test").expect("open");
        assert_eq!(topic.partition(2).unwrap().next_offset(), 1);
        assert_eq!(topic.partition(0).unwrap().next_offset(), 0);
    }

    #[test]
    fn send_to_nonexistent_topic() {
        let mut mgr = setup();
        let mut prod = Producer::new("no-such-topic");
        let result = prod.send(&mut mgr, b"msg");
        assert!(result.is_err());
    }

    #[test]
    fn fnv1a_deterministic() {
        assert_eq!(fnv1a_hash(b"hello"), fnv1a_hash(b"hello"));
        assert_ne!(fnv1a_hash(b"hello"), fnv1a_hash(b"world"));
    }
}