grafos_mq/

group.rs

//! Consumer group: decentralized partition assignment with lease-based liveness.
//!
//! Each consumer in a group claims unclaimed or stale partitions from a shared
//! assignment map. Liveness is tracked via timestamps — if a consumer's
//! heartbeat goes stale beyond `stale_threshold_secs`, its partitions are
//! eligible for rebalancing.

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

use grafos_std::host;

/// Assignment entry for a single partition.
#[derive(Clone, Debug)]
struct Assignment {
    consumer_id: String,
    last_heartbeat: u64,
}

/// Shared state for a consumer group.
pub struct ConsumerGroup {
    group_name: String,
    topic_name: String,
    num_partitions: u32,
    /// Partition index -> assignment.
    assignments: BTreeMap<u32, Assignment>,
    /// Seconds after which a consumer is considered stale.
    stale_threshold_secs: u64,
}

impl ConsumerGroup {
    /// Create a new consumer group.
    pub fn new(
        group_name: &str,
        topic_name: &str,
        num_partitions: u32,
        stale_threshold_secs: u64,
    ) -> Self {
        ConsumerGroup {
            group_name: String::from(group_name),
            topic_name: String::from(topic_name),
            num_partitions,
            assignments: BTreeMap::new(),
            stale_threshold_secs,
        }
    }

    /// Claim unclaimed or stale partitions for the given consumer.
    ///
    /// Returns the list of partition indices now assigned to this consumer.
    pub fn claim(&mut self, consumer_id: &str) -> Vec<u32> {
        let now = host::unix_time_secs();
        let mut claimed = Vec::new();

        for part_idx in 0..self.num_partitions {
            let should_claim = match self.assignments.get(&part_idx) {
                None => true,
                Some(a) => {
                    if a.consumer_id == consumer_id {
                        true
                    } else {
                        now.saturating_sub(a.last_heartbeat) > self.stale_threshold_secs
                    }
                }
            };

            if should_claim {
                self.assignments.insert(
                    part_idx,
                    Assignment {
                        consumer_id: String::from(consumer_id),
                        last_heartbeat: now,
                    },
                );
                claimed.push(part_idx);
            }
        }

        claimed
    }

    /// Return the partitions currently assigned to the given consumer.
    pub fn assigned_to(&self, consumer_id: &str) -> Vec<u32> {
        self.assignments
            .iter()
            .filter(|(_, a)| a.consumer_id == consumer_id)
            .map(|(&idx, _)| idx)
            .collect()
    }

    /// Update the heartbeat for a consumer's assigned partitions.
    pub fn heartbeat(&mut self, consumer_id: &str) {
        let now = host::unix_time_secs();
        for a in self.assignments.values_mut() {
            if a.consumer_id == consumer_id {
                a.last_heartbeat = now;
            }
        }
    }

    /// Release all partitions held by the given consumer.
    pub fn release(&mut self, consumer_id: &str) {
        self.assignments.retain(|_, a| a.consumer_id != consumer_id);
    }

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

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

    /// Returns the number of unassigned partitions.
    pub fn unassigned_count(&self) -> u32 {
        let assigned = self.assignments.len() as u32;
        self.num_partitions.saturating_sub(assigned)
    }
}

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

    fn setup() {
        host::reset_mock();
        host::mock_set_unix_time_secs(1000);
    }

    #[test]
    fn single_consumer_claims_all() {
        setup();
        let mut group = ConsumerGroup::new("g1", "orders", 4, 30);
        let claimed = group.claim("c1");
        assert_eq!(claimed, vec![0, 1, 2, 3]);
        assert_eq!(group.assigned_to("c1"), vec![0, 1, 2, 3]);
        assert_eq!(group.unassigned_count(), 0);
    }

    #[test]
    fn second_consumer_waits_for_stale() {
        setup();
        let mut group = ConsumerGroup::new("g1", "orders", 4, 30);

        // c1 claims all
        group.claim("c1");

        // c2 can't claim anything yet (c1 is still fresh)
        let claimed = group.claim("c2");
        assert!(claimed.is_empty());

        // Advance time past stale threshold
        host::mock_advance_time_secs(31);

        // c2 can now claim stale partitions
        let claimed = group.claim("c2");
        assert_eq!(claimed, vec![0, 1, 2, 3]);
        assert_eq!(group.assigned_to("c2"), vec![0, 1, 2, 3]);
        assert!(group.assigned_to("c1").is_empty());
    }

    #[test]
    fn heartbeat_prevents_stale_takeover() {
        setup();
        let mut group = ConsumerGroup::new("g1", "orders", 4, 30);
        group.claim("c1");

        // Advance 20s, heartbeat
        host::mock_advance_time_secs(20);
        group.heartbeat("c1");

        // Advance another 20s (40s total, but only 20s since heartbeat)
        host::mock_advance_time_secs(20);

        // c2 can't claim (c1 heartbeated recently)
        let claimed = group.claim("c2");
        assert!(claimed.is_empty());
    }

    #[test]
    fn release_frees_partitions() {
        setup();
        let mut group = ConsumerGroup::new("g1", "orders", 4, 30);
        group.claim("c1");

        group.release("c1");
        assert_eq!(group.unassigned_count(), 4);

        // c2 can immediately claim
        let claimed = group.claim("c2");
        assert_eq!(claimed, vec![0, 1, 2, 3]);
    }
}