krabka-client-consumer 0.7.1

Subscribe-style consumer client for Apache Kafka in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

//! Round robin assignor, Kafka's `RoundRobinAssignor`.
//!
//! It sorts every partition of every subscribed topic by topic and partition,
//! and hands them out one by one to the members in member order, skipping a
//! member that does not subscribe to the topic of the partition.

use std::collections::{BTreeSet, HashMap};

use super::MemberKey;

/// Returns `member_id → Vec<(topic, partition)>` assignments.
#[must_use]
#[tracing::instrument(
    name = "consumer.assignor.round_robin",
    level = "info",
    skip_all,
    fields(members = members.len(), topics = topic_partitions.len())
)]
pub fn assign(
    mut members: Vec<(MemberKey, Vec<String>)>,
    topic_partitions: &HashMap<String, i32>,
) -> HashMap<String, Vec<(String, i32)>> {
    let mut out: HashMap<String, Vec<(String, i32)>> = members
        .iter()
        .map(|(member, _)| (member.member_id.clone(), Vec::new()))
        .collect();
    if members.is_empty() {
        return out;
    }
    members.sort_by(|(a, _), (b, _)| a.cmp(b));
    let topics: BTreeSet<&String> = members.iter().flat_map(|(_, topics)| topics).collect();
    let mut next = 0;
    for topic in topics {
        let Some(&count) = topic_partitions.get(topic) else {
            continue;
        };
        for partition in 0..count.max(0) {
            // Kafka's `CircularIterator`: advance to the next member that
            // subscribes to the topic. One member subscribes to it at least.
            while !members[next % members.len()].1.contains(topic) {
                next += 1;
            }
            let member = &members[next % members.len()].0.member_id;
            next += 1;
            if let Some(slot) = out.get_mut(member) {
                slot.push((topic.clone(), partition));
            }
        }
    }
    out
}

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

    fn member(id: &str, instance: Option<&str>, topics: &[&str]) -> (MemberKey, Vec<String>) {
        (
            MemberKey {
                member_id: id.into(),
                group_instance_id: instance.map(str::to_owned),
            },
            topics.iter().map(|topic| (*topic).to_owned()).collect(),
        )
    }

    fn partitions(items: &[(&str, i32)]) -> Vec<(String, i32)> {
        items
            .iter()
            .map(|(topic, p)| ((*topic).to_owned(), *p))
            .collect()
    }

    /// The cases of Kafka's `RoundRobinAssignorTest`.
    #[test]
    fn round_robin_assigns_as_kafka_does() {
        let counts = maplit::hashmap! {"t1".to_owned() => 3, "t2".to_owned() => 3};
        let mut actual = Vec::new();
        let mut wanted = Vec::new();
        for (name, members, expected) in [
            (
                "two members, two topics",
                vec![
                    member("c1", None, &["t1", "t2"]),
                    member("c2", None, &["t1", "t2"]),
                ],
                maplit::hashmap! {"c1".to_owned() => partitions(&[("t1", 0), ("t1", 2), ("t2", 1)]), "c2".to_owned() => partitions(&[("t1", 1), ("t2", 0), ("t2", 2)])},
            ),
            (
                "members with different subscriptions",
                vec![
                    member("c0", None, &["t1"]),
                    member("c1", None, &["t1", "t2"]),
                    member("c2", None, &["t1"]),
                ],
                maplit::hashmap! {"c0".to_owned() => partitions(&[("t1", 0)]), "c1".to_owned() => partitions(&[("t1", 1), ("t2", 0), ("t2", 1), ("t2", 2)]), "c2".to_owned() => partitions(&[("t1", 2)])},
            ),
            (
                "static members sort by group instance id",
                vec![
                    member("b", Some("i2"), &["t1"]),
                    member("a", Some("i1"), &["t1"]),
                ],
                maplit::hashmap! {"a".to_owned() => partitions(&[("t1", 0), ("t1", 2)]), "b".to_owned() => partitions(&[("t1", 1)])},
            ),
            (
                "a topic without metadata",
                vec![member("c1", None, &["t3"])],
                maplit::hashmap! {"c1".to_owned() => Vec::new()},
            ),
        ] {
            actual.push((name, assign(members, &counts)));
            wanted.push((name, expected));
        }
        assert2::assert!(actual == wanted);
    }
}