photon_backend/consumer_group/
coordinator.rs

1//! Consumer group coordinator trait and implementations.
2
3use async_trait::async_trait;
4
5use crate::consumer_group::static_assignment;
6use crate::error::Result;
7
8/// Registered group member.
9#[derive(Debug, Clone)]
10pub struct GroupMember {
11    /// Consumer group id.
12    pub group_id: String,
13    /// Unique member instance id within the group.
14    pub instance_id: String,
15    /// Topic this member consumes.
16    pub topic_name: String,
17    /// Virtual shard count for the topic.
18    pub shard_count: u32,
19}
20
21/// Assigns virtual shards to group members.
22#[async_trait]
23pub trait ConsumerGroupCoordinator: Send + Sync {
24    /// Register member and return assigned shard ids.
25    async fn register(&self, member: GroupMember) -> Result<Vec<u32>>;
26
27    /// Heartbeat to keep assignment alive (fleet lease store).
28    async fn heartbeat(&self, group_id: &str, instance_id: &str) -> Result<()> {
29        let _ = (group_id, instance_id);
30        Ok(())
31    }
32
33    /// Current assignment for a member.
34    async fn assigned_shards(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>>;
35}
36
37/// Env-based static assignment (`PHOTON_GROUP_SHARD_ASSIGNMENT`).
38pub struct StaticGroupCoordinator;
39
40#[async_trait]
41impl ConsumerGroupCoordinator for StaticGroupCoordinator {
42    async fn register(&self, member: GroupMember) -> Result<Vec<u32>> {
43        Ok(static_assignment::static_assigned_shards(member.shard_count))
44    }
45
46    async fn assigned_shards(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>> {
47        let _ = (group_id, instance_id);
48        Ok(static_assignment::static_assigned_shards(
49            static_assignment::shard_count_from_env().unwrap_or(32),
50        ))
51    }
52}
53
54/// Fleet coordinator backed by a [`super::lease_store::LeaseStore`].
55pub struct FleetGroupCoordinator<L: super::lease_store::LeaseStore> {
56    store: L,
57    lease_ttl_secs: u64,
58}
59
60impl<L: super::lease_store::LeaseStore> FleetGroupCoordinator<L> {
61    /// Create a coordinator with the given lease store and lease TTL.
62    pub const fn new(store: L, lease_ttl_secs: u64) -> Self {
63        Self {
64            store,
65            lease_ttl_secs,
66        }
67    }
68
69    fn range_assign(member_index: u32, member_count: u32, shard_count: u32) -> Vec<u32> {
70        if member_count == 0 {
71            return Vec::new();
72        }
73        let per = shard_count.div_ceil(member_count);
74        let start = member_index * per;
75        let end = (start + per).min(shard_count);
76        (start..end).collect()
77    }
78}
79
80#[async_trait]
81impl<L: super::lease_store::LeaseStore> ConsumerGroupCoordinator for FleetGroupCoordinator<L> {
82    async fn register(&self, member: GroupMember) -> Result<Vec<u32>> {
83        let member_index = member
84            .instance_id
85            .parse::<u32>()
86            .unwrap_or_else(|_| {
87                u32::try_from(member.instance_id.len()).unwrap_or(0) % member.shard_count.max(1)
88            });
89        let member_count = static_assignment::member_count_from_env().unwrap_or(2);
90        let shards = Self::range_assign(member_index, member_count, member.shard_count);
91        for shard_id in &shards {
92            self.store
93                .claim(super::lease_store::ConsumerLease {
94                    group_id: member.group_id.clone(),
95                    shard_id: *shard_id,
96                    instance_id: member.instance_id.clone(),
97                    ttl_secs: self.lease_ttl_secs,
98                })
99                .await?;
100        }
101        Ok(shards)
102    }
103
104    async fn heartbeat(&self, group_id: &str, instance_id: &str) -> Result<()> {
105        self.store
106            .renew(group_id, instance_id, self.lease_ttl_secs)
107            .await
108    }
109
110    async fn assigned_shards(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>> {
111        self.store.list_for_instance(group_id, instance_id).await
112    }
113}
photon_backend/consumer_group/coordinator.rs

photon_backend/consumer_group/
coordinator.rs
