diff --git a/src/sources/kafka.rs b/src/sources/kafka.rs index 1a14680bc4ee6..131a4419e68e6 100644 --- a/src/sources/kafka.rs +++ b/src/sources/kafka.rs @@ -1571,6 +1571,35 @@ mod integration_test { const HEADER_KEY: &str = "my header"; const HEADER_VALUE: &str = "my header value"; + fn message_indices(events: &[Event]) -> HashSet { + events + .iter() + .map(|event| { + let key = event.as_log()["message_key"].to_string_lossy(); + key.strip_prefix(&format!("{KEY} ")) + .expect("message_key should have the expected prefix") + .parse() + .expect("message_key suffix should be the message index") + }) + .collect() + } + + fn message_offsets(events: &[Event]) -> HashSet<(i64, i64)> { + events + .iter() + .map(|event| { + let log = event.as_log(); + let partition = log["partition"].to_string_lossy().parse().expect( + "partition should be an integer, since it was pulled from a `sources::kafka::Keys.partition` field", + ); + let offset = log["offset"].to_string_lossy().parse().expect( + "offset should be an integer, since it was pulled from a `sources::kafka::Keys.offset` field", + ); + (partition, offset) + }) + .collect() + } + fn kafka_test_topic() -> String { std::env::var("KAFKA_TEST_TOPIC") .unwrap_or_else(|_| format!("test-topic-{}", random_string(10))) @@ -2034,8 +2063,10 @@ mod integration_test { async fn consume_with_rebalance(rebalance_strategy: String) { // 1. Send N events (if running against a pre-populated kafka topic, use send_count=0 and expect_count=expected number of messages; otherwise just set send_count) + // A larger backlog gives the later consumers a bigger margin against being starved + // (see the `events3` assertion below) as CI runners get faster at draining the topic. let send_count: usize = std::env::var("KAFKA_SEND_COUNT") - .unwrap_or_else(|_| "125000".into()) + .unwrap_or_else(|_| "500000".into()) .parse() .expect("Number of messages to send to kafka."); let expect_count: usize = std::env::var("KAFKA_EXPECT_COUNT") @@ -2141,7 +2172,49 @@ mod integration_test { 0, "The first set of consumers should consume and ack all messages." ); - assert_eq!(total, expect_count); + // Kafka only guarantees at-least-once delivery: a partition revoked mid-rebalance + // can be re-read by the consumer it's reassigned to before the prior consumer's + // offset commit lands, so `total` may legitimately exceed `expect_count` if some + // messages were delivered more than once. + assert!( + total >= expect_count, + "Consumers should not lose any messages: got {total}, expected at least {expect_count}" + ); + // Duplicates alone could mask a real drop behind an inflated `total`, so also check + // that every message was seen by at least one consumer. When we produced the messages + // ourselves via `send_events` (their `message_key` is `"{KEY} {index}"`), we can check + // this by index. A pre-populated topic (`send_count == 0`, see above) can contain + // arbitrary records, so instead we check the number of distinct Kafka + // `(partition, offset)` pairs seen: since each source record has a unique offset, + // duplicates can't inflate that count the way they inflate `total`, so it must equal + // `expect_count` exactly for no message to have been lost. + if send_count > 0 { + let received: HashSet = message_indices(&events1) + .into_iter() + .chain(message_indices(&events2)) + .chain(message_indices(&events3)) + .collect(); + let missing: Vec = (0..expect_count) + .filter(|i| !received.contains(i)) + .collect(); + assert!( + missing.is_empty(), + "Consumers should not lose any messages: got {total} events covering {}/{expect_count} indices, missing: {missing:?}", + received.len(), + ); + } else { + let received: HashSet<(i64, i64)> = message_offsets(&events1) + .into_iter() + .chain(message_offsets(&events2)) + .chain(message_offsets(&events3)) + .collect(); + assert_eq!( + received.len(), + expect_count, + "Consumers should not lose any messages: got {total} events covering only {}/{expect_count} unique offsets", + received.len(), + ); + } } #[tokio::test]