feat: Replace unbounded channel with bounded (#312)

Co-authored-by: Sirius Fang <[email protected]>

Cargo.toml

@@ -16,6 +16,7 @@ description = "Rust client for Apache Pulsar"
 keywords = ["pulsar", "api", "client"]
 
 [dependencies]
+async-channel = "2"
 bytes = "^1.4.0"
 crc = "^3.0.1"
 nom = { version="^7.1.3", default-features=false, features=["alloc"] }

src/client.rs

@@ -179,17 +179,20 @@ impl<Exe: Executor> Pulsar<Exe> {
  connection_retry_parameters: Option<ConnectionRetryOptions>,
  operation_retry_parameters: Option<OperationRetryOptions>,
  tls_options: Option<TlsOptions>,
+ outbound_channel_size: Option<usize>,
  executor: Exe,
  ) -> Result<Self, Error> {
  let url: String = url.into();
  let executor = Arc::new(executor);
  let operation_retry_options = operation_retry_parameters.unwrap_or_default();
+ let outbound_channel_size = outbound_channel_size.unwrap_or(100);
  let manager = ConnectionManager::new(
  url,
  auth,
  connection_retry_parameters,
  operation_retry_options.clone(),
  tls_options,
+ outbound_channel_size,
  executor.clone(),
  )
  .await?;
@@ -252,6 +255,7 @@ impl<Exe: Executor> Pulsar<Exe> {
  connection_retry_options: None,
  operation_retry_options: None,
  tls_options: None,
+ outbound_channel_size: None,
  executor,
  }
  }
@@ -452,6 +456,7 @@ pub struct PulsarBuilder<Exe: Executor> {
  connection_retry_options: Option<ConnectionRetryOptions>,
  operation_retry_options: Option<OperationRetryOptions>,
  tls_options: Option<TlsOptions>,
+ outbound_channel_size: Option<usize>,
  executor: Exe,
 }
 
@@ -549,6 +554,12 @@ impl<Exe: Executor> PulsarBuilder<Exe> {
  Ok(self.with_certificate_chain(v))
  }
 
+ #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
+ pub fn with_outbound_channel_size(mut self, size: usize) -> Result<Self, std::io::Error> {
+ self.outbound_channel_size = Some(size);
+ Ok(self)
+ }
+
  /// creates the Pulsar client and connects it
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
  pub async fn build(self) -> Result<Pulsar<Exe>, Error> {
@@ -558,6 +569,7 @@ impl<Exe: Executor> PulsarBuilder<Exe> {
  connection_retry_options,
  operation_retry_options,
  tls_options,
+ outbound_channel_size,
  executor,
  } = self;
 
@@ -567,6 +579,7 @@ impl<Exe: Executor> PulsarBuilder<Exe> {
  connection_retry_options,
  operation_retry_options,
  tls_options,
+ outbound_channel_size,
  executor,
  )
  .await

src/connection.rs

@@ -99,7 +99,7 @@ impl crate::authentication::Authentication for Authentication {
 
 pub(crate) struct Receiver<S: Stream<Item = Result<Message, ConnectionError>>> {
  inbound: Pin<Box<S>>,
- outbound: mpsc::UnboundedSender<Message>,
+ outbound: async_channel::Sender<Message>,
  error: SharedError,
  pending_requests: BTreeMap<RequestKey, oneshot::Sender<Message>>,
  consumers: BTreeMap<u64, mpsc::UnboundedSender<Message>>,
@@ -114,7 +114,7 @@ impl<S: Stream<Item = Result<Message, ConnectionError>>> Receiver<S> {
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
  pub fn new(
  inbound: S,
- outbound: mpsc::UnboundedSender<Message>,
+ outbound: async_channel::Sender<Message>,
  error: SharedError,
  registrations: mpsc::UnboundedReceiver<Register>,
  shutdown: oneshot::Receiver<()>,
@@ -187,7 +187,9 @@ impl<S: Stream<Item = Result<Message, ConnectionError>>> Future for Receiver<S>
  command: BaseCommand { ping: Some(_), .. },
  ..
  } => {
- let _ = self.outbound.unbounded_send(messages::pong());
+ if let Err(e) = self.outbound.try_send(messages::pong()) {
+ error!("failed to send pong: {}", e);
+ }
  }
  Message {
  command: BaseCommand { pong: Some(_), .. },
@@ -289,7 +291,7 @@ impl SerialId {
 //#[derive(Clone)]
 pub struct ConnectionSender<Exe: Executor> {
  connection_id: Uuid,
- tx: mpsc::UnboundedSender<Message>,
+ tx: async_channel::Sender<Message>,
  registrations: mpsc::UnboundedSender<Register>,
  receiver_shutdown: Option<oneshot::Sender<()>>,
  request_id: SerialId,
@@ -302,7 +304,7 @@ impl<Exe: Executor> ConnectionSender<Exe> {
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
  pub(crate) fn new(
  connection_id: Uuid,
- tx: mpsc::UnboundedSender<Message>,
+ tx: async_channel::Sender<Message>,
  registrations: mpsc::UnboundedSender<Register>,
  receiver_shutdown: oneshot::Sender<()>,
  request_id: SerialId,
@@ -349,9 +351,9 @@ impl<Exe: Executor> ConnectionSender<Exe> {
  match (
  self.registrations
  .unbounded_send(Register::Ping { resolver }),
- self.tx.unbounded_send(messages::ping()),
+ self.tx.try_send(messages::ping())?,
  ) {
- (Ok(_), Ok(_)) => {
+ (Ok(_), ()) => {
  let delay_f = self.executor.delay(self.operation_timeout);
  pin_mut!(response);
  pin_mut!(delay_f);
@@ -526,8 +528,8 @@ impl<Exe: Executor> ConnectionSender<Exe> {
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
  pub fn send_flow(&self, consumer_id: u64, message_permits: u32) -> Result<(), ConnectionError> {
  self.tx
- .unbounded_send(messages::flow(consumer_id, message_permits))
-  .map_err(|_| ConnectionError::Disconnected)
+ .try_send(messages::flow(consumer_id, message_permits))?;
+ Ok(())
  }
 
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
@@ -538,8 +540,8 @@ impl<Exe: Executor> ConnectionSender<Exe> {
  cumulative: bool,
  ) -> Result<(), ConnectionError> {
  self.tx
- .unbounded_send(messages::ack(consumer_id, message_ids, cumulative))
-  .map_err(|_| ConnectionError::Disconnected)
+ .try_send(messages::ack(consumer_id, message_ids, cumulative))?;
+ Ok(())
  }
 
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
@@ -549,11 +551,11 @@ impl<Exe: Executor> ConnectionSender<Exe> {
  message_ids: Vec<proto::MessageIdData>,
  ) -> Result<(), ConnectionError> {
  self.tx
- .unbounded_send(messages::redeliver_unacknowleged_messages(
+ .try_send(messages::redeliver_unacknowleged_messages(
  consumer_id,
  message_ids,
- ))
-  .map_err(|_| ConnectionError::Disconnected)
+ ))?;
+ Ok(())
  }
 
  #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
@@ -661,7 +663,7 @@ impl<Exe: Executor> ConnectionSender<Exe> {
  match (
  self.registrations
  .unbounded_send(Register::Request { key, resolver }),
- self.tx.unbounded_send(msg),
+ self.tx.try_send(msg),
  ) {
  (Ok(_), Ok(_)) => {
  let connection_id = self.connection_id;
@@ -700,6 +702,7 @@ impl<Exe: Executor> ConnectionSender<Exe> {
 
  Ok(fut)
  }
+ (_, Err(e)) if e.is_full() => Err(ConnectionError::SlowDown),
  _ => {
  warn!(
  "connection {} disconnected sending message to the Pulsar server",
@@ -781,6 +784,7 @@ impl<Exe: Executor> Connection<Exe> {
  tls_hostname_verification_enabled: bool,
  connection_timeout: Duration,
  operation_timeout: Duration,
+ outbound_channel_size: usize,
  executor: Arc<Exe>,
  ) -> Result<Connection<Exe>, ConnectionError> {
  if url.scheme() != "pulsar" && url.scheme() != "pulsar+ssl" {
@@ -839,6 +843,7 @@ impl<Exe: Executor> Connection<Exe> {
  tls_hostname_verification_enabled,
  executor.clone(),
  operation_timeout,
+ outbound_channel_size,
  );
  let delay_f = executor.delay(connection_timeout);
 
@@ -916,6 +921,7 @@ impl<Exe: Executor> Connection<Exe> {
  tls_hostname_verification_enabled: bool,
  executor: Arc<Exe>,
  operation_timeout: Duration,
+ outbound_channel_size: usize,
  ) -> Result<ConnectionSender<Exe>, ConnectionError> {
  match executor.kind() {
  #[cfg(feature = "tokio-runtime")]
@@ -945,6 +951,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  } else {
@@ -959,6 +966,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  }
@@ -1007,6 +1015,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  } else {
@@ -1021,6 +1030,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  }
@@ -1053,6 +1063,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  } else {
@@ -1067,6 +1078,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  }
@@ -1119,6 +1131,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  } else {
@@ -1133,6 +1146,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url,
  executor,
  operation_timeout,
+ outbound_channel_size,
  )
  .await
  }
@@ -1155,6 +1169,7 @@ impl<Exe: Executor> Connection<Exe> {
  proxy_to_broker_url: Option<String>,
  executor: Arc<Exe>,
  operation_timeout: Duration,
+ outbound_channel_size: usize,
  ) -> Result<ConnectionSender<Exe>, ConnectionError>
  where
  S: Stream<Item = Result<Message, ConnectionError>>,
@@ -1194,7 +1209,7 @@ impl<Exe: Executor> Connection<Exe> {
  }?;
 
  let (mut sink, stream) = stream.split();
- let (tx, mut rx) = mpsc::unbounded();
+ let (tx, rx) = async_channel::bounded(outbound_channel_size);
  let (registrations_tx, registrations_rx) = mpsc::unbounded();
  let error = SharedError::new();
  let (receiver_shutdown_tx, receiver_shutdown_rx) = oneshot::channel();
@@ -1220,7 +1235,7 @@ impl<Exe: Executor> Connection<Exe> {
 
  let err = error.clone();
  let res = executor.spawn(Box::pin(async move {
- while let Some(msg) = rx.next().await {
+ while let Ok(msg) = rx.recv().await {
  // println!("real sent msg: {:?}", msg);
  if let Err(e) = sink.send(msg).await {
  err.set(e);
@@ -1236,7 +1251,7 @@ impl<Exe: Executor> Connection<Exe> {
  if auth.is_some() {
  let auth_challenge_res = executor.spawn({
  let err = error.clone();
- let mut tx = tx.clone();
+ let tx = tx.clone();
  let auth = auth.clone();
  Box::pin(async move {
  while auth_challenge_rx.next().await.is_some() {
@@ -1796,7 +1811,7 @@ mod tests {
  #[cfg(any(feature = "tokio-runtime", feature = "tokio-rustls-runtime"))]
  async fn receiver_auth_challenge_test() {
  let (message_tx, message_rx) = mpsc::unbounded();
- let (tx, _) = mpsc::unbounded();
+ let (tx, _) = async_channel::bounded(10);
  let (_registrations_tx, registrations_rx) = mpsc::unbounded();
  let error = SharedError::new();
  let (_receiver_shutdown_tx, receiver_shutdown_rx) = oneshot::channel();
@@ -1904,6 +1919,7 @@ mod tests {
  None,
  TokioExecutor.into(),
  Duration::from_secs(10),
+ 100,
  )
  .await;
 

src/connection_manager.rs

@@ -123,6 +123,7 @@ pub struct ConnectionManager<Exe: Executor> {
  pub(crate) operation_retry_options: OperationRetryOptions,
  tls_options: TlsOptions,
  certificate_chain: Vec<Certificate>,
+ outbound_channel_size: usize,
 }
 
 impl<Exe: Executor> ConnectionManager<Exe> {
@@ -133,6 +134,7 @@ impl<Exe: Executor> ConnectionManager<Exe> {
  connection_retry: Option<ConnectionRetryOptions>,
  operation_retry_options: OperationRetryOptions,
  tls: Option<TlsOptions>,
+ outbound_channel_size: usize,
  executor: Arc<Exe>,
  ) -> Result<Self, ConnectionError> {
  let connection_retry_options = connection_retry.unwrap_or_default();
@@ -191,6 +193,7 @@ impl<Exe: Executor> ConnectionManager<Exe> {
  operation_retry_options,
  tls_options,
  certificate_chain,
+ outbound_channel_size,
  };
  let broker_address = BrokerAddress {
  url: url.clone(),
@@ -312,6 +315,7 @@ impl<Exe: Executor> ConnectionManager<Exe> {
  self.tls_options.tls_hostname_verification_enabled,
  self.connection_retry_options.connection_timeout,
  self.operation_retry_options.operation_timeout,
+ self.outbound_channel_size,
  self.executor.clone(),
  )
  .await

src/consumer/engine.rs

@@ -159,6 +159,11 @@ impl<Exe: Executor> ConsumerEngine<Exe> {
  .sender()
  .send_flow(self.id, self.batch_size - self.remaining_messages)?;
  }
+ Err(ConnectionError::SlowDown) => {
+ self.connection
+ .sender()
+ .send_flow(self.id, self.batch_size - self.remaining_messages)?;
+ }
  Err(e) => return Err(e.into()),
  }
  self.remaining_messages = self.batch_size;