fix(producer): prevent send requests with 0 as num_messages (#377)

BewareMyPower · web-flow · commit d963be5198a3 · 2025-12-17T21:53:54.000+08:00
diff --git a/src/lib.rs b/src/lib.rs
@@ -225,6 +225,7 @@ pub mod reader;
 mod retry_op;
 pub mod routing_policy;
 mod service_discovery;
+mod test_utils;
 
 #[cfg(all(
     any(
@@ -248,7 +249,7 @@ mod tests {
 
     use assert_matches::assert_matches;
     use futures::{future::try_join_all, StreamExt};
-    use log::{LevelFilter, Metadata, Record};
+    use log::{Metadata, Record};
     use serde_json::Value;
     #[cfg(any(
         feature = "tokio-runtime",
@@ -263,7 +264,6 @@ mod tests {
         feature = "tokio-rustls-runtime-aws-lc-rs",
         feature = "tokio-rustls-runtime-ring"
     ))]
-    use crate::executor::TokioExecutor;
     use crate::{
         client::SerializeMessage,
         consumer::{InitialPosition, Message},
@@ -332,11 +332,7 @@ mod tests {
         feature = "tokio-rustls-runtime-ring"
     ))]
     async fn round_trip() {
-        let _result = log::set_logger(&TEST_LOGGER);
-        log::set_max_level(LevelFilter::Debug);
-
-        let addr = "pulsar://127.0.0.1:6650";
-        let pulsar: Pulsar<_> = Pulsar::builder(addr, TokioExecutor).build().await.unwrap();
+        let pulsar = test_utils::new_pulsar().await;
 
         // random topic to better allow multiple test runs while debugging
         let topic = format!("test_{}", rand::random::<u16>());
@@ -401,12 +397,8 @@ mod tests {
         feature = "tokio-rustls-runtime-ring"
     ))]
     async fn unsized_data() {
-        let _result = log::set_logger(&TEST_LOGGER);
-        log::set_max_level(LevelFilter::Debug);
-
-        let addr = "pulsar://127.0.0.1:6650";
+        let pulsar = test_utils::new_pulsar().await;
         let test_id: u16 = rand::random();
-        let pulsar: Pulsar<_> = Pulsar::builder(addr, TokioExecutor).build().await.unwrap();
 
         // test &str
         {
@@ -491,13 +483,8 @@ mod tests {
         feature = "tokio-rustls-runtime-ring"
     ))]
     async fn redelivery() {
-        let _result = log::set_logger(&TEST_LOGGER);
-        log::set_max_level(LevelFilter::Debug);
-
-        let addr = "pulsar://127.0.0.1:6650";
+        let pulsar = test_utils::new_pulsar().await;
         let topic = format!("test_redelivery_{}", rand::random::<u16>());
-
-        let pulsar: Pulsar<_> = Pulsar::builder(addr, TokioExecutor).build().await.unwrap();
         pulsar
             .send(&topic, String::from("data"))
             .await
@@ -545,13 +532,8 @@ mod tests {
     async fn batching() {
         use assert_matches::assert_matches;
 
-        let _result = log::set_logger(&TEST_LOGGER);
-        log::set_max_level(LevelFilter::Debug);
-
-        let addr = "pulsar://127.0.0.1:6650";
+        let pulsar = test_utils::new_pulsar().await;
         let topic = format!("test_batching_{}", rand::random::<u16>());
-
-        let pulsar: Pulsar<_> = Pulsar::builder(addr, TokioExecutor).build().await.unwrap();
         let mut consumer: Consumer<String, _> =
             pulsar.consumer().with_topic(&topic).build().await.unwrap();
 
@@ -677,13 +659,8 @@ mod tests {
         feature = "tokio-rustls-runtime-ring"
     ))]
     async fn flush() {
-        let _result = log::set_logger(&TEST_LOGGER);
-        log::set_max_level(LevelFilter::Debug);
-
-        let addr = "pulsar://127.0.0.1:6650";
+        let pulsar = test_utils::new_pulsar().await;
         let topic = format!("test_flush_{}", rand::random::<u16>());
-
-        let pulsar: Pulsar<_> = Pulsar::builder(addr, TokioExecutor).build().await.unwrap();
         let mut consumer: Consumer<String, _> =
             pulsar.consumer().with_topic(&topic).build().await.unwrap();
         let mut producer =
@@ -825,4 +802,33 @@ mod tests {
             panic!("No publishers in the stats");
         }
     }
+
+    #[tokio::test]
+    async fn flush_on_partitioned_topic() {
+        let pulsar = test_utils::new_pulsar().await;
+        let topic = format!("test_flush_on_part_topic_{}", rand::random::<u16>());
+
+        const NUM_PARTITIONS: u32 = 2;
+        test_utils::create_partitioned_topic("public", "default", &topic, NUM_PARTITIONS).await;
+        let mut producer =
+            create_batched_producer(pulsar.clone(), &topic, Some(2), None, None).await;
+        let send_future = producer.send_non_blocking("msg").await.unwrap();
+        producer.send_batch().await.unwrap();
+
+        let msg_id = send_future.await.unwrap().message_id.unwrap();
+        for i in 0..NUM_PARTITIONS {
+            let mut reader = pulsar
+                .reader()
+                .with_topic(format!("{}-partition-{}", topic, i))
+                .build::<String>()
+                .await
+                .unwrap();
+            let last_msg_ids = reader.get_last_message_id().await.unwrap();
+            let last_msg_id = last_msg_ids.first().unwrap();
+            if last_msg_id.ledger_id != u64::MAX && last_msg_id.entry_id != u64::MAX {
+                assert_eq!(last_msg_id.ledger_id, msg_id.ledger_id);
+                assert_eq!(last_msg_id.entry_id, msg_id.entry_id);
+            }
+        }
+    }
 }
diff --git a/src/producer.rs b/src/producer.rs
@@ -1227,6 +1227,12 @@ async fn message_send_loop<Exe>(
                         );
                     }
                 }
+                if counter == 0 {
+                    if let Some(flush_tx) = flush_tx {
+                        let _ = flush_tx.send(());
+                    }
+                    continue;
+                }
 
                 let message = ProducerMessage {
                     payload,
@@ -1435,7 +1441,9 @@ mod tests {
     use log::LevelFilter;
 
     use super::*;
-    use crate::{routing_policy::CustomRoutingPolicy, tests::TEST_LOGGER, TokioExecutor};
+    use crate::{
+        routing_policy::CustomRoutingPolicy, test_utils, tests::TEST_LOGGER, TokioExecutor,
+    };
 
     #[test]
     fn send_future_errors_when_sender_dropped() {
@@ -1582,7 +1590,7 @@ mod tests {
             ..Default::default()
         };
         let partition_count = 3;
-        create_partitioned_topic("public", "default", &topic, partition_count).await;
+        test_utils::create_partitioned_topic("public", "default", &topic, partition_count).await;
 
         let mut producer = pulsar
             .producer()
@@ -1652,7 +1660,7 @@ mod tests {
             ..Default::default()
         };
         let partition_count = 3;
-        create_partitioned_topic("public", "default", &topic, partition_count).await;
+        test_utils::create_partitioned_topic("public", "default", &topic, partition_count).await;
 
         let mut producer = pulsar
             .producer()
@@ -1715,7 +1723,7 @@ mod tests {
             ..Default::default()
         };
         let partition_count = 3;
-        create_partitioned_topic("public", "default", &topic, partition_count).await;
+        test_utils::create_partitioned_topic("public", "default", &topic, partition_count).await;
 
         let mut producer = pulsar
             .producer()
@@ -1755,23 +1763,4 @@ mod tests {
             assert!(send_receipt.producer_id == producer_id);
         }
     }
-
-    async fn create_partitioned_topic(
-        tenant: &str,
-        namespace: &str,
-        topic_name: &str,
-        num_partitions: u32,
-    ) {
-        let create_partitioned_topic_url = format!(
-            "http://127.0.0.1:8080/admin/v2/persistent/{tenant}/{namespace}/{topic_name}/partitions"
-        );
-        let client = reqwest::Client::new();
-        let response = client
-            .put(create_partitioned_topic_url)
-            .json(&num_partitions.to_string())
-            .send()
-            .await
-            .unwrap();
-        assert!(response.status().is_success());
-    }
 }
diff --git a/src/reader.rs b/src/reader.rs
@@ -184,18 +184,21 @@ impl<T: DeserializeMessage, Exe: Executor> Reader<T, Exe> {
 
 #[cfg(test)]
 mod tests {
-    use crate::consumer::{DeadLetterPolicy, InitialPosition, Message};
-    use crate::proto::MessageIdData;
-    use crate::reader::Reader;
-    use crate::{
-        producer, ConsumerOptions, DeserializeMessage, Error, Executor, Payload, Pulsar,
-        SerializeMessage, SubType, TokioExecutor,
-    };
+    use std::time::Duration;
+
     use futures::StreamExt;
     use serde::{Deserialize, Serialize};
-    use std::time::Duration;
     use tokio::time::timeout;
 
+    use crate::{
+        consumer::{DeadLetterPolicy, InitialPosition, Message},
+        producer,
+        proto::MessageIdData,
+        reader::Reader,
+        ConsumerOptions, DeserializeMessage, Error, Executor, Payload, Pulsar, SerializeMessage,
+        SubType, TokioExecutor,
+    };
+
     #[derive(Serialize, Deserialize)]
     struct TestData {
         data: String,
diff --git a/src/test_utils.rs b/src/test_utils.rs
@@ -0,0 +1,45 @@
+#[cfg(test)]
+use crate::{client::Pulsar, executor::TokioExecutor};
+
+/// Wrapper for the Tokio executor
+#[cfg(any(
+    feature = "tokio-runtime",
+    feature = "tokio-rustls-runtime-aws-lc-rs",
+    feature = "tokio-rustls-runtime-ring"
+))]
+#[cfg(test)]
+pub async fn new_pulsar() -> Pulsar<TokioExecutor> {
+    use log::LevelFilter;
+
+    use crate::tests::TEST_LOGGER;
+
+    let _result = log::set_logger(&TEST_LOGGER);
+    log::set_max_level(LevelFilter::Debug);
+
+    Pulsar::builder("pulsar://127.0.0.1:6650", TokioExecutor)
+        .build()
+        .await
+        .unwrap()
+}
+
+#[cfg(test)]
+pub(crate) async fn create_partitioned_topic(
+    tenant: &str,
+    namespace: &str,
+    topic_name: &str,
+    num_partitions: u32,
+) {
+    use reqwest::Client;
+
+    let create_partitioned_topic_url = format!(
+        "http://127.0.0.1:8080/admin/v2/persistent/{tenant}/{namespace}/{topic_name}/partitions"
+    );
+    let client = Client::new();
+    let response = client
+        .put(create_partitioned_topic_url)
+        .json(&num_partitions.to_string())
+        .send()
+        .await
+        .unwrap();
+    assert!(response.status().is_success());
+}
