datum-mq 0.10.9

use std::{
    collections::BTreeMap,
    env,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use super::{
    AutoOffsetReset, FetchTuning, GroupAssignor, KafkaPayloadBatch, NativeKafkaConsumer,
    NativeKafkaConsumerConfig, NativeKafkaControl, NativeKafkaSource, TopicPartitionAssignment,
};
use datum::{Keep, Sink, StreamCompletion, StreamError};
use futures::executor::block_on;
use rdkafka::{
    ClientConfig, Message,
    admin::{AdminClient, AdminOptions, NewTopic, TopicReplication},
    client::DefaultClientContext,
    consumer::{BaseConsumer, Consumer},
    producer::{BaseProducer, BaseRecord, Producer},
    topic_partition_list::{Offset, TopicPartitionList},
    types::RDKafkaErrorCode,
    util::Timeout,
};

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_matches_rdkafka_from_beginning() {
    if skip_if_disabled("native_fetch_matches_rdkafka_from_beginning") {
        return;
    }
    let topic = unique_topic("diff");
    create_topic(&topic, 3);
    produce(&topic, 90, 3);

    let native = native_read(
        &topic,
        (0..3)
            .map(|partition| TopicPartitionAssignment::beginning(topic.clone(), partition))
            .collect(),
        90,
    );
    let stock = rdkafka_read(&topic, 3, 90, Offset::Beginning);
    assert_eq!(native, stock);
}

#[test]
fn native_fetch_matches_rdkafka_large_multibatch() {
    if skip_if_disabled("native_fetch_matches_rdkafka_large_multibatch") {
        return;
    }
    // 8 KiB payloads over 4 partitions: each partition's record set is far
    // larger than the 64 KiB streaming read chunk and spans many record
    // batches, so this exercises KC-14's incremental blob accumulation against
    // a live broker. Native and rdkafka must return byte-identical payloads.
    let topic = unique_topic("largemb");
    create_topic(&topic, 4);
    produce_sized(&topic, 240, 4, 8 * 1024);

    let native = native_read(
        &topic,
        (0..4)
            .map(|partition| TopicPartitionAssignment::beginning(topic.clone(), partition))
            .collect(),
        240,
    );
    let stock = rdkafka_read(&topic, 4, 240, Offset::Beginning);
    assert_eq!(native, stock);
}

#[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 rdkafka_read(
    topic: &str,
    partitions: i32,
    expected: usize,
    offset: Offset,
) -> BTreeMap<(i64, i32), Vec<u8>> {
    let consumer: BaseConsumer = ClientConfig::new()
        .set("bootstrap.servers", bootstrap())
        .set("group.id", format!("datum-mq-native-diff-{}", now_ns()))
        .set("enable.auto.commit", "false")
        .set("enable.partition.eof", "false")
        .create()
        .expect("rdkafka consumer");
    let mut tpl = TopicPartitionList::new();
    for partition in 0..partitions {
        tpl.add_partition_offset(topic, partition, offset)
            .expect("add assignment");
    }
    consumer.assign(&tpl).expect("assign");
    let deadline = std::time::Instant::now() + Duration::from_secs(30);
    let mut out = BTreeMap::new();
    while out.len() < expected && std::time::Instant::now() < deadline {
        if let Some(result) = consumer.poll(Duration::from_millis(100)) {
            let message = result.expect("rdkafka message");
            out.insert(
                (message.offset(), message.partition()),
                message.payload().unwrap_or_default().to_vec(),
            );
        }
    }
    assert_eq!(out.len(), expected, "rdkafka read count");
    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 create_topic(topic: &str, partitions: i32) {
    let admin: AdminClient<DefaultClientContext> = ClientConfig::new()
        .set("bootstrap.servers", bootstrap())
        .create()
        .expect("admin client");
    let new_topic =
        NewTopic::new(topic, partitions, TopicReplication::Fixed(1)).set("retention.ms", "600000");
    let results = block_on(admin.create_topics(&[new_topic], &AdminOptions::new()))
        .expect("topic creation response");
    for result in results {
        match result {
            Ok(_) => {}
            Err((_, RDKafkaErrorCode::TopicAlreadyExists)) => {}
            Err((name, error)) => panic!("failed to create topic {name}: {error:?}"),
        }
    }
}

fn produce(topic: &str, count: usize, partitions: i32) {
    let producer: BaseProducer = ClientConfig::new()
        .set("bootstrap.servers", bootstrap())
        .set(
            "client.id",
            format!("datum-mq-native-producer-{}", now_ns()),
        )
        .set("enable.idempotence", "false")
        .set("acks", "1")
        .set("linger.ms", "0")
        .create()
        .expect("producer");
    for index in 0..count {
        let partition = (index % partitions as usize) as i32;
        let payload = format!("value-{index}").into_bytes();
        loop {
            match producer.send(
                BaseRecord::<(), [u8], ()>::to(topic)
                    .payload(payload.as_slice())
                    .partition(partition),
            ) {
                Ok(()) => break,
                Err((
                    rdkafka::error::KafkaError::MessageProduction(RDKafkaErrorCode::QueueFull),
                    _,
                )) => producer.poll(Duration::from_millis(1)),
                Err((error, _)) => panic!("produce failed: {error}"),
            }
        }
        producer.poll(Duration::ZERO);
    }
    producer
        .flush(Timeout::After(Duration::from_secs(10)))
        .expect("flush producer");
}

fn produce_sized(topic: &str, count: usize, partitions: i32, payload_len: usize) {
    let producer: BaseProducer = ClientConfig::new()
        .set("bootstrap.servers", bootstrap())
        .set(
            "client.id",
            format!("datum-mq-native-producer-{}", now_ns()),
        )
        .set("enable.idempotence", "false")
        .set("acks", "1")
        .set("linger.ms", "0")
        .create()
        .expect("producer");
    for index in 0..count {
        let partition = (index % partitions as usize) as i32;
        let prefix = format!("record-{index}-").into_bytes();
        let payload: Vec<u8> = (0..payload_len)
            .map(|byte| {
                prefix
                    .get(byte)
                    .copied()
                    .unwrap_or_else(|| (index as u8).wrapping_add(byte as u8))
            })
            .collect();
        loop {
            match producer.send(
                BaseRecord::<(), [u8], ()>::to(topic)
                    .payload(payload.as_slice())
                    .partition(partition),
            ) {
                Ok(()) => break,
                Err((
                    rdkafka::error::KafkaError::MessageProduction(RDKafkaErrorCode::QueueFull),
                    _,
                )) => producer.poll(Duration::from_millis(1)),
                Err((error, _)) => panic!("produce failed: {error}"),
            }
        }
        producer.poll(Duration::ZERO);
    }
    producer
        .flush(Timeout::After(Duration::from_secs(10)))
        .expect("flush producer");
}

fn committed_offset_sum_for(group: &str, topic: &str, partitions: i32) -> u64 {
    let consumer: BaseConsumer = ClientConfig::new()
        .set("bootstrap.servers", bootstrap())
        .set("group.id", group)
        .set("enable.auto.commit", "false")
        .create()
        .expect("metadata consumer");
    let mut tpl = TopicPartitionList::new();
    for partition in 0..partitions {
        tpl.add_partition_offset(topic, partition, Offset::Invalid)
            .expect("add partition");
    }
    let committed = consumer
        .committed_offsets(tpl, Timeout::After(Duration::from_secs(5)))
        .expect("committed offsets");
    let mut sum = 0_u64;
    for element in committed.elements() {
        if element.topic() == topic
            && let Offset::Offset(offset) = element.offset()
        {
            sum = sum.saturating_add(offset.max(0) as u64);
        }
    }
    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
}