use std::collections::{BTreeSet, HashMap};
pub(crate) use krabka_client_core::TopicNameOffsetFetch;
use krabka_protocol::{
owned::{
offset_commit_request::{OffsetCommitRequestPartition, OffsetCommitRequestTopic},
offset_fetch_request::{
OffsetFetchRequest, OffsetFetchRequestGroup, OffsetFetchRequestTopic,
OffsetFetchRequestTopics,
},
offset_fetch_response::{
OffsetFetchResponse, OffsetFetchResponsePartition, OffsetFetchResponsePartitions,
},
},
primitives::uuid::Uuid as WireUuid,
};
use crate::{
commit::OffsetAndMetadata,
coordinator::{COORDINATOR_NOT_AVAILABLE, NOT_COORDINATOR},
};
pub(crate) fn build_offset_fetch(
group_id: &str,
by_topic: &HashMap<String, Vec<i32>>,
) -> TopicNameOffsetFetch {
let legacy_topics: Vec<OffsetFetchRequestTopic> = by_topic
.iter()
.map(|(name, parts)| OffsetFetchRequestTopic {
name: name.clone(),
partition_indexes: parts.clone(),
..Default::default()
})
.collect();
let group_topics: Vec<OffsetFetchRequestTopics> = by_topic
.iter()
.map(|(name, parts)| OffsetFetchRequestTopics {
name: name.clone(),
partition_indexes: parts.clone(),
..Default::default()
})
.collect();
TopicNameOffsetFetch::new(OffsetFetchRequest {
group_id: group_id.to_string(),
topics: Some(legacy_topics),
groups: vec![OffsetFetchRequestGroup {
group_id: group_id.to_string(),
topics: Some(group_topics),
..Default::default()
}],
require_stable: true,
..Default::default()
})
}
pub(crate) fn parse_offset_fetch(resp: &OffsetFetchResponse) -> Vec<(String, i32, i64, i32)> {
let mut out = Vec::new();
if resp.groups.is_empty() {
for t in &resp.topics {
append_topic_offsets(&mut out, &t.name, &t.partitions);
}
} else {
for g in &resp.groups {
for t in &g.topics {
append_topic_offsets(&mut out, &t.name, &t.partitions);
}
}
}
out
}
trait OffsetFetchPartitionData {
fn offset_fields(&self) -> (i32, i64, i32);
fn error_code(&self) -> i16;
}
impl OffsetFetchPartitionData for OffsetFetchResponsePartition {
fn offset_fields(&self) -> (i32, i64, i32) {
(
self.partition_index,
self.committed_offset,
self.committed_leader_epoch,
)
}
fn error_code(&self) -> i16 {
self.error_code
}
}
impl OffsetFetchPartitionData for OffsetFetchResponsePartitions {
fn offset_fields(&self) -> (i32, i64, i32) {
(
self.partition_index,
self.committed_offset,
self.committed_leader_epoch,
)
}
fn error_code(&self) -> i16 {
self.error_code
}
}
fn append_topic_offsets<P: OffsetFetchPartitionData>(
out: &mut Vec<(String, i32, i64, i32)>,
topic: &str,
partitions: &[P],
) {
out.extend(partitions.iter().map(|partition| {
let (index, offset, leader_epoch) = partition.offset_fields();
(topic.to_owned(), index, offset, leader_epoch)
}));
}
fn append_topic_errors<'a, P: OffsetFetchPartitionData>(
out: &mut Vec<(&'a str, i16)>,
topic: &'a str,
partitions: &[P],
) {
out.extend(
partitions
.iter()
.map(|partition| (topic, partition.error_code())),
);
}
pub(crate) fn parse_committed_offsets(
resp: &OffsetFetchResponse,
) -> HashMap<(String, i32), Option<OffsetAndMetadata>> {
let committed = |error_code: i16, offset: i64, leader_epoch: i32, metadata: &Option<String>| {
(error_code == 0 && offset >= 0).then(|| OffsetAndMetadata {
offset,
leader_epoch: (leader_epoch >= 0).then_some(leader_epoch),
metadata: metadata.clone().unwrap_or_default(),
})
};
let mut out = HashMap::new();
for t in &resp.topics {
for p in &t.partitions {
out.insert(
(t.name.clone(), p.partition_index),
committed(
p.error_code,
p.committed_offset,
p.committed_leader_epoch,
&p.metadata,
),
);
}
}
for g in &resp.groups {
for t in &g.topics {
for p in &t.partitions {
out.insert(
(t.name.clone(), p.partition_index),
committed(
p.error_code,
p.committed_offset,
p.committed_leader_epoch,
&p.metadata,
),
);
}
}
}
out
}
const UNKNOWN_TOPIC_OR_PARTITION: i16 = 3;
const TOPIC_AUTHORIZATION_FAILED: i16 = 29;
const GROUP_AUTHORIZATION_FAILED: i16 = 30;
const UNSTABLE_OFFSET_COMMIT: i16 = 88;
const UNKNOWN_TOPIC_ID: i16 = 100;
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum OffsetFetchAction {
Complete,
RetryUnknownTopic(i16),
Retry(i16),
FindCoordinator(i16),
GroupAuthorizationFailed,
TopicAuthorizationFailed(BTreeSet<String>),
Fatal(i16),
}
pub(crate) fn classify_offset_fetch(resp: &OffsetFetchResponse) -> OffsetFetchAction {
let group_error = std::iter::once(resp.error_code)
.chain(resp.groups.iter().map(|g| g.error_code))
.find(|code| *code != 0);
if let Some(code) = group_error {
return classify_group_error(code);
}
let mut partitions = Vec::new();
for topic in &resp.topics {
append_topic_errors(&mut partitions, &topic.name, &topic.partitions);
}
for group in &resp.groups {
for topic in &group.topics {
append_topic_errors(&mut partitions, &topic.name, &topic.partitions);
}
}
let mut unauthorized_topics = BTreeSet::new();
let mut unstable = false;
let mut unknown_topic = None;
for (topic, code) in partitions {
match code {
0 => {}
UNKNOWN_TOPIC_OR_PARTITION | UNKNOWN_TOPIC_ID => {
unknown_topic.get_or_insert(code);
}
TOPIC_AUTHORIZATION_FAILED => {
unauthorized_topics.insert(topic.to_string());
}
UNSTABLE_OFFSET_COMMIT => unstable = true,
code => return OffsetFetchAction::Fatal(code),
}
}
if !unauthorized_topics.is_empty() {
OffsetFetchAction::TopicAuthorizationFailed(unauthorized_topics)
} else if unstable {
OffsetFetchAction::Retry(UNSTABLE_OFFSET_COMMIT)
} else if let Some(code) = unknown_topic {
OffsetFetchAction::RetryUnknownTopic(code)
} else {
OffsetFetchAction::Complete
}
}
fn classify_group_error(code: i16) -> OffsetFetchAction {
match code {
COORDINATOR_NOT_AVAILABLE | NOT_COORDINATOR => OffsetFetchAction::FindCoordinator(code),
GROUP_AUTHORIZATION_FAILED => OffsetFetchAction::GroupAuthorizationFailed,
code if is_retriable_error(code) => OffsetFetchAction::Retry(code),
code => OffsetFetchAction::Fatal(code),
}
}
pub(crate) fn is_retriable_error(code: i16) -> bool {
matches!(
code,
2 | 3
| 5
| 6
| 7
| 9
| 13
| 14
| 15
| 16
| 19
| 20
| 41
| 51
| 56
| 70
| 71
| 72
| 74
| 75
| 78
| 80
| 83
| 84
| 88
| 89
| 100
| 103
| 106
| 122
| 123
| 133
)
}
pub(crate) use krabka_client_core::TopicNameOffsetCommit;
pub(crate) fn build_commit_topics(
offsets: HashMap<(String, i32), OffsetAndMetadata>,
) -> Vec<OffsetCommitRequestTopic> {
let mut by_topic: HashMap<String, Vec<(i32, OffsetAndMetadata)>> = HashMap::new();
for ((t, p), offset) in offsets {
by_topic.entry(t).or_default().push((p, offset));
}
by_topic
.into_iter()
.map(|(name, parts)| OffsetCommitRequestTopic {
name,
partitions: parts
.into_iter()
.map(|(p, offset)| OffsetCommitRequestPartition {
partition_index: p,
committed_offset: offset.offset,
committed_leader_epoch: offset.leader_epoch.unwrap_or(-1),
committed_metadata: Some(offset.metadata),
..Default::default()
})
.collect(),
..Default::default()
})
.collect()
}
pub(crate) fn id_to_name(topic_ids: &HashMap<String, WireUuid>) -> HashMap<WireUuid, String> {
topic_ids.iter().map(|(n, id)| (*id, n.clone())).collect()
}
#[cfg(test)]
mod tests {
use krabka_protocol::{
UnknownTaggedFields,
owned::offset_fetch_response::{
OffsetFetchResponse, OffsetFetchResponseGroup, OffsetFetchResponsePartition,
OffsetFetchResponsePartitions, OffsetFetchResponseTopic, OffsetFetchResponseTopics,
},
};
use super::*;
use crate::test_support::wire_uuid;
#[test]
fn build_offset_fetch_populates_legacy_and_groups() {
let mut by_topic = HashMap::new();
by_topic.insert("t".to_string(), vec![0, 1]);
let req = build_offset_fetch("g", &by_topic);
assert2::assert!(
req == TopicNameOffsetFetch::new(OffsetFetchRequest {
group_id: "g".to_string(),
topics: Some(vec![OffsetFetchRequestTopic {
name: "t".to_string(),
partition_indexes: vec![0, 1],
unknown_tagged_fields: UnknownTaggedFields(vec![]),
}]),
groups: vec![OffsetFetchRequestGroup {
group_id: "g".to_string(),
member_id: None,
member_epoch: -1,
topics: Some(vec![OffsetFetchRequestTopics {
name: "t".to_string(),
topic_id: WireUuid::ZERO,
partition_indexes: vec![0, 1],
unknown_tagged_fields: UnknownTaggedFields(vec![]),
}]),
unknown_tagged_fields: UnknownTaggedFields(vec![]),
}],
require_stable: true,
unknown_tagged_fields: UnknownTaggedFields(vec![]),
})
);
}
fn topic_name(n: u8) -> String {
match n {
1 => "orders".into(),
2 => "payments".into(),
_ => "shipments".into(),
}
}
trait MissingOffsetFetchPartition {
fn with_error_code(error_code: i16) -> Self;
}
impl MissingOffsetFetchPartition for OffsetFetchResponsePartition {
fn with_error_code(error_code: i16) -> Self {
Self {
committed_offset: -1,
error_code,
..Default::default()
}
}
}
impl MissingOffsetFetchPartition for OffsetFetchResponsePartitions {
fn with_error_code(error_code: i16) -> Self {
Self {
committed_offset: -1,
error_code,
..Default::default()
}
}
}
fn grouped(group_error: i16, rows: &[(u8, i16)]) -> OffsetFetchResponse {
OffsetFetchResponse {
groups: vec![OffsetFetchResponseGroup {
group_id: "g".into(),
error_code: group_error,
topics: rows
.iter()
.map(|(topic, error_code)| OffsetFetchResponseTopics {
name: topic_name(*topic),
partitions: vec![OffsetFetchResponsePartitions::with_error_code(
*error_code,
)],
..Default::default()
})
.collect(),
..Default::default()
}],
..Default::default()
}
}
fn legacy(group_error: i16, rows: &[(&str, i16)]) -> OffsetFetchResponse {
OffsetFetchResponse {
error_code: group_error,
topics: rows
.iter()
.map(|(name, error_code)| OffsetFetchResponseTopic {
name: (*name).into(),
partitions: vec![OffsetFetchResponsePartition::with_error_code(*error_code)],
..Default::default()
})
.collect(),
..Default::default()
}
}
#[test]
fn classify_offset_fetch_maps_codes_as_kafka_does() {
use OffsetFetchAction::{
Complete, Fatal, FindCoordinator, GroupAuthorizationFailed, Retry, RetryUnknownTopic,
TopicAuthorizationFailed,
};
let topics = |names: &[&str]| {
TopicAuthorizationFailed(names.iter().map(|name| (*name).to_string()).collect())
};
for (name, response, expected) in [
("no error", grouped(0, &[(1, 0), (2, 0)]), Complete),
("no rows", grouped(0, &[]), Complete),
("coordinator load in progress", grouped(14, &[]), Retry(14)),
(
"coordinator not available",
grouped(15, &[]),
FindCoordinator(15),
),
("not coordinator", grouped(16, &[]), FindCoordinator(16)),
("request timed out", grouped(7, &[]), Retry(7)),
("network exception", grouped(13, &[]), Retry(13)),
(
"group authorization failed",
grouped(30, &[]),
GroupAuthorizationFailed,
),
("unknown member id", grouped(25, &[]), Fatal(25)),
("stale member epoch", grouped(113, &[]), Fatal(113)),
("group id not found", grouped(69, &[]), Fatal(69)),
("unknown server error", grouped(-1, &[]), Fatal(-1)),
(
"a group code hides the partition codes",
grouped(14, &[(1, 29)]),
Retry(14),
),
(
"unknown topic or partition",
grouped(0, &[(1, 0), (2, 3)]),
RetryUnknownTopic(3),
),
(
"unknown topic id",
grouped(0, &[(1, 100)]),
RetryUnknownTopic(100),
),
(
"topic authorization failed names each topic",
grouped(0, &[(1, 29), (2, 29), (3, 29)]),
topics(&["orders", "payments", "shipments"]),
),
(
"unstable offset commit",
grouped(0, &[(1, 88)]),
Retry(UNSTABLE_OFFSET_COMMIT),
),
(
"not leader or follower is unexpected",
grouped(0, &[(1, 6)]),
Fatal(6),
),
(
"unknown member id on a partition is unexpected",
grouped(0, &[(1, 25)]),
Fatal(25),
),
(
"topic authorization comes before unstable offsets",
grouped(0, &[(1, 88), (2, 29)]),
topics(&["payments"]),
),
(
"unstable offsets come before unknown topics",
grouped(0, &[(1, 100), (2, 88)]),
Retry(UNSTABLE_OFFSET_COMMIT),
),
(
"an unexpected code comes before collected codes",
grouped(0, &[(1, 29), (2, 6)]),
Fatal(6),
),
(
"the first unexpected code wins",
grouped(0, &[(1, 6), (2, 25)]),
Fatal(6),
),
(
"legacy group code",
legacy(16, &[("orders", 0)]),
FindCoordinator(16),
),
(
"legacy topic authorization failed",
legacy(0, &[("orders", 29)]),
topics(&["orders"]),
),
(
"legacy unknown topic or partition",
legacy(0, &[("orders", 3)]),
RetryUnknownTopic(3),
),
] {
assert2::check!(classify_offset_fetch(&response) == expected, "case {name}");
}
}
#[test]
fn is_retriable_error_matches_kafka_retriable_exceptions() {
let retriable = (-1..=140)
.filter(|code| is_retriable_error(*code))
.collect::<Vec<_>>();
assert2::assert!(
retriable
== vec![
2, 3, 5, 6, 7, 9, 13, 14, 15, 16, 19, 20, 41, 51, 56, 70, 71, 72, 74, 75, 78,
80, 83, 84, 88, 89, 100, 103, 106, 122, 123, 133,
]
);
}
#[test]
fn parse_offset_fetch_reads_groups_with_name() {
let resp = OffsetFetchResponse {
groups: vec![OffsetFetchResponseGroup {
group_id: "g".into(),
topics: vec![OffsetFetchResponseTopics {
name: "t".into(),
partitions: vec![OffsetFetchResponsePartitions {
partition_index: 3,
committed_offset: 42,
..Default::default()
}],
..Default::default()
}],
..Default::default()
}],
..Default::default()
};
let out = parse_offset_fetch(&resp);
assert2::assert!(out == vec![("t".to_string(), 3, 42, -1)]);
}
#[test]
fn parse_offset_fetch_falls_back_to_legacy_topics() {
let resp = OffsetFetchResponse {
topics: vec![OffsetFetchResponseTopic {
name: "legacy".into(),
partitions: vec![OffsetFetchResponsePartition {
partition_index: 1,
committed_offset: 11,
..Default::default()
}],
..Default::default()
}],
..Default::default()
};
let out = parse_offset_fetch(&resp);
assert2::assert!(out == vec![("legacy".to_string(), 1, 11, -1)]);
}
#[test]
fn build_commit_topics_names_each_topic() {
let mut offsets = HashMap::new();
offsets.insert(
("t".to_string(), 3),
OffsetAndMetadata {
offset: 100,
leader_epoch: Some(5),
metadata: "note".into(),
},
);
let topics = build_commit_topics(offsets);
assert2::assert!(
topics
== vec![OffsetCommitRequestTopic {
name: "t".to_string(),
topic_id: WireUuid::ZERO,
partitions: vec![OffsetCommitRequestPartition {
partition_index: 3,
committed_offset: 100,
committed_leader_epoch: 5,
committed_metadata: Some("note".into()),
unknown_tagged_fields: UnknownTaggedFields(vec![]),
}],
unknown_tagged_fields: UnknownTaggedFields(vec![]),
}]
);
}
#[test]
fn id_to_name_inverts_the_map() {
let mut ids = HashMap::new();
ids.insert("t".to_string(), wire_uuid(7));
let inv = id_to_name(&ids);
assert2::assert!(inv.get(&wire_uuid(7)) == Some(&"t".to_string()));
}
}