use std::{
collections::BTreeMap,
env,
path::PathBuf,
process::Command,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use super::{
AutoOffsetReset, FetchTuning, GroupAssignor, KafkaPayloadBatch, NativeKafkaConsumer,
NativeKafkaConsumerConfig, NativeKafkaControl, NativeKafkaSource, TopicPartitionAssignment,
};
use crate::{KafkaProducerSettings, KafkaSink, ProducerRecord};
use bytes::Bytes;
use datum::{Keep, Sink, Source, StreamCompletion, StreamError};
fn kafka_enabled() -> bool {
env::var("DATUM_MQ_TEST_KAFKA").is_ok_and(|value| value == "1")
}
fn skip_if_disabled(test: &str) -> bool {
if kafka_enabled() {
false
} else {
eprintln!("skipping {test}; set DATUM_MQ_TEST_KAFKA=1 and run baselines/mq/up.sh");
true
}
}
fn bootstrap() -> String {
env::var("MQ_BOOTSTRAP_SERVERS").unwrap_or_else(|_| "127.0.0.1:9092".to_owned())
}
fn unique_topic(prefix: &str) -> String {
format!("datum-mq-native-{prefix}-{}", now_ns())
}
#[test]
fn native_fetch_supports_specific_offsets_and_reassignment() {
if skip_if_disabled("native_fetch_supports_specific_offsets_and_reassignment") {
return;
}
let topic = unique_topic("assign");
create_topic(&topic, 2);
produce(&topic, 40, 2);
let partition_zero = native_read(
&topic,
vec![TopicPartitionAssignment::absolute(topic.clone(), 0, 5)],
15,
);
assert!(partition_zero.keys().all(|(_, partition)| *partition == 0));
assert!(partition_zero.keys().all(|(offset, _)| *offset >= 5));
let partition_one = native_read(
&topic,
vec![TopicPartitionAssignment::absolute(topic.clone(), 1, 5)],
15,
);
assert!(partition_one.keys().all(|(_, partition)| *partition == 1));
assert!(partition_one.keys().all(|(offset, _)| *offset >= 5));
}
#[test]
fn native_fetch_latest_starts_at_end() {
if skip_if_disabled("native_fetch_latest_starts_at_end") {
return;
}
let topic = unique_topic("latest");
create_topic(&topic, 1);
produce(&topic, 10, 1);
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime");
let fetched = runtime.block_on(async {
let mut config = consumer_config(vec![TopicPartitionAssignment::end(topic.clone(), 0)]);
config.fetch.max_wait_ms = 25;
let mut consumer = NativeKafkaConsumer::connect(config)
.await
.expect("native consumer");
tokio::time::timeout(Duration::from_millis(250), consumer.fetch_batch())
.await
.expect("fetch timeout")
.expect("fetch result")
});
assert!(fetched.is_none());
}
#[test]
fn native_source_commits_offsets_to_broker_and_restarts_from_committed() {
if skip_if_disabled("native_source_commits_offsets_to_broker_and_restarts_from_committed") {
return;
}
let topic = unique_topic("commit");
let group = format!("{topic}-group");
create_topic(&topic, 2);
produce(&topic, 24, 2);
let consumed = source_consume(
&topic,
&group,
(0..2)
.map(|partition| TopicPartitionAssignment::beginning(topic.clone(), partition))
.collect(),
24,
true,
);
assert_eq!(consumed, 24);
assert_eq!(committed_offset_sum_for(&group, &topic, 2), 24);
produce(&topic, 6, 2);
let resumed = source_consume(
&topic,
&group,
(0..2)
.map(|partition| TopicPartitionAssignment::committed(topic.clone(), partition))
.collect(),
6,
true,
);
assert_eq!(resumed, 6);
assert_eq!(committed_offset_sum_for(&group, &topic, 2), 30);
}
#[test]
fn native_source_replays_when_shutdown_happens_before_commit() {
if skip_if_disabled("native_source_replays_when_shutdown_happens_before_commit") {
return;
}
let topic = unique_topic("replay");
let group = format!("{topic}-group");
create_topic(&topic, 1);
produce(&topic, 12, 1);
let consumed_before_kill = source_consume(
&topic,
&group,
vec![TopicPartitionAssignment::beginning(topic.clone(), 0)],
4,
false,
);
assert!(consumed_before_kill >= 4);
assert_eq!(committed_offset_sum_for(&group, &topic, 1), 0);
let replayed = source_consume(
&topic,
&group,
vec![TopicPartitionAssignment::committed(topic.clone(), 0)],
12,
true,
);
assert_eq!(replayed, 12);
assert_eq!(committed_offset_sum_for(&group, &topic, 1), 12);
}
#[test]
fn native_group_subscription_two_members_commit_and_leave() {
if skip_if_disabled("native_group_subscription_two_members_commit_and_leave") {
return;
}
let topic = unique_topic("group-two");
let group = format!("{topic}-group");
create_topic(&topic, 4);
produce(&topic, 80, 4);
let counter = Arc::new(TestCounter::new());
let first = start_group_source(&topic, &group, "member-a", Arc::clone(&counter), 8);
let second = start_group_source(&topic, &group, "member-b", Arc::clone(&counter), 8);
assert!(
wait_until(Duration::from_secs(45), || counter.unique() >= 80),
"native group consumers did not consume all records"
);
first
.0
.drain_and_shutdown(Duration::from_secs(30))
.expect("first drains");
second
.0
.drain_and_shutdown(Duration::from_secs(30))
.expect("second drains");
first.1.wait().expect("first completes");
second.1.wait().expect("second completes");
assert_eq!(counter.unique(), 80);
assert_eq!(committed_offset_sum_for(&group, &topic, 4), 80);
}
#[test]
fn native_group_static_member_restarts_from_committed() {
if skip_if_disabled("native_group_static_member_restarts_from_committed") {
return;
}
let topic = unique_topic("static");
let group = format!("{topic}-group");
create_topic(&topic, 2);
produce(&topic, 20, 2);
let first_counter = Arc::new(TestCounter::new());
let first = start_group_source(&topic, &group, "static-a", Arc::clone(&first_counter), 4);
assert!(wait_until(Duration::from_secs(30), || first_counter
.unique()
>= 20));
first
.0
.drain_and_shutdown(Duration::from_secs(30))
.expect("first static drains");
first.1.wait().expect("first static completes");
assert_eq!(committed_offset_sum_for(&group, &topic, 2), 20);
produce(&topic, 10, 2);
let second_counter = Arc::new(TestCounter::new());
let second = start_group_source(&topic, &group, "static-a", Arc::clone(&second_counter), 4);
assert!(wait_until(Duration::from_secs(30), || second_counter
.unique()
>= 10));
second
.0
.drain_and_shutdown(Duration::from_secs(30))
.expect("second static drains");
second.1.wait().expect("second static completes");
assert_eq!(committed_offset_sum_for(&group, &topic, 2), 30);
}
#[test]
fn native_group_backpressure_pause_still_drains() {
if skip_if_disabled("native_group_backpressure_pause_still_drains") {
return;
}
let topic = unique_topic("bp");
let group = format!("{topic}-group");
create_topic(&topic, 2);
produce(&topic, 32, 2);
let counter = Arc::new(TestCounter::new());
let sink_counter = Arc::clone(&counter);
let mut config = NativeKafkaConsumerConfig::subscribe(bootstrap(), vec![topic.to_owned()])
.with_group_id(group.clone())
.with_group_instance_id("member-a")
.with_group_assignor(GroupAssignor::CooperativeSticky);
config.client.client_id = format!("datum-mq-native-group-{}", now_ns());
config.auto_offset_reset = AutoOffsetReset::Earliest;
config.commit_batch_size = 2;
config.commit_interval = Duration::from_millis(10);
config.fetch = FetchTuning {
max_wait_ms: 50,
min_bytes: 1,
output_records: 4,
output_payload_bytes: 4096,
include_timestamps: false,
high_watermark: 4,
low_watermark: 1,
..FetchTuning::default()
};
let (control, completion): (NativeKafkaControl, StreamCompletion<datum::NotUsed>) =
NativeKafkaSource::payload_batches(config)
.to_mat(
Sink::foreach_result(move |batch: KafkaPayloadBatch| {
std::thread::sleep(Duration::from_millis(5));
sink_counter.observe(&batch);
batch.commit().map_err(StreamError::from)?;
Ok(())
}),
Keep::both,
)
.run()
.expect("native backpressure group source materializes");
assert!(
wait_until(Duration::from_secs(30), || counter.unique() >= 32),
"native group source did not consume through backpressure"
);
control
.drain_and_shutdown(Duration::from_secs(30))
.expect("backpressured source drains");
completion.wait().expect("backpressured source completes");
assert_eq!(committed_offset_sum_for(&group, &topic, 2), 32);
}
#[test]
fn native_fetch_recovers_offset_out_of_range_with_metric() {
if skip_if_disabled("native_fetch_recovers_offset_out_of_range_with_metric") {
return;
}
let topic = unique_topic("oor");
create_topic(&topic, 1);
produce(&topic, 5, 1);
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime");
let metrics = runtime.block_on(async {
let mut config = consumer_config(vec![TopicPartitionAssignment::absolute(
topic.clone(),
0,
10_000,
)]);
config.auto_offset_reset = AutoOffsetReset::Earliest;
config.fetch.max_wait_ms = 25;
let metrics = config.metrics.clone();
let mut consumer = NativeKafkaConsumer::connect(config)
.await
.expect("native consumer");
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
let mut records = 0_usize;
while records < 5 && tokio::time::Instant::now() < deadline {
if let Some(batch) = consumer.fetch_batch().await.expect("native fetch") {
records += batch.len();
}
}
assert_eq!(records, 5);
metrics
});
assert_eq!(metrics.snapshot().offset_out_of_range_resets, 1);
}
#[test]
fn broker_restart_mid_fetch_is_env_gated() {
if env::var("DATUM_KAFKA_CLIENT_TEST_RESTART").ok().as_deref() != Some("1") {
eprintln!(
"skipping broker_restart_mid_fetch_is_env_gated; set DATUM_KAFKA_CLIENT_TEST_RESTART=1 to run the destructive broker-restart check"
);
return;
}
if skip_if_disabled("broker_restart_mid_fetch_is_env_gated") {
return;
}
let topic = unique_topic("restart");
create_topic(&topic, 1);
produce(&topic, 100, 1);
let read = native_read(
&topic,
vec![TopicPartitionAssignment::beginning(topic.clone(), 0)],
100,
);
assert_eq!(read.len(), 100);
}
fn native_read(
topic: &str,
assignments: Vec<TopicPartitionAssignment>,
expected: usize,
) -> BTreeMap<(i64, i32), Vec<u8>> {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime");
runtime.block_on(async move {
let mut consumer = NativeKafkaConsumer::connect(consumer_config(assignments))
.await
.expect("native consumer");
let mut out = BTreeMap::new();
let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
while out.len() < expected && tokio::time::Instant::now() < deadline {
if let Some(batch) = consumer.fetch_batch().await.expect("native fetch") {
for record in batch.records() {
out.insert(
(record.offset, record.partition),
batch.payload(record).to_vec(),
);
}
}
}
assert_eq!(out.len(), expected, "native read count for topic {topic}");
out
})
}
fn consumer_config(assignments: Vec<TopicPartitionAssignment>) -> NativeKafkaConsumerConfig {
let mut config = NativeKafkaConsumerConfig::new(bootstrap(), assignments);
config.client.client_id = format!("datum-mq-native-test-{}", now_ns());
config.group_id = format!("datum-mq-native-group-{}", now_ns());
config.fetch = FetchTuning {
max_wait_ms: 100,
min_bytes: 1,
output_records: 64,
output_payload_bytes: 16 * 1024,
include_timestamps: false,
..FetchTuning::default()
};
config
}
fn source_consume(
topic: &str,
group: &str,
assignments: Vec<TopicPartitionAssignment>,
expected: usize,
commit: bool,
) -> usize {
let consumed = Arc::new(AtomicUsize::new(0));
let mut config = consumer_config(assignments);
config.group_id = group.to_owned();
config.auto_offset_reset = AutoOffsetReset::Earliest;
config.commit_batch_size = 4;
config.commit_interval = Duration::from_millis(10);
config.fetch.output_records = 4;
config.fetch.output_payload_bytes = 1024;
let sink_count = Arc::clone(&consumed);
let (control, completion): (NativeKafkaControl, StreamCompletion<datum::NotUsed>) =
NativeKafkaSource::payload_batches(config)
.to_mat(
Sink::foreach_result(move |batch: KafkaPayloadBatch| {
sink_count.fetch_add(batch.len(), Ordering::Relaxed);
if commit {
batch.commit().map_err(StreamError::from)?;
}
Ok(())
}),
Keep::both,
)
.run()
.expect("native source materializes");
assert!(
wait_until(Duration::from_secs(30), || {
consumed.load(Ordering::Relaxed) >= expected
}),
"native source consumed fewer than {expected} records for {topic}"
);
if commit {
control
.drain_and_shutdown(Duration::from_secs(30))
.expect("native source drains");
} else {
control.shutdown_now();
}
completion.wait().expect("native source completes");
consumed.load(Ordering::Relaxed)
}
fn start_group_source(
topic: &str,
group: &str,
instance_id: &str,
counter: Arc<TestCounter>,
output_records: usize,
) -> (NativeKafkaControl, StreamCompletion<datum::NotUsed>) {
let mut config = NativeKafkaConsumerConfig::subscribe(bootstrap(), vec![topic.to_owned()])
.with_group_id(group.to_owned())
.with_group_instance_id(instance_id.to_owned())
.with_group_assignor(GroupAssignor::CooperativeSticky);
config.client.client_id = format!("datum-mq-native-group-{}", now_ns());
config.auto_offset_reset = AutoOffsetReset::Earliest;
config.commit_batch_size = output_records;
config.commit_interval = Duration::from_millis(10);
config.fetch = FetchTuning {
max_wait_ms: 50,
min_bytes: 1,
output_records,
output_payload_bytes: 4096,
include_timestamps: false,
high_watermark: output_records,
low_watermark: output_records.saturating_sub(1),
..FetchTuning::default()
};
NativeKafkaSource::payload_batches(config)
.to_mat(
Sink::foreach_result(move |batch: KafkaPayloadBatch| {
counter.observe(&batch);
batch.commit().map_err(StreamError::from)?;
Ok(())
}),
Keep::both,
)
.run()
.expect("native group source materializes")
}
#[derive(Debug, Default)]
struct TestCounter {
seen: std::sync::Mutex<BTreeMap<(i64, i32), ()>>,
}
impl TestCounter {
fn new() -> Self {
Self::default()
}
fn observe(&self, batch: &KafkaPayloadBatch) {
let mut seen = self.seen.lock().expect("test counter lock");
for record in batch.records() {
seen.insert((record.offset, record.partition), ());
}
}
fn unique(&self) -> usize {
self.seen.lock().expect("test counter lock").len()
}
}
fn mq_root() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../baselines/mq")
}
fn create_topic(topic: &str, partitions: i32) {
let status = Command::new(mq_root().join(".runtime/kafka_2.13-4.2.0/bin/kafka-topics.sh"))
.arg("--bootstrap-server")
.arg(bootstrap())
.arg("--create")
.arg("--if-not-exists")
.arg("--topic")
.arg(topic)
.arg("--partitions")
.arg(partitions.to_string())
.arg("--replication-factor")
.arg("1")
.status()
.expect("run kafka-topics.sh");
assert!(status.success(), "create Kafka topic {topic}");
}
fn produce(topic: &str, count: usize, partitions: i32) {
let records = (0..count).map(|index| {
ProducerRecord::new(topic, Some(Bytes::from(format!("value-{index}"))))
.with_partition((index % partitions as usize) as i32)
});
let settings = KafkaProducerSettings::new(bootstrap())
.with("enable.idempotence", "false")
.with("acks", "1")
.with("linger.ms", "0")
.with_drain_timeout(Duration::from_secs(20));
let control = Source::from_iterable(records)
.run_with(KafkaSink::plain(settings))
.expect("native producer materializes");
control
.drain_and_shutdown()
.expect("native producer drains");
}
fn committed_offset_sum_for(group: &str, topic: &str, _partitions: i32) -> u64 {
let output =
Command::new(mq_root().join(".runtime/kafka_2.13-4.2.0/bin/kafka-consumer-groups.sh"))
.arg("--bootstrap-server")
.arg(bootstrap())
.arg("--describe")
.arg("--group")
.arg(group)
.output()
.expect("run kafka-consumer-groups.sh");
assert!(
output.status.success(),
"describe Kafka consumer group {group}"
);
String::from_utf8(output.stdout)
.expect("consumer group output is UTF-8")
.lines()
.filter_map(|line| {
let columns = line.split_whitespace().collect::<Vec<_>>();
(columns.len() >= 4 && columns[0] == group && columns[1] == topic)
.then(|| columns[3].parse::<u64>().unwrap_or(0))
})
.sum()
}
fn wait_until<F>(timeout: Duration, mut predicate: F) -> bool
where
F: FnMut() -> bool,
{
let deadline = std::time::Instant::now() + timeout;
while std::time::Instant::now() < deadline {
if predicate() {
return true;
}
std::thread::sleep(Duration::from_millis(10));
}
predicate()
}
fn now_ns() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system time after unix epoch")
.as_nanos()
.min(u128::from(u64::MAX)) as u64
}