use std::collections::{HashMap, HashSet, VecDeque};
use bytes::Bytes;
use kacrab_protocol::{
KafkaUuid,
generated::{
ApiKey, ErrorCode, FetchRequestData, FetchResponseData,
fetch_request::{FetchPartition, FetchTopic, ForgottenTopic, ReplicaState},
},
record::decode_next_batch,
};
use super::{
config::ConsumerRuntimeConfig,
error::{ConsumerError, Result},
offsets::partition_leader,
record::{ConsumerRecord, TimestampType},
subscription::{FetchPosition, SubscriptionState},
};
use crate::{
common::TopicPartition,
wire::{ClusterMetadata, WireClient},
};
const FETCH_NAME_KEYED_VERSION: i16 = 12;
const FETCH_TOPIC_ID_MIN_VERSION: i16 = 13;
const fn select_fetch_version(negotiated: i16, all_topics_have_ids: bool) -> (i16, bool) {
if all_topics_have_ids && negotiated >= FETCH_TOPIC_ID_MIN_VERSION {
(negotiated, true)
} else if negotiated < FETCH_NAME_KEYED_VERSION {
(negotiated, false)
} else {
(FETCH_NAME_KEYED_VERSION, false)
}
}
const LOG_APPEND_TIME_BIT: i16 = 0x0008;
const INVALID_SESSION_ID: i32 = 0;
const INITIAL_SESSION_EPOCH: i32 = 0;
#[derive(Debug, Default, Clone)]
struct BrokerFetchSession {
session_id: i32,
epoch: i32,
sent: HashMap<(String, i32), FetchPosition>,
uses_topic_ids: bool,
topic_ids: HashMap<String, KafkaUuid>,
}
impl BrokerFetchSession {
const fn next_epoch(&self) -> i32 {
match self.epoch.checked_add(1) {
Some(next) => next,
None => 1,
}
}
const fn is_full(&self) -> bool {
self.epoch == INITIAL_SESSION_EPOCH
}
fn reset(&mut self) {
self.session_id = INVALID_SESSION_ID;
self.epoch = INITIAL_SESSION_EPOCH;
self.sent.clear();
self.uses_topic_ids = false;
self.topic_ids.clear();
}
fn is_compatible(&self, use_topic_ids: bool, topic_ids: &HashMap<String, KafkaUuid>) -> bool {
if self.uses_topic_ids != use_topic_ids {
return false;
}
!use_topic_ids
|| topic_ids.iter().all(|(topic, id)| {
self.topic_ids
.get(topic)
.is_none_or(|session_id| session_id == id)
})
}
fn advance(
&mut self,
response_session_id: i32,
entries: &[(TopicPartition, FetchPosition)],
topic_ids: &HashMap<String, KafkaUuid>,
) {
self.session_id = response_session_id;
if response_session_id == INVALID_SESSION_ID {
self.epoch = INITIAL_SESSION_EPOCH;
self.sent.clear();
self.uses_topic_ids = false;
self.topic_ids.clear();
} else {
self.epoch = self.next_epoch();
self.sent = entries
.iter()
.map(|(partition, position)| {
((partition.topic.clone(), partition.partition), *position)
})
.collect();
self.uses_topic_ids = !topic_ids.is_empty();
self.topic_ids
.extend(topic_ids.iter().map(|(topic, id)| (topic.clone(), *id)));
self.topic_ids
.retain(|topic, _| self.sent.keys().any(|(sent, _)| sent == topic));
}
}
}
#[derive(Debug, Default)]
pub(super) struct FetchSessions {
by_broker: HashMap<i32, BrokerFetchSession>,
}
#[derive(Debug)]
pub(super) struct PartitionFetch {
pub partition: TopicPartition,
pub records: Vec<ConsumerRecord>,
pub next_offset: i64,
pub next_leader_epoch: Option<i32>,
}
#[derive(Debug)]
pub(super) struct RawPartitionFetch {
pub partition: TopicPartition,
pub fetch_position: FetchPosition,
pub records: Bytes,
}
#[derive(Debug, Default)]
pub(super) struct FetchProgress {
pub partitions: Vec<RawPartitionFetch>,
pub resets: Vec<TopicPartition>,
pub stale: Vec<TopicPartition>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PartitionErrorAction {
Reset,
Retriable,
Fatal,
}
const fn classify_partition_error(error: ErrorCode) -> PartitionErrorAction {
match error {
ErrorCode::OffsetOutOfRange => PartitionErrorAction::Reset,
ErrorCode::NotLeaderOrFollower
| ErrorCode::FencedLeaderEpoch
| ErrorCode::UnknownLeaderEpoch
| ErrorCode::UnknownTopicOrPartition
| ErrorCode::UnknownTopicId
| ErrorCode::InconsistentTopicId
| ErrorCode::ReplicaNotAvailable
| ErrorCode::LeaderNotAvailable
| ErrorCode::KafkaStorageError
| ErrorCode::OffsetNotAvailable => PartitionErrorAction::Retriable,
_ => PartitionErrorAction::Fatal,
}
}
#[derive(Debug, Clone, Copy)]
pub(super) struct FetchContext<'a> {
pub wire: &'a WireClient,
pub config: &'a ConsumerRuntimeConfig,
pub metadata: &'a ClusterMetadata,
pub max_wait_ms: i32,
}
pub(super) async fn fetch(
context: &FetchContext<'_>,
fetchable: &[(TopicPartition, FetchPosition)],
sessions: &mut FetchSessions,
) -> Result<FetchProgress> {
let mut progress = FetchProgress::default();
if fetchable.is_empty() {
return Ok(progress);
}
let FetchContext {
wire,
config,
metadata,
max_wait_ms,
} = *context;
let mut by_leader: HashMap<i32, Vec<(TopicPartition, FetchPosition)>> = HashMap::new();
for (partition, position) in fetchable {
let Some(leader) = partition_leader(metadata, &partition.topic, partition.partition) else {
continue;
};
by_leader
.entry(leader)
.or_default()
.push((partition.clone(), *position));
}
for (leader, entries) in by_leader {
let mut topic_ids: HashMap<String, KafkaUuid> = HashMap::new();
let mut all_have_ids = true;
for (partition, _) in &entries {
if topic_ids.contains_key(&partition.topic) {
continue;
}
match metadata.topic(&partition.topic).map(|topic| topic.topic_id) {
Some(id) if !id.is_nil() => {
let _previous = topic_ids.insert(partition.topic.clone(), id);
},
_ => all_have_ids = false,
}
}
let negotiated = wire
.negotiated_version(leader, ApiKey::Fetch)
.unwrap_or(FETCH_NAME_KEYED_VERSION);
let (version, use_topic_ids) = select_fetch_version(negotiated, all_have_ids);
if !use_topic_ids {
topic_ids.clear();
}
let session = sessions.by_broker.entry(leader).or_default();
if !session.is_full() && !session.is_compatible(use_topic_ids, &topic_ids) {
session.reset();
}
let request = build_fetch_request(
config,
session,
&entries,
max_wait_ms,
use_topic_ids.then_some(&topic_ids),
);
let response: FetchResponseData = wire
.send_to_broker(leader, ApiKey::Fetch, version, &request)
.await?;
let top_level = ErrorCode::from(response.error_code);
if matches!(
top_level,
ErrorCode::InvalidFetchSessionEpoch | ErrorCode::FetchSessionIdNotFound
) {
session.reset();
continue;
}
if top_level == ErrorCode::FetchSessionTopicIdError {
session.reset();
progress
.stale
.extend(entries.iter().map(|(partition, _)| partition.clone()));
continue;
}
if top_level.is_error() {
return Err(ConsumerError::broker(
"fetch",
top_level,
"fetch request rejected",
));
}
session.advance(response.session_id, &entries, &topic_ids);
let want: HashMap<(String, i32), FetchPosition> = entries
.iter()
.map(|(tp, pos)| ((tp.topic.clone(), tp.partition), *pos))
.collect();
let topic_names: HashMap<KafkaUuid, &str> = topic_ids
.iter()
.map(|(name, id)| (*id, name.as_str()))
.collect();
collect_fetches(response, &want, &topic_names, &mut progress)?;
}
Ok(progress)
}
fn collect_fetches(
response: FetchResponseData,
want: &HashMap<(String, i32), FetchPosition>,
topic_names: &HashMap<KafkaUuid, &str>,
progress: &mut FetchProgress,
) -> Result<()> {
for topic in response.responses {
let name = if topic.topic.as_str().is_empty() {
let Some(name) = topic_names.get(&topic.topic_id) else {
continue;
};
(*name).to_owned()
} else {
topic.topic.as_str().to_owned()
};
for partition in topic.partitions {
let tp = TopicPartition::new(name.clone(), partition.partition_index);
let Some(position) = want.get(&(tp.topic.clone(), tp.partition)).copied() else {
continue;
};
let error = ErrorCode::from(partition.error_code);
if error.is_error() {
match classify_partition_error(error) {
PartitionErrorAction::Reset => progress.resets.push(tp),
PartitionErrorAction::Retriable => progress.stale.push(tp),
PartitionErrorAction::Fatal => {
return Err(ConsumerError::broker(
"fetch",
error,
format!("{}-{} fetch failed", tp.topic, tp.partition),
));
},
}
continue;
}
let Some(records) = partition.records.filter(|blob| !blob.is_empty()) else {
continue;
};
progress.partitions.push(RawPartitionFetch {
partition: tp,
fetch_position: position,
records,
});
}
}
Ok(())
}
#[derive(Debug, Default)]
pub(super) struct FetchBuffer {
buffered: VecDeque<BufferedFetch>,
}
#[derive(Debug)]
enum BufferedFetch {
Raw(RawPartitionFetch),
Decoded(DecodedFetch),
}
impl BufferedFetch {
const fn partition(&self) -> &TopicPartition {
match self {
Self::Raw(raw) => &raw.partition,
Self::Decoded(decoded) => &decoded.partition,
}
}
const fn position_offset(&self) -> i64 {
match self {
Self::Raw(raw) => raw.fetch_position.offset,
Self::Decoded(decoded) => decoded.position_offset,
}
}
}
#[derive(Debug)]
struct DecodedFetch {
partition: TopicPartition,
topic: std::sync::Arc<str>,
blob: Bytes,
records: VecDeque<ConsumerRecord>,
position_offset: i64,
next_offset: i64,
next_leader_epoch: Option<i32>,
}
impl DecodedFetch {
fn new(raw: RawPartitionFetch) -> Self {
Self {
topic: std::sync::Arc::from(raw.partition.topic.as_str()),
partition: raw.partition,
blob: raw.records,
records: VecDeque::new(),
position_offset: raw.fetch_position.offset,
next_offset: raw.fetch_position.offset,
next_leader_epoch: None,
}
}
fn is_exhausted(&self) -> bool {
self.records.is_empty() && self.blob.is_empty()
}
fn refill(&mut self, budget: usize) -> Result<()> {
while self.records.len() < budget && !self.blob.is_empty() {
let batch = decode_next_batch(&mut self.blob).map_err(|_error| {
ConsumerError::InvalidState("failed to decode fetched record batch")
})?;
let Some(batch) = batch else {
self.blob = Bytes::new();
break;
};
let leader_epoch =
(batch.partition_leader_epoch >= 0).then_some(batch.partition_leader_epoch);
self.next_leader_epoch = leader_epoch;
let log_append_time = batch.attributes & LOG_APPEND_TIME_BIT != 0;
let timestamp_type = if log_append_time {
TimestampType::LogAppendTime
} else {
TimestampType::CreateTime
};
for record in batch.records {
let offset = batch
.base_offset
.saturating_add(i64::from(record.offset_delta));
if offset < self.position_offset {
continue;
}
let timestamp = if log_append_time {
batch.max_timestamp
} else {
batch.first_timestamp.saturating_add(record.timestamp_delta)
};
self.records.push_back(ConsumerRecord {
topic: std::sync::Arc::clone(&self.topic),
partition: self.partition.partition,
offset,
timestamp,
timestamp_type,
key: record.key,
value: record.value,
headers: record.headers,
leader_epoch,
});
self.next_offset = offset.saturating_add(1);
}
}
Ok(())
}
}
impl FetchBuffer {
pub(super) fn has(&self, partition: &TopicPartition) -> bool {
self.buffered
.iter()
.any(|entry| entry.partition() == partition)
}
pub(super) fn partitions(&self) -> impl Iterator<Item = &TopicPartition> {
self.buffered.iter().map(BufferedFetch::partition)
}
pub(super) fn push(&mut self, raw: RawPartitionFetch) {
self.buffered.push_back(BufferedFetch::Raw(raw));
}
pub(super) fn drain(
&mut self,
subscription: &SubscriptionState,
max_records: usize,
) -> Result<Vec<PartitionFetch>> {
let mut out = Vec::new();
let mut budget = max_records;
let mut visits_left = self.buffered.len();
while budget > 0 && visits_left > 0 {
visits_left = visits_left.saturating_sub(1);
let Some(entry) = self.buffered.pop_front() else {
break;
};
let partition = entry.partition().clone();
let Some(position) = subscription.position(&partition) else {
continue;
};
if position.offset != entry.position_offset() {
continue;
}
if subscription.is_paused(&partition) {
self.buffered.push_back(entry);
continue;
}
let mut decoded = match entry {
BufferedFetch::Decoded(decoded) => decoded,
BufferedFetch::Raw(raw) => DecodedFetch::new(raw),
};
decoded.refill(budget)?;
if decoded.records.is_empty() {
continue;
}
let take = budget.min(decoded.records.len());
let records: Vec<ConsumerRecord> = decoded.records.drain(..take).collect();
budget = budget.saturating_sub(take);
let finished = decoded.is_exhausted();
let (next_offset, next_leader_epoch) = if finished {
(decoded.next_offset, decoded.next_leader_epoch)
} else {
records
.last()
.map_or((decoded.next_offset, decoded.next_leader_epoch), |last| {
(last.offset.saturating_add(1), last.leader_epoch)
})
};
out.push(PartitionFetch {
partition: decoded.partition.clone(),
records,
next_offset,
next_leader_epoch,
});
if !finished {
decoded.position_offset = next_offset;
self.buffered.push_front(BufferedFetch::Decoded(decoded));
}
}
Ok(out)
}
}
fn build_fetch_request(
config: &ConsumerRuntimeConfig,
session: &BrokerFetchSession,
entries: &[(TopicPartition, FetchPosition)],
max_wait_ms: i32,
topic_ids: Option<&HashMap<String, KafkaUuid>>,
) -> FetchRequestData {
let full = session.is_full();
let mut grouped: Vec<(&str, Vec<FetchPartition>)> = Vec::new();
for (partition, position) in entries {
if !full
&& session
.sent
.get(&(partition.topic.clone(), partition.partition))
== Some(position)
{
continue;
}
let wire_partition = FetchPartition {
partition: partition.partition,
current_leader_epoch: position.leader_epoch.unwrap_or(-1),
fetch_offset: position.offset,
last_fetched_epoch: -1,
log_start_offset: -1,
partition_max_bytes: config.max_partition_fetch_bytes,
replica_directory_id: KafkaUuid::default(),
high_watermark: i64::MAX,
_unknown_tagged_fields: Vec::new(),
};
if let Some((_, partitions)) = grouped
.iter_mut()
.find(|(topic, _)| *topic == partition.topic)
{
partitions.push(wire_partition);
} else {
grouped.push((partition.topic.as_str(), vec![wire_partition]));
}
}
let topics: Vec<FetchTopic> = grouped
.into_iter()
.map(|(name, partitions)| match topic_ids {
Some(ids) => FetchTopic {
topic: kacrab_protocol::KafkaString::default(),
topic_id: ids.get(name).copied().unwrap_or_default(),
partitions,
_unknown_tagged_fields: Vec::new(),
},
None => FetchTopic {
topic: name.to_owned().into(),
topic_id: KafkaUuid::default(),
partitions,
_unknown_tagged_fields: Vec::new(),
},
})
.collect();
let forgotten = if full {
Vec::new()
} else {
build_forgotten(session, entries)
};
FetchRequestData {
cluster_id: None,
replica_id: -1,
replica_state: ReplicaState {
replica_id: -1,
replica_epoch: -1,
_unknown_tagged_fields: Vec::new(),
},
max_wait_ms,
min_bytes: config.fetch_min_bytes,
max_bytes: config.fetch_max_bytes,
isolation_level: config.isolation_level.wire(),
session_id: session.session_id,
session_epoch: session.epoch,
topics,
forgotten_topics_data: forgotten,
rack_id: config.client_rack.clone().into(),
_unknown_tagged_fields: Vec::new(),
}
}
fn build_forgotten(
session: &BrokerFetchSession,
entries: &[(TopicPartition, FetchPosition)],
) -> Vec<ForgottenTopic> {
let current: HashSet<(&str, i32)> = entries
.iter()
.map(|(partition, _)| (partition.topic.as_str(), partition.partition))
.collect();
let mut grouped: Vec<(&str, Vec<i32>)> = Vec::new();
for (topic, partition) in session.sent.keys() {
if current.contains(&(topic.as_str(), *partition)) {
continue;
}
if let Some((_, partitions)) = grouped.iter_mut().find(|(name, _)| name == topic) {
partitions.push(*partition);
} else {
grouped.push((topic.as_str(), vec![*partition]));
}
}
grouped
.into_iter()
.map(|(name, partitions)| {
if session.uses_topic_ids {
ForgottenTopic {
topic: kacrab_protocol::KafkaString::default(),
topic_id: session.topic_ids.get(name).copied().unwrap_or_default(),
partitions,
_unknown_tagged_fields: Vec::new(),
}
} else {
ForgottenTopic {
topic: name.to_owned().into(),
topic_id: KafkaUuid::default(),
partitions,
_unknown_tagged_fields: Vec::new(),
}
}
})
.collect()
}
#[cfg(test)]
mod tests {
use bytes::BytesMut;
use kacrab_protocol::record::{Record, RecordBatch};
use super::*;
fn record(offset_delta: i32, key: &str, value: &str) -> Record {
Record {
attributes: 0,
timestamp_delta: i64::from(offset_delta),
offset_delta,
key: Some(Bytes::copy_from_slice(key.as_bytes())),
value: Some(Bytes::copy_from_slice(value.as_bytes())),
headers: Vec::new(),
}
}
fn encode_batch(base_offset: i64, count: i32) -> Bytes {
let records: Vec<Record> = (0..count)
.map(|i| record(i, &format!("k{i}"), &format!("v{i}")))
.collect();
let batch = RecordBatch {
base_offset,
partition_leader_epoch: 4,
magic: 2,
attributes: 0,
last_offset_delta: count.saturating_sub(1),
first_timestamp: 1_000,
max_timestamp: 1_000_i64.saturating_add(i64::from(count.saturating_sub(1))),
producer_id: -1,
producer_epoch: -1,
base_sequence: -1,
records,
};
let mut buf = BytesMut::new();
batch.encode(&mut buf).expect("encode batch");
buf.freeze()
}
fn tp() -> TopicPartition {
TopicPartition::new("t", 0)
}
#[test]
fn decodes_records_with_absolute_offsets() {
let mut decoded = DecodedFetch::new(RawPartitionFetch {
partition: tp(),
fetch_position: FetchPosition::new(100, None),
records: encode_batch(100, 3),
});
decoded.refill(10).unwrap();
assert_eq!(decoded.records.len(), 3);
assert_eq!(decoded.records[0].offset, 100);
assert_eq!(decoded.records[2].offset, 102);
assert_eq!(decoded.records[0].value.as_deref(), Some(b"v0".as_ref()));
assert_eq!(decoded.records[0].leader_epoch, Some(4));
assert_eq!(decoded.next_offset, 103);
assert!(decoded.blob.is_empty());
}
#[test]
fn skips_records_before_the_fetch_offset() {
let mut decoded = DecodedFetch::new(RawPartitionFetch {
partition: tp(),
fetch_position: FetchPosition::new(101, None),
records: encode_batch(100, 3),
});
decoded.refill(10).unwrap();
assert_eq!(decoded.records.len(), 2);
assert_eq!(decoded.records[0].offset, 101);
assert_eq!(decoded.next_offset, 103);
}
#[test]
fn refill_decodes_lazily_per_batch() {
let mut blob = BytesMut::new();
blob.extend_from_slice(&encode_batch(100, 3));
blob.extend_from_slice(&encode_batch(103, 3));
let mut decoded = DecodedFetch::new(RawPartitionFetch {
partition: tp(),
fetch_position: FetchPosition::new(100, None),
records: blob.freeze(),
});
decoded.refill(2).unwrap();
assert_eq!(decoded.records.len(), 3);
assert!(!decoded.blob.is_empty());
decoded.refill(6).unwrap();
assert_eq!(decoded.records.len(), 6);
assert!(decoded.blob.is_empty());
assert_eq!(decoded.next_offset, 106);
}
#[test]
fn truncated_trailing_batch_is_dropped_cleanly() {
let whole = encode_batch(100, 3);
let mut blob = BytesMut::new();
blob.extend_from_slice(&whole);
blob.extend_from_slice(&whole.slice(..whole.len() / 2));
let mut buffer = FetchBuffer::default();
buffer.push(RawPartitionFetch {
partition: tp(),
fetch_position: FetchPosition::new(100, None),
records: blob.freeze(),
});
let subscription = subscription_at("t", 0, 100);
let drained = buffer.drain(&subscription, 10).expect("drain");
assert_eq!(drained.len(), 1);
assert_eq!(drained[0].records.len(), 3);
assert_eq!(drained[0].next_offset, 103);
assert!(!buffer.has(&tp()));
}
fn entry(topic: &str, partition: i32, offset: i64) -> (TopicPartition, FetchPosition) {
(
TopicPartition::new(topic, partition),
FetchPosition::new(offset, None),
)
}
#[test]
fn session_opens_full_then_goes_incremental() {
let mut session = BrokerFetchSession::default();
assert!(session.is_full());
assert_eq!(session.session_id, INVALID_SESSION_ID);
let entries = vec![entry("t", 0, 10), entry("t", 1, 20)];
session.advance(42, &entries, &HashMap::new());
assert!(!session.is_full());
assert_eq!(session.session_id, 42);
assert_eq!(session.epoch, 1);
assert_eq!(session.sent.len(), 2);
assert!(!session.uses_topic_ids);
session.reset();
assert!(session.is_full());
assert!(session.sent.is_empty());
}
#[test]
fn session_without_broker_id_stays_full() {
let mut session = BrokerFetchSession::default();
session.advance(INVALID_SESSION_ID, &[entry("t", 0, 5)], &HashMap::new());
assert!(session.is_full());
assert!(session.sent.is_empty());
}
fn test_config() -> ConsumerRuntimeConfig {
let client: crate::config::ClientConfig =
[("bootstrap.servers", "127.0.0.1:9092"), ("group.id", "g")]
.into_iter()
.collect();
ConsumerRuntimeConfig::from_config(&client.consumer_config().expect("config"))
.expect("runtime")
}
#[test]
fn full_fetch_sends_all_partitions_then_incremental_sends_only_changes() {
let config = test_config();
let mut session = BrokerFetchSession::default();
let entries = vec![entry("t", 0, 10), entry("t", 1, 20)];
let full = build_fetch_request(&config, &session, &entries, 500, None);
assert_eq!(full.session_epoch, INITIAL_SESSION_EPOCH);
assert_eq!(
full.topics
.iter()
.map(|t| t.partitions.len())
.sum::<usize>(),
2
);
assert!(full.forgotten_topics_data.is_empty());
session.advance(99, &entries, &HashMap::new());
let changed = vec![entry("t", 0, 10), entry("t", 1, 25)];
let incremental = build_fetch_request(&config, &session, &changed, 500, None);
assert_eq!(incremental.session_id, 99);
assert_eq!(incremental.session_epoch, 1);
let sent: Vec<i32> = incremental
.topics
.iter()
.flat_map(|topic| topic.partitions.iter().map(|p| p.partition))
.collect();
assert_eq!(sent, vec![1]);
assert!(incremental.forgotten_topics_data.is_empty());
}
use kacrab_protocol::generated::{
FetchResponseData,
fetch_response::{FetchableTopicResponse, PartitionData},
};
fn partition_data(index: i32, error: ErrorCode, records: Option<Bytes>) -> PartitionData {
PartitionData {
partition_index: index,
error_code: error.code(),
records,
..PartitionData::default()
}
}
fn fetch_response(topic: &str, partitions: Vec<PartitionData>) -> FetchResponseData {
FetchResponseData {
responses: vec![FetchableTopicResponse {
topic: topic.to_owned().into(),
partitions,
..FetchableTopicResponse::default()
}],
..FetchResponseData::default()
}
}
fn want_two() -> HashMap<(String, i32), FetchPosition> {
[
(("t".to_owned(), 0), FetchPosition::new(100, None)),
(("t".to_owned(), 1), FetchPosition::new(0, None)),
]
.into_iter()
.collect()
}
#[test]
fn out_of_range_partition_resets_while_healthy_partition_survives() {
let response = fetch_response(
"t",
vec![
partition_data(0, ErrorCode::OffsetOutOfRange, None),
partition_data(1, ErrorCode::None, Some(encode_batch(0, 2))),
],
);
let mut progress = FetchProgress::default();
collect_fetches(response, &want_two(), &HashMap::new(), &mut progress)
.expect("no fatal error");
assert_eq!(progress.resets, vec![TopicPartition::new("t", 0)]);
assert!(progress.stale.is_empty());
assert_eq!(progress.partitions.len(), 1);
assert_eq!(
progress.partitions[0].partition,
TopicPartition::new("t", 1)
);
assert_eq!(progress.partitions[0].fetch_position.offset, 0);
assert!(!progress.partitions[0].records.is_empty());
}
#[test]
fn stale_leader_partition_is_flagged_not_fatal() {
let response = fetch_response(
"t",
vec![partition_data(0, ErrorCode::NotLeaderOrFollower, None)],
);
let mut progress = FetchProgress::default();
collect_fetches(response, &want_two(), &HashMap::new(), &mut progress).expect("retriable");
assert_eq!(progress.stale, vec![TopicPartition::new("t", 0)]);
assert!(progress.resets.is_empty());
assert!(progress.partitions.is_empty());
}
#[test]
fn genuinely_fatal_partition_error_propagates() {
let response = fetch_response(
"t",
vec![partition_data(0, ErrorCode::CorruptMessage, None)],
);
let mut progress = FetchProgress::default();
assert!(collect_fetches(response, &want_two(), &HashMap::new(), &mut progress).is_err());
}
fn buffer_with(topic: &str, partition: i32, offset: i64, count: i32) -> FetchBuffer {
let mut buffer = FetchBuffer::default();
buffer.push(RawPartitionFetch {
partition: TopicPartition::new(topic, partition),
fetch_position: FetchPosition::new(offset, None),
records: encode_batch(offset, count),
});
buffer
}
fn subscription_at(topic: &str, partition: i32, offset: i64) -> SubscriptionState {
let mut subscription = SubscriptionState::new(crate::consumer::AutoOffsetReset::Earliest);
let tp = TopicPartition::new(topic, partition);
subscription.assign(std::slice::from_ref(&tp));
subscription.set_position(&tp, FetchPosition::new(offset, None));
subscription
}
#[test]
fn buffer_drains_in_max_poll_slices_across_polls() {
let mut buffer = buffer_with("t", 0, 100, 5);
let mut subscription = subscription_at("t", 0, 100);
let tp = TopicPartition::new("t", 0);
assert!(buffer.has(&tp));
let first = buffer.drain(&subscription, 2).expect("drain");
assert_eq!(first.len(), 1);
assert_eq!(first[0].records.len(), 2);
assert_eq!(first[0].next_offset, 102);
subscription.advance_position(&tp, first[0].next_offset, first[0].next_leader_epoch);
assert!(buffer.has(&tp));
let second = buffer.drain(&subscription, 2).expect("drain");
assert_eq!(second[0].records.len(), 2);
assert_eq!(second[0].records[0].offset, 102);
subscription.advance_position(&tp, second[0].next_offset, second[0].next_leader_epoch);
let third = buffer.drain(&subscription, 2).expect("drain");
assert_eq!(third[0].records.len(), 1);
assert_eq!(third[0].next_offset, 105);
assert!(!buffer.has(&tp));
}
#[test]
fn seek_invalidates_buffered_records() {
let mut buffer = buffer_with("t", 0, 100, 5);
let subscription = subscription_at("t", 0, 42);
let drained = buffer.drain(&subscription, 10).expect("drain");
assert!(drained.is_empty());
assert!(!buffer.has(&TopicPartition::new("t", 0)));
}
#[test]
fn revoked_partition_buffer_is_dropped() {
let mut buffer = buffer_with("t", 0, 100, 5);
let mut subscription = subscription_at("u", 1, 0);
subscription.set_position(&TopicPartition::new("u", 1), FetchPosition::new(0, None));
let drained = buffer.drain(&subscription, 10).expect("drain");
assert!(drained.is_empty());
assert!(!buffer.has(&TopicPartition::new("t", 0)));
}
#[test]
fn paused_partition_buffer_is_retained_until_resumed() {
let mut buffer = buffer_with("t", 0, 100, 3);
let mut subscription = subscription_at("t", 0, 100);
let tp = TopicPartition::new("t", 0);
subscription.pause(std::slice::from_ref(&tp));
let drained = buffer.drain(&subscription, 10).expect("drain");
assert!(drained.is_empty());
assert!(buffer.has(&tp));
subscription.resume(std::slice::from_ref(&tp));
let drained = buffer.drain(&subscription, 10).expect("drain");
assert_eq!(drained[0].records.len(), 3);
assert!(!buffer.has(&tp));
}
#[test]
fn paused_entry_does_not_starve_other_partitions() {
let mut buffer = buffer_with("t", 0, 100, 3);
buffer.push(RawPartitionFetch {
partition: TopicPartition::new("t", 1),
fetch_position: FetchPosition::new(0, None),
records: encode_batch(0, 2),
});
let mut subscription = subscription_at("t", 0, 100);
let other = TopicPartition::new("t", 1);
subscription.assign(&[TopicPartition::new("t", 0), other.clone()]);
subscription.set_position(&TopicPartition::new("t", 0), FetchPosition::new(100, None));
subscription.set_position(&other, FetchPosition::new(0, None));
subscription.pause(&[TopicPartition::new("t", 0)]);
let drained = buffer.drain(&subscription, 10).expect("drain");
assert_eq!(drained.len(), 1);
assert_eq!(drained[0].partition, other);
assert_eq!(drained[0].records.len(), 2);
assert!(buffer.has(&TopicPartition::new("t", 0)));
}
#[test]
fn blob_entirely_before_position_is_dropped() {
let mut buffer = FetchBuffer::default();
buffer.push(RawPartitionFetch {
partition: TopicPartition::new("t", 0),
fetch_position: FetchPosition::new(103, None),
records: encode_batch(100, 3),
});
let subscription = subscription_at("t", 0, 103);
let drained = buffer.drain(&subscription, 10).expect("drain");
assert!(drained.is_empty());
assert!(!buffer.has(&TopicPartition::new("t", 0)));
}
fn uuid(byte: u8) -> KafkaUuid {
KafkaUuid::from_parts(u64::from(byte), u64::from(byte))
}
fn ids(pairs: &[(&str, KafkaUuid)]) -> HashMap<String, KafkaUuid> {
pairs
.iter()
.map(|(name, id)| ((*name).to_owned(), *id))
.collect()
}
#[test]
fn version_selection_mirrors_java_downgrade() {
assert_eq!(select_fetch_version(17, true), (17, true));
assert_eq!(select_fetch_version(13, true), (13, true));
assert_eq!(select_fetch_version(17, false), (12, false));
assert_eq!(select_fetch_version(12, true), (12, false));
assert_eq!(select_fetch_version(4, false), (4, false));
}
#[test]
fn topic_id_request_omits_names_and_encodes_at_v17() {
let config = test_config();
let session = BrokerFetchSession::default();
let entries = vec![entry("t", 0, 10), entry("u", 1, 20)];
let topic_ids = ids(&[("t", uuid(1)), ("u", uuid(2))]);
let request = build_fetch_request(&config, &session, &entries, 500, Some(&topic_ids));
for topic in &request.topics {
assert!(topic.topic.as_str().is_empty());
assert!(!topic.topic_id.is_nil());
}
let sent: HashSet<KafkaUuid> = request.topics.iter().map(|t| t.topic_id).collect();
assert_eq!(sent, [uuid(1), uuid(2)].into_iter().collect());
let mut buf = BytesMut::new();
request.write(&mut buf, 17).expect("v17 encode");
assert_eq!(buf.len(), request.encoded_len(17).expect("v17 len"));
assert!(request.write(&mut BytesMut::new(), 12).is_err());
}
#[test]
fn name_keyed_request_encodes_at_v12_but_not_v17() {
let config = test_config();
let session = BrokerFetchSession::default();
let entries = vec![entry("t", 0, 10)];
let request = build_fetch_request(&config, &session, &entries, 500, None);
assert_eq!(request.topics[0].topic.as_str(), "t");
assert!(request.topics[0].topic_id.is_nil());
request.write(&mut BytesMut::new(), 12).expect("v12 encode");
assert!(request.write(&mut BytesMut::new(), 17).is_err());
}
#[test]
fn session_compatibility_tracks_ids_and_mode() {
let mut session = BrokerFetchSession::default();
let entries = vec![entry("t", 0, 10)];
session.advance(42, &entries, &ids(&[("t", uuid(1))]));
assert!(session.uses_topic_ids);
assert!(session.is_compatible(true, &ids(&[("t", uuid(1))])));
assert!(session.is_compatible(true, &ids(&[("t", uuid(1)), ("u", uuid(2))])));
assert!(!session.is_compatible(true, &ids(&[("t", uuid(9))])));
assert!(!session.is_compatible(false, &HashMap::new()));
session.reset();
session.advance(42, &entries, &HashMap::new());
assert!(!session.is_compatible(true, &ids(&[("t", uuid(1))])));
assert!(session.is_compatible(false, &HashMap::new()));
}
#[test]
fn forgotten_topics_carry_the_session_ids() {
let mut session = BrokerFetchSession::default();
session.advance(
7,
&[entry("t", 0, 1), entry("u", 0, 1)],
&ids(&[("t", uuid(1)), ("u", uuid(2))]),
);
let forgotten = build_forgotten(&session, &[entry("t", 0, 1)]);
assert_eq!(forgotten.len(), 1);
assert!(forgotten[0].topic.as_str().is_empty());
assert_eq!(forgotten[0].topic_id, uuid(2));
assert_eq!(forgotten[0].partitions, vec![0]);
}
#[test]
fn collect_fetches_resolves_topics_by_id() {
let mut response = fetch_response(
"",
vec![partition_data(
0,
ErrorCode::None,
Some(encode_batch(100, 2)),
)],
);
response.responses[0].topic_id = uuid(1);
let topic_names: HashMap<KafkaUuid, &str> = std::iter::once((uuid(1), "t")).collect();
let mut progress = FetchProgress::default();
collect_fetches(response, &want_two(), &topic_names, &mut progress).expect("collect");
assert_eq!(progress.partitions.len(), 1);
assert_eq!(
progress.partitions[0].partition,
TopicPartition::new("t", 0)
);
let mut unknown = fetch_response(
"",
vec![partition_data(
0,
ErrorCode::None,
Some(encode_batch(100, 2)),
)],
);
unknown.responses[0].topic_id = uuid(9);
let mut progress = FetchProgress::default();
collect_fetches(unknown, &want_two(), &topic_names, &mut progress).expect("collect");
assert!(progress.partitions.is_empty());
}
#[test]
fn unknown_topic_id_errors_are_retriable() {
for error in [ErrorCode::UnknownTopicId, ErrorCode::InconsistentTopicId] {
let response = fetch_response("t", vec![partition_data(0, error, None)]);
let mut progress = FetchProgress::default();
collect_fetches(response, &want_two(), &HashMap::new(), &mut progress)
.expect("retriable");
assert_eq!(progress.stale, vec![TopicPartition::new("t", 0)]);
}
}
#[test]
fn forgotten_lists_partitions_dropped_from_the_session() {
let mut session = BrokerFetchSession::default();
session.advance(
7,
&[entry("t", 0, 1), entry("t", 1, 1), entry("u", 0, 1)],
&HashMap::new(),
);
let forgotten = build_forgotten(&session, &[entry("t", 0, 1)]);
let mut pairs: Vec<(String, i32)> = forgotten
.into_iter()
.flat_map(|topic| {
let name = topic.topic.as_str().to_owned();
topic
.partitions
.into_iter()
.map(move |partition| (name.clone(), partition))
})
.collect();
pairs.sort();
assert_eq!(pairs, vec![("t".to_owned(), 1), ("u".to_owned(), 0)]);
}
}