From 8653e53f5012f0212e92225c79d2522538376334 Mon Sep 17 00:00:00 2001 From: Thomas Date: Tue, 7 Jul 2026 15:15:46 -0400 Subject: [PATCH 1/4] fix(sources): reduce flakiness in kafka rebalance acknowledgement tests --- src/sources/kafka.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/sources/kafka.rs b/src/sources/kafka.rs index 1a14680bc4ee6..4a4a45e016fdf 100644 --- a/src/sources/kafka.rs +++ b/src/sources/kafka.rs @@ -2034,8 +2034,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 +2143,13 @@ 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`. + assert!( + total >= expect_count, + "Consumers should not lose any messages: got {total}, expected at least {expect_count}" + ); } #[tokio::test] From 9d458c5489b55e4f1fe9dba5ad741b4c0234d515 Mon Sep 17 00:00:00 2001 From: Thomas Date: Tue, 7 Jul 2026 15:23:42 -0400 Subject: [PATCH 2/4] fix(sources): assert exact per-message delivery in kafka rebalance test --- src/sources/kafka.rs | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/src/sources/kafka.rs b/src/sources/kafka.rs index 4a4a45e016fdf..0f06049496d7e 100644 --- a/src/sources/kafka.rs +++ b/src/sources/kafka.rs @@ -1571,6 +1571,19 @@ 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 kafka_test_topic() -> String { std::env::var("KAFKA_TEST_TOPIC") .unwrap_or_else(|_| format!("test-topic-{}", random_string(10))) @@ -2145,11 +2158,27 @@ mod integration_test { ); // 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`. + // 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 index was seen by at least one consumer. + 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(), + ); } #[tokio::test] From b93269ffad426dbf91db3a4b027325db3a546ece Mon Sep 17 00:00:00 2001 From: Thomas Date: Tue, 7 Jul 2026 15:56:40 -0400 Subject: [PATCH 3/4] fix(sources): skip exact index check for pre-populated kafka topics --- src/sources/kafka.rs | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/src/sources/kafka.rs b/src/sources/kafka.rs index 0f06049496d7e..731af4954150b 100644 --- a/src/sources/kafka.rs +++ b/src/sources/kafka.rs @@ -2165,20 +2165,25 @@ mod integration_test { "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 index was seen by at least one consumer. - 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(), - ); + // that every message index was seen by at least one consumer. This only applies + // when we produced the messages ourselves via `send_events` (their `message_key` + // is `"{KEY} {index}"`); a pre-populated topic (`send_count == 0`, see above) can + // contain arbitrary records, so fall back to the count-only check in that case. + 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(), + ); + } } #[tokio::test] From b846c02af0c61ff095482aa10a1cc90849ef104c Mon Sep 17 00:00:00 2001 From: Thomas Date: Wed, 8 Jul 2026 15:05:21 -0400 Subject: [PATCH 4/4] fix(sources): detect duplicate-masked message loss for pre-populated kafka topics --- src/sources/kafka.rs | 39 +++++++++++++++++++++++++++++++++++---- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/src/sources/kafka.rs b/src/sources/kafka.rs index 731af4954150b..131a4419e68e6 100644 --- a/src/sources/kafka.rs +++ b/src/sources/kafka.rs @@ -1584,6 +1584,22 @@ mod integration_test { .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))) @@ -2165,10 +2181,13 @@ mod integration_test { "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 index was seen by at least one consumer. This only applies - // when we produced the messages ourselves via `send_events` (their `message_key` - // is `"{KEY} {index}"`); a pre-populated topic (`send_count == 0`, see above) can - // contain arbitrary records, so fall back to the count-only check in that case. + // 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() @@ -2183,6 +2202,18 @@ mod integration_test { "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(), + ); } }