photon_backend/consumer_group/
static_assignment.rs

1//! Static shard assignment from environment (lab / CI).
2
3use std::env;
4
5/// Parsed static assignment configuration.
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct StaticAssignmentConfig {
8    /// Member instance id.
9    pub instance_id: String,
10    /// Shards assigned to this instance.
11    pub shard_ids: Vec<u32>,
12}
13
14/// `PHOTON_GROUP_SHARD_ASSIGNMENT` — comma list or inclusive range `0-15`.
15#[must_use]
16pub fn parse_shard_assignment(raw: &str, shard_count: u32) -> Vec<u32> {
17    let raw = raw.trim();
18    if raw.is_empty() {
19        return Vec::new();
20    }
21    if let Some((start, end)) = raw.split_once('-') {
22        if let (Ok(s), Ok(e)) = (start.trim().parse::<u32>(), end.trim().parse::<u32>()) {
23            return (s..=e.min(shard_count.saturating_sub(1))).collect();
24        }
25    }
26    raw.split(',')
27        .filter_map(|p| p.trim().parse::<u32>().ok())
28        .filter(|id| *id < shard_count)
29        .collect()
30}
31
32/// `PHOTON_GROUP_SHARD_COUNT` when set and parseable.
33#[must_use]
34pub fn shard_count_from_env() -> Option<u32> {
35    env::var("PHOTON_GROUP_SHARD_COUNT")
36        .ok()
37        .and_then(|v| v.parse().ok())
38}
39
40/// `PHOTON_GROUP_MEMBER_COUNT` when set and parseable.
41#[must_use]
42pub fn member_count_from_env() -> Option<u32> {
43    env::var("PHOTON_GROUP_MEMBER_COUNT")
44        .ok()
45        .and_then(|v| v.parse().ok())
46}
47
48/// `PHOTON_GROUP_INSTANCE_ID` when set.
49#[must_use]
50pub fn instance_id_from_env() -> Option<String> {
51    env::var("PHOTON_GROUP_INSTANCE_ID").ok()
52}
53
54/// Shards assigned to `instance_index` when splitting `shard_count` across `member_count`.
55#[must_use]
56pub fn round_robin_shards_for_member(
57    shard_count: u32,
58    member_count: u32,
59    instance_index: u32,
60) -> Vec<u32> {
61    let member_count = member_count.max(1);
62    let per = shard_count.div_ceil(member_count);
63    let start = instance_index * per;
64    let end = (start + per).min(shard_count);
65    (start..end).collect()
66}
67
68/// Shards for this process: explicit env range, else round-robin by instance index.
69#[must_use]
70pub fn static_assigned_shards(shard_count: u32) -> Vec<u32> {
71    if let Ok(raw) = env::var("PHOTON_GROUP_SHARD_ASSIGNMENT") {
72        let ids = parse_shard_assignment(&raw, shard_count);
73        if !ids.is_empty() {
74            return ids;
75        }
76    }
77    let member_count = member_count_from_env().unwrap_or(1).max(1);
78    let instance_index = instance_id_from_env()
79        .and_then(|id| id.parse::<u32>().ok())
80        .unwrap_or(0);
81    round_robin_shards_for_member(shard_count, member_count, instance_index)
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn parse_range() {
90        assert_eq!(parse_shard_assignment("0-3", 8), vec![0, 1, 2, 3]);
91    }
92
93    #[test]
94    fn parse_list() {
95        assert_eq!(parse_shard_assignment("1,3,5", 8), vec![1, 3, 5]);
96    }
97
98    #[test]
99    fn round_robin_two_members_covers_eight_shards() {
100        assert_eq!(round_robin_shards_for_member(8, 2, 0), vec![0, 1, 2, 3]);
101        assert_eq!(round_robin_shards_for_member(8, 2, 1), vec![4, 5, 6, 7]);
102    }
103}
photon_backend/consumer_group/static_assignment.rs

photon_backend/consumer_group/
static_assignment.rs
