feat: support blocking send_non_blocking() and other RPCs (#355)

Author: BewareMyPower

src/client.rs

@@ -554,6 +554,7 @@ impl<Exe: Executor> PulsarBuilder<Exe> {
         Ok(self.with_certificate_chain(v))
     }
 
+    /// The internal pending queue size for each producer on a topic partition. (default: 100)
     #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
     pub fn with_outbound_channel_size(mut self, size: usize) -> Self {
         self.outbound_channel_size = Some(size);

src/connection.rs

@@ -319,12 +319,13 @@ impl<Exe: Executor> ConnectionSender<Exe> {
     }
 
     #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
-    pub(crate) fn send(
+    pub(crate) async fn send(
         &self,
         producer_id: u64,
         producer_name: String,
         sequence_id: u64,
         message: producer::ProducerMessage,
+        block_if_queue_full: bool,
     ) -> Result<
         impl Future<Output = Result<proto::CommandSendReceipt, ConnectionError>>,
         ConnectionError,
@@ -334,7 +335,13 @@ impl<Exe: Executor> ConnectionSender<Exe> {
             sequence_id,
         };
         let msg = messages::send(producer_id, producer_name, sequence_id, message);
-        self.send_message_non_blocking(msg, key, |resp| resp.command.send_receipt)
+        self.send_message_non_blocking(
+            msg,
+            key,
+            |resp| resp.command.send_receipt,
+            block_if_queue_full,
+        )
+        .await
     }
 
     #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
@@ -633,15 +640,22 @@ impl<Exe: Executor> ConnectionSender<Exe> {
     where
         F: FnOnce(Message) -> Option<R> + 'static,
     {
-        self.send_message_non_blocking(msg, key, extract)?.await
+        // This method is called for RPCs other than CommandSend. If the queue is full due to too
+        // many CommandSend RPCs not processed in time, we should wait rather than fail fast
+        // because it's a client side issue that can be recovered later. Hence, set
+        // block_if_queue_full to true here.
+        self.send_message_non_blocking(msg, key, extract, true)
+            .await?
+            .await
     }
 
     #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
-    fn send_message_non_blocking<R: Debug, F>(
+    async fn send_message_non_blocking<R: Debug, F>(
         &self,
         msg: Message,
         key: RequestKey,
         extract: F,
+        block_if_queue_full: bool,
     ) -> Result<impl Future<Output = Result<R, ConnectionError>>, ConnectionError>
     where
         F: FnOnce(Message) -> Option<R> + 'static,
@@ -668,62 +682,76 @@ impl<Exe: Executor> ConnectionSender<Exe> {
                 })?
         };
 
-        match (
-            self.registrations
-                .unbounded_send(Register::Request { key, resolver }),
-            self.tx.try_send(msg),
-        ) {
-            (Ok(_), Ok(_)) => {
-                let connection_id = self.connection_id;
-                let error = self.error.clone();
-                let delay_f = self.executor.delay(self.operation_timeout);
-                trace!(
-                    "Create timeout futures with operation timeout at {:?}",
-                    self.operation_timeout
+        self.registrations
+            .unbounded_send(Register::Request { key, resolver })
+            .map_err(|e| {
+                warn!(
+                    "connection {} disconnected when sending the Request: {}",
+                    self.connection_id, e
                 );
-                let fut = async move {
-                    pin_mut!(response);
-                    pin_mut!(delay_f);
-                    match select(response, delay_f).await {
-                        Either::Left((res, _)) => {
-                            debug!("Received response: {:?}", res);
-                            res
-                        }
-                        Either::Right(_) => {
-                            warn!(
-                                "connection {} timedout sending message to the Pulsar server",
-                                connection_id
-                            );
-                            error.set(ConnectionError::Io(std::io::Error::new(
-                                std::io::ErrorKind::TimedOut,
-                                format!(
-                                    " connection {} timedout sending message to the Pulsar server",
-                                    connection_id
-                                ),
-                            )));
-                            Err(ConnectionError::Io(std::io::Error::new(
-                                std::io::ErrorKind::TimedOut,
-                                format!(
-                                    " connection {} timedout sending message to the Pulsar server",
-                                    connection_id
-                                ),
-                            )))
-                        }
-                    }
-                };
-
-                Ok(fut)
-            }
-            (_, Err(e)) if e.is_full() => Err(ConnectionError::SlowDown),
-            _ => {
+                ConnectionError::Disconnected
+            })?;
+        if block_if_queue_full {
+            self.tx.send(msg).await.map_err(|e| {
                 warn!(
-                    "connection {} disconnected sending message to the Pulsar server",
-                    self.connection_id
+                    "connection {} disconnected when sending the message: {}",
+                    self.connection_id, e
                 );
-                self.error.set(ConnectionError::Disconnected);
-                Err(ConnectionError::Disconnected)
+                ConnectionError::Disconnected
+            })?;
+        } else {
+            self.tx.try_send(msg).map_err(|e| {
+                if e.is_full() {
+                    ConnectionError::SlowDown
+                } else {
+                    warn!(
+                        "connection {} disconnected when sending the message: {}",
+                        self.connection_id, e
+                    );
+                    ConnectionError::Disconnected
+                }
+            })?;
+        };
+
+        let connection_id = self.connection_id;
+        let error = self.error.clone();
+        let delay_f = self.executor.delay(self.operation_timeout);
+        trace!(
+            "Create timeout futures with operation timeout at {:?}",
+            self.operation_timeout
+        );
+        let fut = async move {
+            pin_mut!(response);
+            pin_mut!(delay_f);
+            match select(response, delay_f).await {
+                Either::Left((res, _)) => {
+                    debug!("Received response: {:?}", res);
+                    res
+                }
+                Either::Right(_) => {
+                    warn!(
+                        "connection {} timedout sending message to the Pulsar server",
+                        connection_id
+                    );
+                    error.set(ConnectionError::Io(std::io::Error::new(
+                        std::io::ErrorKind::TimedOut,
+                        format!(
+                            " connection {} timedout sending message to the Pulsar server",
+                            connection_id
+                        ),
+                    )));
+                    Err(ConnectionError::Io(std::io::Error::new(
+                        std::io::ErrorKind::TimedOut,
+                        format!(
+                            " connection {} timedout sending message to the Pulsar server",
+                            connection_id
+                        ),
+                    )))
+                }
             }
-        }
+        };
+
+        Ok(fut)
     }
 
     /// wait for desired message(commandproducersuccess with ready field true)

src/producer.rs

@@ -157,6 +157,10 @@ pub struct ProducerOptions {
     /// producer access mode: shared = 0, exclusive = 1, waitforexclusive =2,
     /// exclusivewithoutfencing =3
     pub access_mode: Option<i32>,
+    /// Whether to block if the internal pending queue, whose size is configured by
+    /// [`crate::client::PulsarBuilder::with_outbound_channel_size`] is full, when awaiting
+    /// [`Producer::send_non_blocking`]. (default: false)
+    pub block_queue_if_full: bool,
 }
 
 impl ProducerOptions {
@@ -378,6 +382,26 @@ impl<Exe: Executor> Producer<Exe> {
     /// - the producer is batching messages, so this function must return immediately, and the
     ///   receipt will come when the batched messages are actually sent
     ///
+    /// If [`ProducerOptions::block_queue_if_full`] is false (by default) and the internal pending
+    /// queue is full, which means the send rate is too fast,
+    /// [`crate::error::ConnectionError::SlowDown`] will be returned. You should handle the error
+    /// like:
+    ///
+    /// ```rust,no_run
+    /// use pulsar::error::{ConnectionError, Error, ProducerError};
+    ///
+    /// # async fn run(mut producer: pulsar::Producer<pulsar::TokioExecutor>) -> Result<(), pulsar::Error> {
+    /// match producer.send_non_blocking("msg").await {
+    ///     Ok(future) => { /* handle the send future */ }
+    ///     Err(Error::Producer(ProducerError::Connection(ConnectionError::SlowDown))) => {
+    ///         /* wait for a while and resent */
+    ///     }
+    ///     Err(e) => { /* handle other errors */ }
+    /// }
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
     /// Usage:
     ///
     /// ```rust,no_run
@@ -818,12 +842,17 @@ where
 {
     loop {
         let message = message.clone();
-        match connection.sender().send(
-            producer_id,
-            producer_name.clone(),
-            sequence_id.get(),
-            message,
-        ) {
+        match connection
+            .sender()
+            .send(
+                producer_id,
+                producer_name.clone(),
+                sequence_id.get(),
+                message,
+                options.block_queue_if_full,
+            )
+            .await
+        {
             Ok(fut) => {
                 let fut = async move {
                     let res = fut.await;
@@ -1379,8 +1408,10 @@ impl<T: SerializeMessage + Sized, Exe: Executor> MessageBuilder<'_, T, Exe> {
 #[cfg(test)]
 mod tests {
     use futures::executor::block_on;
+    use log::LevelFilter;
 
     use super::*;
+    use crate::{tests::TEST_LOGGER, TokioExecutor};
 
     #[test]
     fn send_future_errors_when_sender_dropped() {
@@ -1436,4 +1467,58 @@ mod tests {
         assert!(pm.compression.is_none());
         assert!(pm.uncompressed_size.is_none());
     }
+
+    #[tokio::test]
+    async fn block_if_queue_full() {
+        let _result = log::set_logger(&TEST_LOGGER);
+        log::set_max_level(LevelFilter::Debug);
+        let pulsar: Pulsar<_> = Pulsar::builder("pulsar://127.0.0.1:6650", TokioExecutor)
+            .with_outbound_channel_size(3)
+            .build()
+            .await
+            .unwrap();
+        let mut producer = pulsar
+            .producer()
+            .with_topic(format!("block_queue_if_full_{}", rand::random::<u16>()))
+            .build()
+            .await
+            .unwrap();
+        let mut send_results = Vec::with_capacity(10);
+        for i in 0..10 {
+            send_results.push(producer.send_non_blocking(format!("msg-{i}")).await);
+        }
+        let mut failed_indexes = vec![];
+        for (i, result) in send_results.into_iter().enumerate() {
+            match result {
+                Ok(_) => {}
+                Err(Error::Producer(ProducerError::Connection(ConnectionError::SlowDown))) => {
+                    failed_indexes.push(i);
+                }
+                Err(e) => panic!("failed to send {}: {}", i, e),
+            }
+        }
+        info!("Messages failed due to SlowDown: {:?}", &failed_indexes);
+        assert!(!failed_indexes.is_empty());
+
+        let mut producer = pulsar
+            .producer()
+            .with_topic(format!("block_queue_if_full_{}", rand::random::<u16>()))
+            .with_options(ProducerOptions {
+                block_queue_if_full: true,
+                ..Default::default()
+            })
+            .build()
+            .await
+            .unwrap();
+        let mut send_results = Vec::with_capacity(10);
+        for i in 0..10 {
+            send_results.push(producer.send_non_blocking(format!("msg-{i}")).await);
+        }
+        for (i, result) in send_results.into_iter().enumerate() {
+            match result {
+                Ok(_) => {}
+                Err(e) => panic!("failed to send {}: {}", i, e),
+            }
+        }
+    }
 }