fix(kafka source): reduce flakiness in rebalance acknowledgement tests

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<usize> {
+        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<usize> = message_indices(&events1)
+                .into_iter()
+                .chain(message_indices(&events2))
+                .chain(message_indices(&events3))
+                .collect();
+            let missing: Vec<usize> = (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]