diff --git a/README.md b/README.md index 641eaff9..de6ee932 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,8 @@ Features: - Automatic reconnection with exponential back off; - Message batching; - Compression with LZ4, zlib, zstd or Snappy (can be deactivated with Cargo features); -- Telemetry using [tracing](https://github.com/tokio-rs/tracing) crate (can be activated with Cargo features). +- Telemetry using [tracing](https://github.com/tokio-rs/tracing) crate (can be activated with Cargo features); +- Configurable negative-acknowledgment redelivery delay with pluggable backoff policies ## Getting Started @@ -31,6 +32,7 @@ Try out [examples](examples): - [producer](examples/producer.rs) - [consumer](examples/consumer.rs) - [reader](examples/reader.rs) +- [negative_ack](examples/negative_ack.rs) ## Project Maintainers diff --git a/examples/negative_ack.rs b/examples/negative_ack.rs new file mode 100644 index 00000000..dd7d41ec --- /dev/null +++ b/examples/negative_ack.rs @@ -0,0 +1,103 @@ +#[macro_use] +extern crate serde; +use std::{env, time::Duration}; + +use futures::TryStreamExt; +use pulsar::{ + authentication::{basic::BasicAuthentication, oauth2::OAuth2Authentication}, + Authentication, Consumer, DeserializeMessage, MultiplierRedeliveryBackoff, Payload, Pulsar, + SubType, TokioExecutor, +}; + +#[derive(Serialize, Deserialize)] +struct TestData { + data: String, +} + +impl DeserializeMessage for TestData { + type Output = Result; + + fn deserialize_message(payload: &Payload) -> Self::Output { + serde_json::from_slice(&payload.data) + } +} + +#[tokio::main] +async fn main() -> Result<(), pulsar::Error> { + env_logger::init(); + + let addr = env::var("PULSAR_ADDRESS") + .ok() + .unwrap_or_else(|| "pulsar://127.0.0.1:6650".to_string()); + let topic = env::var("PULSAR_TOPIC") + .ok() + .unwrap_or_else(|| "non-persistent://public/default/test".to_string()); + + let mut builder = Pulsar::builder(addr, TokioExecutor); + + if let Ok(token) = env::var("PULSAR_TOKEN") { + let authentication = Authentication { + name: "token".to_string(), + data: token.into_bytes(), + }; + + builder = builder.with_auth(authentication); + } else if let Ok(oauth2_cfg) = env::var("PULSAR_OAUTH2") { + builder = builder.with_auth_provider(OAuth2Authentication::client_credentials( + serde_json::from_str(oauth2_cfg.as_str()) + .unwrap_or_else(|_| panic!("invalid oauth2 config [{}]", oauth2_cfg.as_str())), + )); + } else if let (Ok(username), Ok(password)) = ( + env::var("PULSAR_BASIC_USERNAME"), + env::var("PULSAR_BASIC_PASSWORD"), + ) { + builder = builder.with_auth_provider(BasicAuthentication::new(&username, &password)) + } + + let pulsar: Pulsar<_> = builder.build().await?; + let negative_ack_backoff = MultiplierRedeliveryBackoff::builder() + .min_delay(Duration::from_secs(1)) + .max_delay(Duration::from_secs(30)) + .multiplier(2.0) + .build()?; + + let mut consumer: Consumer = pulsar + .consumer() + .with_topic(topic) + .with_consumer_name("negative_ack_consumer") + .with_subscription_type(SubType::Exclusive) + .with_subscription("negative_ack_subscription") + .with_nack_redelivery_delay(Duration::from_secs(5)) + .with_negative_ack_backoff(negative_ack_backoff) + .build() + .await?; + + let mut counter = 0usize; + while let Some(msg) = consumer.try_next().await? { + log::info!("metadata: {:?}", msg.metadata()); + log::info!("id: {:?}", msg.message_id()); + log::info!("redelivery_count: {}", msg.redelivery_count()); + + let data = match msg.deserialize() { + Ok(data) => data, + Err(e) => { + log::error!("could not deserialize message: {:?}", e); + consumer.nack(&msg).await?; + break; + } + }; + + log::info!("payload data: {}", data.data); + consumer.nack(&msg).await?; + + counter += 1; + log::info!("nacked {} messages", counter); + + if counter > 10 { + consumer.close().await.expect("Unable to close consumer"); + break; + } + } + + Ok(()) +} diff --git a/src/consumer/builder.rs b/src/consumer/builder.rs index 8d734e3b..1ca6e23a 100644 --- a/src/consumer/builder.rs +++ b/src/consumer/builder.rs @@ -1,6 +1,7 @@ use std::{ collections::{BTreeMap, VecDeque}, - time::Duration, + sync::Arc, + time::{Duration, Instant}, }; use futures::{future::try_join_all, StreamExt}; @@ -10,13 +11,19 @@ use regex::Regex; use crate::{ consumer::{ config::ConsumerConfig, data::DeadLetterPolicy, multi::MultiTopicConsumer, - options::ConsumerOptions, topic::TopicConsumer, InnerConsumer, + negative_ack_backoff::NegativeAckBackoff, options::ConsumerOptions, topic::TopicConsumer, + InnerConsumer, }, message::proto::command_subscribe::SubType, reader::{Reader, State}, BrokerAddress, Consumer, DeserializeMessage, Error, Executor, Pulsar, }; +const NACK_REDELIVERY_DELAY_EXCEEDS_MILLISECONDS_RANGE: &str = + "negative-ack redelivery delay is too large: it must be less than u64::MAX milliseconds"; +const NACK_REDELIVERY_DELAY_EXCEEDS_INSTANT_RANGE: &str = + "negative-ack redelivery delay is too large to schedule from the current instant"; + /// Builder structure for consumers /// /// This is the main way to create a [Consumer] or a [Reader] @@ -35,6 +42,24 @@ pub struct ConsumerBuilder { consumer_options: Option, namespace: Option, topic_refresh: Option, + nack_redelivery_delay: Option, + negative_ack_backoff: Option>, +} + +fn check_nack_delay_duration(delay: Duration) -> Result<(), Error> { + if delay.as_millis() >= u64::MAX as u128 { + return Err(Error::Custom(format!( + "{NACK_REDELIVERY_DELAY_EXCEEDS_MILLISECONDS_RANGE}: {delay:?}" + ))); + } + + if Instant::now().checked_add(delay).is_none() { + return Err(Error::Custom(format!( + "{NACK_REDELIVERY_DELAY_EXCEEDS_INSTANT_RANGE}: {delay:?}" + ))); + } + + Ok(()) } impl ConsumerBuilder { @@ -56,6 +81,8 @@ impl ConsumerBuilder { consumer_options: None, namespace: None, topic_refresh: None, + nack_redelivery_delay: None, + negative_ack_backoff: None, } } @@ -175,6 +202,43 @@ impl ConsumerBuilder { self } + /// Sets the fixed delay before redelivering negatively acknowledged messages. + /// + /// If this option is not set, negative acknowledgments use a 60 seconds default delay before + /// broker redelivery. + /// + /// `Duration::ZERO` is valid and requests immediate redelivery. Repeated calls are + /// last-call-wins. + /// + /// The configured fixed delay is used for ID-only [`Self::nack_with_id`] calls, which carry no + /// redelivery count. It is also used as the delay for message-based [`Consumer::nack`] calls + /// when no backoff policy has been configured. + #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] + pub fn with_nack_redelivery_delay(mut self, delay: Duration) -> Self { + self.nack_redelivery_delay = Some(delay); + self + } + + /// Configures an optional negative-ack backoff policy. + /// + /// Message-based negative acknowledgments use the message's redelivery count (`0` on first + /// delivery) to compute the backoff delay. ID-only `nack_with_id` calls have no redelivery + /// count, so they use + /// [`Self::with_nack_redelivery_delay`] or the 60 seconds default instead. + /// + /// A backoff that returns `Duration::ZERO` requests immediate redelivery. + /// + /// This policy affects only negative acknowledgments. Unacked-message redelivery and + /// dead-letter policy behavior remain configured separately. + #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] + pub fn with_negative_ack_backoff(mut self, backoff: B) -> Self + where + B: NegativeAckBackoff + 'static, + { + self.negative_ack_backoff = Some(Arc::new(backoff) as Arc); + self + } + // Checks the builder for inconsistencies // returns a config and a list of topics with associated brokers #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] @@ -193,6 +257,8 @@ impl ConsumerBuilder { dead_letter_policy, namespace: _, topic_refresh: _, + nack_redelivery_delay, + negative_ack_backoff, } = self; if consumer_name.is_none() { @@ -208,6 +274,9 @@ impl ConsumerBuilder { "Cannot create consumer with no topics and no topic regex".into(), )); } + if let Some(delay) = nack_redelivery_delay { + check_nack_delay_duration(delay)?; + } let topics: Vec<(String, BrokerAddress)> = try_join_all( topics @@ -263,6 +332,8 @@ impl ConsumerBuilder { unacked_message_redelivery_delay: unacked_message_resend_delay, options: consumer_options.unwrap_or_default(), dead_letter_policy, + nack_redelivery_delay, + negative_ack_backoff, }; Ok((config, topics)) } @@ -347,3 +418,102 @@ impl ConsumerBuilder { }) } } + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use super::*; + + #[test] + fn test_nack_delay_validation_oversized() { + let delay = Duration::MAX; + let result = check_nack_delay_duration(delay); + + assert_nack_delay_error( + result, + NACK_REDELIVERY_DELAY_EXCEEDS_MILLISECONDS_RANGE, + delay, + ); + } + + #[test] + fn test_nack_delay_validation_rejects_millisecond_range_boundary() { + let delay = Duration::from_millis(u64::MAX); + let result = check_nack_delay_duration(delay); + + assert_nack_delay_error( + result, + NACK_REDELIVERY_DELAY_EXCEEDS_MILLISECONDS_RANGE, + delay, + ); + } + + fn assert_nack_delay_error(result: Result<(), Error>, expected_prefix: &str, delay: Duration) { + match result { + Err(Error::Custom(message)) => { + assert!( + message.starts_with(expected_prefix), + "expected error to start with {expected_prefix:?}, got {message:?}" + ); + assert!( + message.contains(&format!("{delay:?}")), + "expected error to include the rejected Duration, got {message:?}" + ); + } + other => panic!("expected Error::Custom({expected_prefix:?}), got {other:?}"), + } + } + + #[test] + fn test_nack_delay_validation_zero_is_valid() { + let result = check_nack_delay_duration(Duration::ZERO); + + assert!(result.is_ok()); + } + + #[test] + fn test_consumer_config_nack_delay_default_is_none() { + let config = ConsumerConfig::default(); + + assert!(config.nack_redelivery_delay.is_none()); + assert!(config.negative_ack_backoff.is_none()); + } + + #[test] + fn test_consumer_config_stores_nack_delay() { + let config = ConsumerConfig { + nack_redelivery_delay: Some(Duration::from_secs(5)), + ..ConsumerConfig::default() + }; + + assert_eq!(config.nack_redelivery_delay, Some(Duration::from_secs(5))); + } + + #[test] + fn test_consumer_config_stores_both_nack_configs_independently() { + let config = ConsumerConfig { + nack_redelivery_delay: Some(Duration::from_secs(5)), + negative_ack_backoff: Some(Arc::new( + crate::consumer::negative_ack_backoff::MultiplierRedeliveryBackoff::default(), + )), + ..ConsumerConfig::default() + }; + + assert!(config.nack_redelivery_delay.is_some()); + assert!(config.negative_ack_backoff.is_some()); + } + + #[test] + fn test_consumer_config_backoff_only_stores_none_for_delay() { + let config = ConsumerConfig { + negative_ack_backoff: Some(Arc::new( + crate::consumer::negative_ack_backoff::MultiplierRedeliveryBackoff::default(), + )), + ..ConsumerConfig::default() + }; + + assert!(config.nack_redelivery_delay.is_none()); + assert!(config.negative_ack_backoff.is_some()); + } +} diff --git a/src/consumer/config.rs b/src/consumer/config.rs index 8fba5fbc..7426aeec 100644 --- a/src/consumer/config.rs +++ b/src/consumer/config.rs @@ -1,12 +1,14 @@ -use std::time::Duration; +use std::{fmt, sync::Arc, time::Duration}; use crate::{ - consumer::{data::DeadLetterPolicy, options::ConsumerOptions}, + consumer::{ + data::DeadLetterPolicy, negative_ack_backoff::NegativeAckBackoff, options::ConsumerOptions, + }, message::proto::command_subscribe::SubType, }; /// the complete configuration of a consumer -#[derive(Debug, Clone, Default)] +#[derive(Clone, Default)] pub struct ConsumerConfig { /// subscription name pub(crate) subscription: String, @@ -28,4 +30,31 @@ pub struct ConsumerConfig { pub(crate) options: ConsumerOptions, /// dead letter policy pub(crate) dead_letter_policy: Option, + /// fixed negative-ack redelivery delay + pub(crate) nack_redelivery_delay: Option, + /// negative-ack redelivery backoff policy + pub(crate) negative_ack_backoff: Option>, +} + +impl fmt::Debug for ConsumerConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ConsumerConfig") + .field("subscription", &self.subscription) + .field("sub_type", &self.sub_type) + .field("batch_size", &self.batch_size) + .field("consumer_name", &self.consumer_name) + .field("consumer_id", &self.consumer_id) + .field( + "unacked_message_redelivery_delay", + &self.unacked_message_redelivery_delay, + ) + .field("options", &self.options) + .field("dead_letter_policy", &self.dead_letter_policy) + .field("nack_redelivery_delay", &self.nack_redelivery_delay) + .field( + "negative_ack_backoff", + &self.negative_ack_backoff.as_ref().map(|_| "configured"), + ) + .finish() + } } diff --git a/src/consumer/data.rs b/src/consumer/data.rs index 38200c00..aa52018c 100644 --- a/src/consumer/data.rs +++ b/src/consumer/data.rs @@ -9,6 +9,9 @@ use crate::{ }; pub type MessageIdDataReceiver = mpsc::Receiver>; +pub(crate) type InternalMessageIdDataResult = Result<(MessageIdData, Payload, Option), Error>; +pub(crate) type InternalMessageIdDataSender = mpsc::Sender; +pub(crate) type InternalMessageIdDataReceiver = mpsc::Receiver; pub enum EngineEvent { Message(Option), @@ -22,6 +25,20 @@ pub enum EngineMessage { GetConnection(oneshot::Sender>>), } +pub(crate) enum InternalEngineEvent { + Message(Option), + EngineMessage(Option>), +} + +pub(crate) enum InternalEngineMessage { + Ack(MessageIdData, bool), + Nack(MessageIdData, Option), + NegativeAckRedelivery, + UnackedRedelivery, + GetConnection(oneshot::Sender>>), + ClearNegativeAcks, +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct MessageData { pub id: MessageIdData, @@ -36,3 +53,90 @@ pub struct DeadLetterPolicy { /// Name of the dead topic where the failing messages will be sent. pub dead_letter_topic: String, } + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::proto; + + fn message_id() -> proto::MessageIdData { + proto::MessageIdData { + ledger_id: 1, + entry_id: 1, + partition: None, + batch_index: None, + ack_set: vec![], + batch_size: None, + first_chunk_message_id: None, + } + } + + #[test] + fn test_nack_engine_message_with_count() { + let message = InternalEngineMessage::::Nack( + message_id(), + Some(42u32), + ); + + match message { + InternalEngineMessage::Nack(_, count) => assert_eq!(count, Some(42u32)), + _ => panic!("expected nack engine message"), + } + } + + #[test] + fn test_nack_engine_message_id_only_carries_none() { + let message = + InternalEngineMessage::::Nack(message_id(), None); + + match message { + InternalEngineMessage::Nack(_, count) => assert_eq!(count, None), + _ => panic!("expected nack engine message"), + } + } + + #[test] + fn public_nack_engine_message_shape_stays_id_only() { + let message = EngineMessage::::Nack(message_id()); + + match message { + EngineMessage::Nack(_) => {} + _ => panic!("expected public nack engine message"), + } + } + + #[test] + fn public_nack_engine_message_accepts_default_message_id_with_one_field() { + let message = + EngineMessage::::Nack(MessageIdData::default()); + + match message { + EngineMessage::Nack(_) => {} + _ => panic!("expected public nack engine message"), + } + } + + #[test] + fn internal_engine_message_supports_private_negative_ack_due_event() { + let message = + InternalEngineMessage::::NegativeAckRedelivery; + + match message { + InternalEngineMessage::NegativeAckRedelivery => {} + _ => panic!("expected private negative ack redelivery event"), + } + } + + #[test] + fn public_receiver_alias_preserves_upstream_two_tuple_shape() { + type ExpectedPublicReceiver = mpsc::Receiver>; + + let _: Option = None::; + } + + #[test] + fn internal_receiver_alias_preserves_optional_redelivery_count_shape() { + let _: Option = + None::>; + } +} diff --git a/src/consumer/engine.rs b/src/consumer/engine.rs index 46144205..277bd7a8 100644 --- a/src/consumer/engine.rs +++ b/src/consumer/engine.rs @@ -1,7 +1,10 @@ use std::{ collections::{HashMap, HashSet}, future::Future, - sync::Arc, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, time::{Duration, Instant}, }; @@ -14,7 +17,14 @@ use crate::{ connection::Connection, consumer::{ batched_message_iterator::BatchedMessageIterator, - data::{DeadLetterPolicy, EngineEvent, EngineMessage}, + data::{ + DeadLetterPolicy, InternalEngineEvent, InternalEngineMessage, + InternalMessageIdDataSender, + }, + negative_ack_backoff::NegativeAckBackoff, + negative_ack_tracker::{ + NegativeAckSchedule, NegativeAckTracker, NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + }, options::ConsumerOptions, }, error::{ConnectionError, ConsumerError}, @@ -32,6 +42,72 @@ use crate::{ const SYSTEM_PROPERTY_REAL_TOPIC: &str = "REAL_TOPIC"; const SYSTEM_PROPERTY_ORIGIN_MESSAGE_ID: &str = "ORIGIN_MESSAGE_ID"; +#[async_trait::async_trait] +trait NegativeAckRedeliveryRequester { + async fn request_negative_ack_redelivery( + &mut self, + consumer_id: u64, + ids: Vec, + ) -> Result<(), ConnectionError>; +} + +struct BrokerNegativeAckRedeliveryRequester { + connection: Arc>, +} + +#[async_trait::async_trait] +impl NegativeAckRedeliveryRequester for BrokerNegativeAckRedeliveryRequester { + async fn request_negative_ack_redelivery( + &mut self, + consumer_id: u64, + ids: Vec, + ) -> Result<(), ConnectionError> { + self.connection + .sender() + .send_redeliver_unacknowleged_messages(consumer_id, ids) + .await + } +} + +async fn send_negative_ack_redelivery_with_requester( + requester: &mut R, + consumer_id: u64, + ids: Vec, +) -> Result<(), ConnectionError> +where + R: NegativeAckRedeliveryRequester + ?Sized, +{ + let ids = ids + .into_iter() + .map(|id| NegativeAckTracker::redelivery_message_id(&id)) + .collect(); + requester + .request_negative_ack_redelivery(consumer_id, ids) + .await +} + +async fn dispatch_negative_ack_redelivery_with_requester( + tracker: &mut NegativeAckTracker, + requester: &mut R, + consumer_id: u64, + ids: &[MessageIdData], + retry_at: Instant, +) -> Result<(), ConnectionError> +where + R: NegativeAckRedeliveryRequester + ?Sized, +{ + match send_negative_ack_redelivery_with_requester(requester, consumer_id, ids.to_vec()).await { + Ok(()) => { + tracker.mark_dispatched(ids); + Ok(()) + } + Err(e) => { + tracker.mark_retry_pending(ids, retry_at); + Err(e) + } + } +} + pub struct ConsumerEngine { client: Pulsar, connection: Arc>, @@ -40,15 +116,20 @@ pub struct ConsumerEngine { sub_type: SubType, id: u64, name: Option, - tx: mpsc::Sender>, + tx: InternalMessageIdDataSender, messages_rx: Option>, - engine_rx: Option>>, - event_rx: mpsc::UnboundedReceiver>, - event_tx: UnboundedSender>, + engine_rx: Option>>, + event_rx: mpsc::UnboundedReceiver>, + event_tx: UnboundedSender>, batch_size: u32, remaining_messages: i64, unacked_message_redelivery_delay: Option, unacked_messages: HashMap, + nack_redelivery_delay: Option, + negative_ack_backoff: Option>, + negative_ack_tracker: Option, + negative_ack_ticker_running: Option>, + negative_ack_due_event_pending: Arc, dead_letter_policy: Option, options: ConsumerOptions, } @@ -63,11 +144,13 @@ impl ConsumerEngine { sub_type: SubType, id: u64, name: Option, - tx: mpsc::Sender>, + tx: InternalMessageIdDataSender, messages_rx: mpsc::UnboundedReceiver, - engine_rx: mpsc::UnboundedReceiver>, + engine_rx: mpsc::UnboundedReceiver>, batch_size: u32, unacked_message_redelivery_delay: Option, + nack_redelivery_delay: Option, + negative_ack_backoff: Option>, dead_letter_policy: Option, options: ConsumerOptions, ) -> ConsumerEngine { @@ -89,6 +172,11 @@ impl ConsumerEngine { remaining_messages: batch_size as i64, unacked_message_redelivery_delay, unacked_messages: HashMap::new(), + nack_redelivery_delay, + negative_ack_backoff, + negative_ack_tracker: None, + negative_ack_ticker_running: None, + negative_ack_due_event_pending: Arc::new(AtomicBool::new(false)), dead_letter_policy, options, } @@ -97,7 +185,7 @@ impl ConsumerEngine { fn register_source(&self, mut rx: mpsc::UnboundedReceiver, mapper: M) -> Result<(), ()> where E: Send + 'static, - M: Fn(Option) -> EngineEvent + Send + Sync + 'static, + M: Fn(Option) -> InternalEngineEvent + Send + Sync + 'static, { let mut event_tx = self.event_tx.clone(); @@ -138,14 +226,14 @@ impl ConsumerEngine { } if let Some(messages_rx) = self.messages_rx.take() { - self.register_source(messages_rx, |msg| EngineEvent::Message(msg)) + self.register_source(messages_rx, |msg| InternalEngineEvent::Message(msg)) .map_err(|_| { Error::Custom(String::from("Error registering messages_rx source")) })?; } if let Some(engine_rx) = self.engine_rx.take() { - self.register_source(engine_rx, |msg| EngineEvent::EngineMessage(msg)) + self.register_source(engine_rx, |msg| InternalEngineEvent::EngineMessage(msg)) .map_err(|_| { Error::Custom(String::from("Error registering engine_rx source")) })?; @@ -214,14 +302,14 @@ impl ConsumerEngine { // ``` debug!("consumer engine: timeout (1s)"); } - Ok(Some(EngineEvent::Message(msg))) => { + Ok(Some(InternalEngineEvent::Message(msg))) => { debug!("consumer engine: received message, {:?}", msg); let out = self.handle_message_opt(msg).await; if let Some(res) = out { return res; } } - Ok(Some(EngineEvent::EngineMessage(msg))) => { + Ok(Some(InternalEngineEvent::EngineMessage(msg))) => { debug!("consumer engine: received engine message"); let continue_loop = self.handle_ack_opt(msg).await; if !continue_loop { @@ -276,31 +364,25 @@ impl ConsumerEngine { } } - async fn handle_ack_opt(&mut self, ack_opt: Option>) -> bool { + async fn handle_ack_opt(&mut self, ack_opt: Option>) -> bool { match ack_opt { None => { trace!("ack channel was closed"); false } - Some(EngineMessage::Ack(message_id, cumulative)) => { + Some(InternalEngineMessage::Ack(message_id, cumulative)) => { self.ack(message_id, cumulative).await; true } - Some(EngineMessage::Nack(message_id)) => { - if let Err(e) = self - .connection - .sender() - .send_redeliver_unacknowleged_messages(self.id, vec![message_id.clone()]) - .await - { - error!( - "could not ask for redelivery for message {:?}: {:?}", - message_id, e - ); - } + Some(InternalEngineMessage::Nack(message_id, redelivery_count)) => { + self.handle_negative_ack(message_id, redelivery_count).await; true } - Some(EngineMessage::UnackedRedelivery) => { + Some(InternalEngineMessage::NegativeAckRedelivery) => { + self.handle_negative_ack_redelivery().await; + true + } + Some(InternalEngineMessage::UnackedRedelivery) => { let mut h = HashSet::new(); let now = Instant::now(); // info!("unacked messages length: {}", self.unacked_messages.len()); @@ -328,7 +410,7 @@ impl ConsumerEngine { } true } - Some(EngineMessage::GetConnection(sender)) => { + Some(InternalEngineMessage::GetConnection(sender)) => { let _ = sender.send(self.connection.clone()).map_err(|_| { error!( "consumer requested the engine's connection but dropped the \ @@ -337,13 +419,258 @@ impl ConsumerEngine { }); true } + Some(InternalEngineMessage::ClearNegativeAcks) => { + self.clear_negative_ack_state(); + true + } + } + } + + fn make_negative_ack_tracker(&self) -> NegativeAckTracker { + NegativeAckTracker::new( + self.nack_redelivery_delay, + self.negative_ack_backoff.clone(), + ) + } + + fn ensure_negative_ack_tracker(&mut self) -> &mut NegativeAckTracker { + if self.negative_ack_tracker.is_none() { + self.negative_ack_tracker = Some(self.make_negative_ack_tracker()); + } + self.negative_ack_tracker + .as_mut() + .expect("negative ack tracker was just initialized") + } + + fn start_negative_ack_ticker(&mut self) -> Result<(), ()> { + if self + .negative_ack_ticker_running + .as_ref() + .map(|ticker_running| ticker_running.load(Ordering::SeqCst)) + .unwrap_or(false) + { + return Ok(()); + } + + let ticker_running = Arc::new(AtomicBool::new(true)); + self.negative_ack_ticker_running = Some(ticker_running.clone()); + let due_event_pending = self.negative_ack_due_event_pending.clone(); + let mut event_tx = self.event_tx.clone(); + let mut interval = self + .client + .executor + .interval(NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL); + + let spawn_result = self.client.executor.spawn(Box::pin(async move { + while interval.next().await.is_some() { + if !ticker_running.load(Ordering::SeqCst) { + break; + } + + if due_event_pending + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + continue; + } + + if event_tx + .send(InternalEngineEvent::EngineMessage(Some( + InternalEngineMessage::NegativeAckRedelivery, + ))) + .await + .is_err() + { + due_event_pending.store(false, Ordering::SeqCst); + ticker_running.store(false, Ordering::SeqCst); + break; + } + } + })); + + if spawn_result.is_err() { + if let Some(ticker_running) = &self.negative_ack_ticker_running { + ticker_running.store(false, Ordering::SeqCst); + } + self.negative_ack_due_event_pending + .store(false, Ordering::SeqCst); + } + + spawn_result + } + + fn stop_negative_ack_ticker_if_idle(&mut self) { + let tracker_empty = match &self.negative_ack_tracker { + Some(tracker) => tracker.is_empty(), + None => true, + }; + + if tracker_empty { + if let Some(ticker_running) = &self.negative_ack_ticker_running { + ticker_running.store(false, Ordering::SeqCst); + } + self.negative_ack_ticker_running = None; + self.negative_ack_due_event_pending + .store(false, Ordering::SeqCst); + } + } + + fn clear_negative_ack_state(&mut self) { + if let Some(tracker) = self.negative_ack_tracker.as_mut() { + tracker.clear(); + } + if let Some(ticker_running) = &self.negative_ack_ticker_running { + ticker_running.store(false, Ordering::SeqCst); + } + self.negative_ack_ticker_running = None; + self.negative_ack_due_event_pending + .store(false, Ordering::SeqCst); + } + + async fn handle_negative_ack( + &mut self, + message_id: MessageIdData, + redelivery_count: Option, + ) { + self.remove_negative_ack_from_unacked(&message_id); + + let schedule = self.ensure_negative_ack_tracker().schedule( + message_id.clone(), + redelivery_count, + Instant::now(), + ); + + match schedule { + NegativeAckSchedule::Immediate => { + let mut requester = BrokerNegativeAckRedeliveryRequester { + connection: self.connection.clone(), + }; + let consumer_id = self.id; + let retry_at = Instant::now() + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL; + match dispatch_negative_ack_redelivery_with_requester( + self.ensure_negative_ack_tracker(), + &mut requester, + consumer_id, + std::slice::from_ref(&message_id), + retry_at, + ) + .await + { + Ok(()) => { + self.stop_negative_ack_ticker_if_idle(); + } + Err(e) => { + error!( + "could not ask for immediate negative ack redelivery: {:?}", + e + ); + if self.start_negative_ack_ticker().is_err() { + error!("could not start negative ack redelivery ticker after immediate redelivery failure"); + let mut requester = BrokerNegativeAckRedeliveryRequester { + connection: self.connection.clone(), + }; + let retry_at = Instant::now() + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL; + if let Err(e) = dispatch_negative_ack_redelivery_with_requester( + self.ensure_negative_ack_tracker(), + &mut requester, + consumer_id, + std::slice::from_ref(&message_id), + retry_at, + ) + .await + { + error!( + "could not ask for fallback immediate negative ack redelivery after ticker failure: {:?}", + e + ); + } + self.stop_negative_ack_ticker_if_idle(); + } + } + } + } + NegativeAckSchedule::Scheduled | NegativeAckSchedule::DuplicateEarlier => { + if self.start_negative_ack_ticker().is_err() { + error!("could not start negative ack redelivery ticker"); + let mut requester = BrokerNegativeAckRedeliveryRequester { + connection: self.connection.clone(), + }; + let consumer_id = self.id; + let retry_at = Instant::now() + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL; + if let Err(e) = dispatch_negative_ack_redelivery_with_requester( + self.ensure_negative_ack_tracker(), + &mut requester, + consumer_id, + std::slice::from_ref(&message_id), + retry_at, + ) + .await + { + error!( + "could not ask for immediate negative ack redelivery after ticker failure: {:?}", + e + ); + } + self.stop_negative_ack_ticker_if_idle(); + } + } + NegativeAckSchedule::RetryPending => {} + } + } + + fn remove_negative_ack_from_unacked(&mut self, message_id: &MessageIdData) { + self.unacked_messages + .retain(|id, _| !NegativeAckTracker::is_same_redelivery_entry(id, message_id)); + } + + async fn handle_negative_ack_redelivery(&mut self) { + self.negative_ack_due_event_pending + .store(false, Ordering::SeqCst); + + let ids = match self.negative_ack_tracker.as_mut() { + Some(tracker) => tracker.collect_due(Instant::now()), + None => return, + }; + + if ids.is_empty() { + self.stop_negative_ack_ticker_if_idle(); + return; + } + + let mut requester = BrokerNegativeAckRedeliveryRequester { + connection: self.connection.clone(), + }; + let consumer_id = self.id; + let retry_at = Instant::now() + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL; + if let Some(tracker) = self.negative_ack_tracker.as_mut() { + if let Err(e) = dispatch_negative_ack_redelivery_with_requester( + tracker, + &mut requester, + consumer_id, + &ids, + retry_at, + ) + .await + { + error!("could not ask for delayed negative ack redelivery: {:?}", e); + } } + + self.stop_negative_ack_ticker_if_idle(); } #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] async fn ack(&mut self, message_id: MessageIdData, cumulative: bool) { // FIXME: this does not handle cumulative acks self.unacked_messages.remove(&message_id); + if let Some(tracker) = self.negative_ack_tracker.as_mut() { + if cumulative { + tracker.cancel_cumulative_ack(&message_id); + } else { + tracker.cancel_ack(&message_id); + } + } + self.stop_negative_ack_ticker_if_idle(); let res = self .connection .sender() @@ -540,6 +867,7 @@ impl ConsumerEngine { } else { vec![(message.message_id, payload)] }; + let redelivery_count = message.redelivery_count; for (message_id, payload) in payloads { match (message.redelivery_count, &self.dead_letter_policy) { (Some(redelivery_count), Some(dead_letter_policy)) @@ -586,15 +914,14 @@ impl ConsumerEngine { self.client .send(&dead_letter_policy.dead_letter_topic, message) .await? - .await - .map_err(|e| { - error!("One shot cancelled {:?}", e); - Error::Custom("DLQ send error".to_string()) - })?; + .await?; self.ack(message_id, false).await; } - _ => self.send_to_consumer(message_id, payload).await?, + _ => { + self.send_to_consumer(message_id, payload, redelivery_count) + .await? + } } } Ok(()) @@ -605,10 +932,11 @@ impl ConsumerEngine { &mut self, message_id: MessageIdData, payload: Payload, + redelivery_count: Option, ) -> Result<(), Error> { let now = Instant::now(); self.tx - .send(Ok((message_id.clone(), payload))) + .send(Ok((message_id.clone(), payload, redelivery_count))) .await .map_err(|e| { error!("tx returned {:?}", e); @@ -692,6 +1020,8 @@ impl ConsumerEngine { impl std::ops::Drop for ConsumerEngine { fn drop(&mut self) { + self.clear_negative_ack_state(); + let conn = self.connection.clone(); let id = self.id; let name = self.name.clone(); @@ -706,3 +1036,249 @@ impl std::ops::Drop for ConsumerEngine { })); } } + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use futures::executor::block_on; + + use super::*; + use crate::{ + consumer::negative_ack_tracker::{NegativeAckSchedule, NegativeAckTracker}, + message::proto::MessageIdData, + }; + + #[test] + fn nack_handling_uses_tracker_with_zero_delay_immediate_redelivery() { + let now = Instant::now(); + let id = message_id(1, 2, Some(0)); + let mut tracker = NegativeAckTracker::new(Some(Duration::ZERO), None); + let mut requester = FakeNegativeAckRequester::default(); + + assert_eq!( + tracker.schedule(id.clone(), Some(1), now), + NegativeAckSchedule::Immediate + ); + block_on(dispatch_negative_ack_redelivery_with_requester( + &mut tracker, + &mut requester, + 7, + std::slice::from_ref(&id), + now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + )) + .unwrap(); + + assert_eq!(requester.redelivery_requests, vec![message_id(1, 2, None)]); + assert!(tracker + .collect_due(now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL) + .is_empty()); + } + + #[test] + fn due_redelivery_dispatch_is_in_flight_retry_safe() { + let now = Instant::now(); + let id = message_id(1, 2, Some(3)); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(5)), None); + let mut requester = FakeNegativeAckRequester { + fail_redelivery: true, + ..FakeNegativeAckRequester::default() + }; + + assert_eq!( + tracker.schedule(id.clone(), Some(1), now), + NegativeAckSchedule::Scheduled + ); + let due = tracker.collect_due(now + Duration::from_secs(5)); + assert_eq!(due, vec![message_id(1, 2, None)]); + + let result = block_on(dispatch_negative_ack_redelivery_with_requester( + &mut tracker, + &mut requester, + 7, + &due, + now + Duration::from_secs(5) + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + )); + + assert!(matches!(result, Err(ConnectionError::Disconnected))); + assert_eq!(requester.redelivery_requests, vec![message_id(1, 2, None)]); + let retried = tracker + .collect_due(now + Duration::from_secs(5) + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL); + assert_eq!(retried, vec![message_id(1, 2, None)]); + } + + #[test] + fn immediate_nack_failure_keeps_entry_retryable() { + let now = Instant::now(); + let id = message_id(5, 6, None); + let mut tracker = NegativeAckTracker::new(Some(Duration::ZERO), None); + let mut requester = FakeNegativeAckRequester { + fail_redelivery: true, + ..FakeNegativeAckRequester::default() + }; + + assert_eq!( + tracker.schedule(id.clone(), Some(99), now), + NegativeAckSchedule::Immediate + ); + let result = block_on(dispatch_negative_ack_redelivery_with_requester( + &mut tracker, + &mut requester, + 7, + std::slice::from_ref(&id), + now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + )); + + assert!(matches!(result, Err(ConnectionError::Disconnected))); + assert_eq!(requester.redelivery_requests, vec![id]); + let retried = tracker.collect_due(now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL); + assert_eq!(retried, vec![message_id(5, 6, None)]); + } + + #[derive(Default)] + struct FakeNegativeAckRequester { + redelivery_requests: Vec, + fail_redelivery: bool, + } + + #[async_trait::async_trait] + impl NegativeAckRedeliveryRequester for FakeNegativeAckRequester { + async fn request_negative_ack_redelivery( + &mut self, + _consumer_id: u64, + ids: Vec, + ) -> Result<(), ConnectionError> { + self.redelivery_requests.extend(ids); + if self.fail_redelivery { + Err(ConnectionError::Disconnected) + } else { + Ok(()) + } + } + } + + fn message_id(ledger_id: u64, entry_id: u64, batch_index: Option) -> MessageIdData { + message_id_for_partition(ledger_id, entry_id, Some(0), batch_index) + } + + fn message_id_for_partition( + ledger_id: u64, + entry_id: u64, + partition: Option, + batch_index: Option, + ) -> MessageIdData { + MessageIdData { + ledger_id, + entry_id, + partition, + batch_index, + ack_set: vec![], + batch_size: None, + first_chunk_message_id: None, + } + } + + #[test] + fn immediate_negative_ack_requests_broker_redelivery_only() { + let now = Instant::now(); + let id = message_id(1, 2, None); + let mut tracker = NegativeAckTracker::new(Some(Duration::ZERO), None); + let mut requester = FakeNegativeAckRequester::default(); + + assert_eq!( + tracker.schedule(id.clone(), Some(99), now), + NegativeAckSchedule::Immediate + ); + block_on(dispatch_negative_ack_redelivery_with_requester( + &mut tracker, + &mut requester, + 7, + std::slice::from_ref(&id), + now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + )) + .unwrap(); + + assert_eq!(requester.redelivery_requests, vec![id.clone()]); + assert!(tracker + .collect_due(now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL) + .is_empty()); + } + + #[test] + fn delayed_negative_ack_requests_broker_redelivery_when_due() { + let now = Instant::now(); + let id = message_id(1, 2, Some(3)); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(5)), None); + let mut requester = FakeNegativeAckRequester::default(); + + assert_eq!( + tracker.schedule(id.clone(), Some(99), now), + NegativeAckSchedule::Scheduled + ); + assert!(tracker.collect_due(now + Duration::from_secs(4)).is_empty()); + + let due = tracker.collect_due(now + Duration::from_secs(5)); + block_on(dispatch_negative_ack_redelivery_with_requester( + &mut tracker, + &mut requester, + 7, + &due, + now + Duration::from_secs(5) + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + )) + .unwrap(); + + assert_eq!(requester.redelivery_requests, vec![message_id(1, 2, None)]); + assert!(tracker + .collect_due(now + Duration::from_secs(5) + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL) + .is_empty()); + } + + #[test] + fn id_only_negative_ack_uses_fixed_delay_behavior() { + let now = Instant::now(); + let id = message_id(3, 4, None); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(7)), None); + let mut requester = FakeNegativeAckRequester::default(); + + assert_eq!( + tracker.schedule(id.clone(), None, now), + NegativeAckSchedule::Scheduled + ); + + let due = tracker.collect_due(now + Duration::from_secs(7)); + block_on(dispatch_negative_ack_redelivery_with_requester( + &mut tracker, + &mut requester, + 7, + &due, + now + Duration::from_secs(7) + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + )) + .unwrap(); + + assert_eq!(requester.redelivery_requests, vec![id.clone()]); + assert!(tracker + .collect_due(now + Duration::from_secs(7) + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL) + .is_empty()); + } + + #[test] + fn clear_negative_ack_state_drains_tracker_and_stops_ticker() { + let now = Instant::now(); + let id = message_id(1, 2, Some(0)); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(60)), None); + let ticker_running = Arc::new(AtomicBool::new(false)); + + // Schedule a nack so the tracker is non-empty + tracker.schedule(id.clone(), Some(0), now); + ticker_running.store(true, Ordering::SeqCst); + + // Simulate what clear_negative_ack_state does + tracker.clear(); + ticker_running.store(false, Ordering::SeqCst); + + assert!(tracker + .collect_due(now + Duration::from_secs(120)) + .is_empty()); + assert!(!ticker_running.load(Ordering::SeqCst)); + } +} diff --git a/src/consumer/message.rs b/src/consumer/message.rs index 73f25bc8..65d9b4d1 100644 --- a/src/consumer/message.rs +++ b/src/consumer/message.rs @@ -17,15 +17,26 @@ pub struct Message { pub payload: Payload, /// contains the message's id and batch size data pub message_id: MessageData, + redelivery_count: Option, pub(super) _phantom: PhantomData, } impl Message { pub fn new(topic: &str, message_id: MessageData, payload: Payload) -> Self { + Self::new_with_redelivery_count(topic, message_id, payload, None) + } + + pub(crate) fn new_with_redelivery_count( + topic: &str, + message_id: MessageData, + payload: Payload, + redelivery_count: Option, + ) -> Self { Message { topic: topic.to_string(), message_id, payload, + redelivery_count, _phantom: PhantomData, } } @@ -42,6 +53,14 @@ impl Message { &self.message_id.id } + /// Returns the broker-provided redelivery count for this message. Returns `0` when the broker + /// did not supply a redelivery count (first delivery or broker omitted the field), or when the + /// message was constructed directly with [`Message::new`]. + #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] + pub fn redelivery_count(&self) -> u32 { + self.redelivery_count.unwrap_or_default() + } + /// Get message key (partition key) #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub fn key(&self) -> Option { @@ -56,3 +75,66 @@ impl Message { T::deserialize_message(&self.payload) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::proto; + + fn message_id() -> proto::MessageIdData { + proto::MessageIdData { + ledger_id: 1, + entry_id: 1, + partition: None, + batch_index: None, + ack_set: vec![], + batch_size: None, + first_chunk_message_id: None, + } + } + + fn payload() -> Payload { + Payload { + metadata: proto::MessageMetadata::default(), + data: vec![], + } + } + + #[test] + fn test_redelivery_count_returns_broker_value() { + let message_id = MessageData { + id: message_id(), + batch_size: None, + }; + let msg = + Message::<()>::new_with_redelivery_count("test-topic", message_id, payload(), Some(7)); + + assert_eq!(msg.redelivery_count(), 7); + assert_eq!(msg.redelivery_count, Some(7)); + } + + #[test] + fn test_missing_redelivery_count_defaults_public_value_to_zero() { + let message_id = MessageData { + id: message_id(), + batch_size: None, + }; + let msg = Message::<()>::new("test-topic", message_id, payload()); + + assert_eq!(msg.redelivery_count(), 0); + assert_eq!(msg.redelivery_count, None); + } + + #[test] + fn test_explicit_zero_redelivery_count_remains_available_internally() { + let message_id = MessageData { + id: message_id(), + batch_size: None, + }; + let msg = + Message::<()>::new_with_redelivery_count("test-topic", message_id, payload(), Some(0)); + + assert_eq!(msg.redelivery_count(), 0); + assert_eq!(msg.redelivery_count, Some(0)); + } +} diff --git a/src/consumer/mod.rs b/src/consumer/mod.rs index 3111c8df..4f317e23 100644 --- a/src/consumer/mod.rs +++ b/src/consumer/mod.rs @@ -8,6 +8,8 @@ mod engine; pub(crate) mod initial_position; pub mod message; mod multi; +pub mod negative_ack_backoff; +mod negative_ack_tracker; pub mod options; pub mod topic; @@ -28,6 +30,9 @@ use futures::{ pub use initial_position::InitialPosition; pub use message::Message; use multi::MultiTopicConsumer; +pub use negative_ack_backoff::{ + MultiplierRedeliveryBackoff, MultiplierRedeliveryBackoffBuilder, NegativeAckBackoff, +}; pub use options::ConsumerOptions; pub use topic::TopicConsumer; use url::Url; @@ -152,9 +157,12 @@ impl Consumer { } } - /// negative acknowledgement + /// Negative acknowledgement. /// - /// the message will be sent again on the subscription + /// The message will be redelivered on the subscription after the configured negative-ack + /// redelivery delay (60 seconds by default). If a backoff policy is configured via + /// [`ConsumerBuilder::with_negative_ack_backoff`], the delay is computed from the message's + /// redelivery count instead. #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn nack(&mut self, msg: &Message) -> Result<(), ConsumerError> { match &mut self.inner { @@ -163,9 +171,11 @@ impl Consumer { } } - /// negative acknowledgement + /// Negative acknowledgement by message ID. /// - /// the message with the given ID will be sent again on the subscription + /// The message with the given ID will be redelivered on the subscription after the configured + /// negative-ack redelivery delay (60 seconds by default). Because no redelivery count is + /// available, a backoff policy is not applied. #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn nack_with_id( &mut self, @@ -886,6 +896,7 @@ mod tests { .with_subscription("nack") .with_subscription_type(SubType::Shared) .with_dead_letter_policy(dead_letter_policy) + .with_nack_redelivery_delay(Duration::from_millis(1)) .build() .await .unwrap(); @@ -1005,6 +1016,7 @@ mod tests { .with_subscription("nack") .with_subscription_type(SubType::Shared) .with_dead_letter_policy(dead_letter_policy) + .with_nack_redelivery_delay(Duration::from_millis(1)) .build() .await .unwrap(); diff --git a/src/consumer/negative_ack_backoff.rs b/src/consumer/negative_ack_backoff.rs new file mode 100644 index 00000000..b9c8ad82 --- /dev/null +++ b/src/consumer/negative_ack_backoff.rs @@ -0,0 +1,275 @@ +use std::time::{Duration, Instant}; + +use crate::Error; + +const MIN_DELAY_MUST_NOT_EXCEED_MAX_DELAY: &str = + "min_delay must be less than or equal to max_delay"; +const MULTIPLIER_MUST_BE_GREATER_THAN_ONE: &str = "multiplier must be greater than 1.0"; +const MULTIPLIER_MUST_BE_FINITE: &str = "multiplier must be finite"; +const DELAY_DURATION_TOO_LARGE: &str = "delay duration is too large"; + +/// Policy for computing a negative-acknowledgment redelivery delay from a redelivery count. +pub trait NegativeAckBackoff: Send + Sync { + /// Returns the delay to use for the provided redelivery count. + fn next(&self, redelivery_count: u32) -> Duration; +} + +/// Built-in negative-acknowledgment redelivery backoff using a multiplier. +#[derive(Clone, Debug)] +pub struct MultiplierRedeliveryBackoff { + min_delay: Duration, + max_delay: Duration, + multiplier: f64, +} + +impl Default for MultiplierRedeliveryBackoff { + fn default() -> Self { + Self { + min_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(600), + multiplier: 2.0, + } + } +} + +impl MultiplierRedeliveryBackoff { + /// Creates a builder for configuring a multiplier redelivery backoff policy. + pub fn builder() -> MultiplierRedeliveryBackoffBuilder { + MultiplierRedeliveryBackoffBuilder::default() + } +} + +impl NegativeAckBackoff for MultiplierRedeliveryBackoff { + fn next(&self, redelivery_count: u32) -> Duration { + if self.min_delay == Duration::ZERO { + return Duration::ZERO; + } + + if redelivery_count == 0 { + return self.min_delay; + } + + if self.min_delay >= self.max_delay { + return self.max_delay; + } + + let min_delay_secs = self.min_delay.as_secs_f64(); + let max_delay_secs = self.max_delay.as_secs_f64(); + let multiplier_log = self.multiplier.ln(); + let cap_after = (max_delay_secs / min_delay_secs).ln() / multiplier_log; + + if !cap_after.is_finite() || redelivery_count as f64 >= cap_after { + return self.max_delay; + } + + let delay_secs = min_delay_secs * self.multiplier.powf(redelivery_count as f64); + if !delay_secs.is_finite() || delay_secs >= max_delay_secs { + return self.max_delay; + } + + match Duration::try_from_secs_f64(delay_secs) { + Ok(delay) if delay < self.max_delay => delay, + _ => self.max_delay, + } + } +} + +/// Builder for [`MultiplierRedeliveryBackoff`]. +#[derive(Clone, Debug)] +pub struct MultiplierRedeliveryBackoffBuilder { + min_delay: Duration, + max_delay: Duration, + multiplier: f64, +} + +impl Default for MultiplierRedeliveryBackoffBuilder { + fn default() -> Self { + Self { + min_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(600), + multiplier: 2.0, + } + } +} + +impl MultiplierRedeliveryBackoffBuilder { + /// Sets the minimum delay returned for redelivery count `0`. + pub fn min_delay(mut self, min_delay: Duration) -> Self { + self.min_delay = min_delay; + self + } + + /// Sets the maximum delay cap for the backoff sequence. + pub fn max_delay(mut self, max_delay: Duration) -> Self { + self.max_delay = max_delay; + self + } + + /// Sets the multiplier applied for each redelivery count. + pub fn multiplier(mut self, multiplier: f64) -> Self { + self.multiplier = multiplier; + self + } + + /// Builds a validated [`MultiplierRedeliveryBackoff`]. + pub fn build(self) -> Result { + validate_delay_duration(self.min_delay)?; + validate_delay_duration(self.max_delay)?; + + if self.min_delay > self.max_delay { + return Err(Error::Custom( + MIN_DELAY_MUST_NOT_EXCEED_MAX_DELAY.to_string(), + )); + } + + if !self.multiplier.is_finite() { + return Err(Error::Custom(MULTIPLIER_MUST_BE_FINITE.to_string())); + } + + if self.multiplier <= 1.0 { + return Err(Error::Custom( + MULTIPLIER_MUST_BE_GREATER_THAN_ONE.to_string(), + )); + } + + Ok(MultiplierRedeliveryBackoff { + min_delay: self.min_delay, + max_delay: self.max_delay, + multiplier: self.multiplier, + }) + } +} + +fn validate_delay_duration(delay: Duration) -> Result<(), Error> { + if delay.as_millis() >= u64::MAX as u128 || Instant::now().checked_add(delay).is_none() { + return Err(Error::Custom(DELAY_DURATION_TOO_LARGE.to_string())); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn custom_backoff() -> MultiplierRedeliveryBackoff { + MultiplierRedeliveryBackoff::builder() + .min_delay(Duration::from_secs(2)) + .max_delay(Duration::from_secs(30)) + .multiplier(3.0) + .build() + .unwrap() + } + + fn assert_custom_error(result: Result, expected: &str) { + match result { + Err(Error::Custom(message)) => assert_eq!(message, expected), + other => panic!("expected Error::Custom({expected:?}), got {other:?}"), + } + } + + #[test] + fn default_multiplier_backoff_sequence_matches_pulsar_vectors() { + let backoff = MultiplierRedeliveryBackoff::default(); + + assert_eq!(backoff.next(0), Duration::from_secs(1)); + assert_eq!(backoff.next(1), Duration::from_secs(2)); + assert_eq!(backoff.next(2), Duration::from_secs(4)); + assert_eq!(backoff.next(9), Duration::from_secs(512)); + assert_eq!(backoff.next(10), Duration::from_secs(600)); + assert_eq!(backoff.next(u32::MAX), Duration::from_secs(600)); + } + + #[test] + fn custom_multiplier_backoff_sequence_caps_at_max_delay() { + let backoff = custom_backoff(); + + assert_eq!(backoff.next(0), Duration::from_secs(2)); + assert_eq!(backoff.next(1), Duration::from_secs(6)); + assert_eq!(backoff.next(2), Duration::from_secs(18)); + assert_eq!(backoff.next(3), Duration::from_secs(30)); + } + + #[test] + fn zero_min_delay_returns_zero_for_all_counts() { + let backoff = MultiplierRedeliveryBackoff::builder() + .min_delay(Duration::ZERO) + .build() + .unwrap(); + + assert_eq!(backoff.next(0), Duration::ZERO); + assert_eq!(backoff.next(1), Duration::ZERO); + assert_eq!(backoff.next(10), Duration::ZERO); + assert_eq!(backoff.next(u32::MAX), Duration::ZERO); + } + + #[test] + fn validation_rejects_min_delay_greater_than_max_delay() { + let result = MultiplierRedeliveryBackoff::builder() + .min_delay(Duration::from_secs(2)) + .max_delay(Duration::from_secs(1)) + .build(); + + assert_custom_error(result, MIN_DELAY_MUST_NOT_EXCEED_MAX_DELAY); + } + + #[test] + fn validation_rejects_multiplier_less_than_or_equal_to_one() { + let result = MultiplierRedeliveryBackoff::builder() + .multiplier(1.0) + .build(); + + assert_custom_error(result, MULTIPLIER_MUST_BE_GREATER_THAN_ONE); + } + + #[test] + fn validation_rejects_non_finite_multiplier() { + let nan_result = MultiplierRedeliveryBackoff::builder() + .multiplier(f64::NAN) + .build(); + assert_custom_error(nan_result, MULTIPLIER_MUST_BE_FINITE); + + let infinite_result = MultiplierRedeliveryBackoff::builder() + .multiplier(f64::INFINITY) + .build(); + assert_custom_error(infinite_result, MULTIPLIER_MUST_BE_FINITE); + } + + #[test] + fn validation_rejects_oversized_duration() { + let oversized = Duration::from_millis(u64::MAX) + .checked_add(Duration::from_millis(1)) + .unwrap(); + let result = MultiplierRedeliveryBackoff::builder() + .min_delay(oversized) + .max_delay(oversized) + .build(); + + assert_custom_error(result, DELAY_DURATION_TOO_LARGE); + } + + #[test] + fn validation_rejects_unschedulable_max_millis_duration() { + let result = MultiplierRedeliveryBackoff::builder() + .min_delay(Duration::from_millis(u64::MAX)) + .max_delay(Duration::from_millis(u64::MAX)) + .build(); + + assert_custom_error(result, DELAY_DURATION_TOO_LARGE); + } + + #[test] + fn custom_negative_ack_backoff_can_be_implemented_without_engine_dependencies() { + struct FixedBackoff(Duration); + + impl NegativeAckBackoff for FixedBackoff { + fn next(&self, _redelivery_count: u32) -> Duration { + self.0 + } + } + + let backoff = FixedBackoff(Duration::from_secs(42)); + + assert_eq!(backoff.next(123), Duration::from_secs(42)); + } +} diff --git a/src/consumer/negative_ack_tracker.rs b/src/consumer/negative_ack_tracker.rs new file mode 100644 index 00000000..9943234e --- /dev/null +++ b/src/consumer/negative_ack_tracker.rs @@ -0,0 +1,742 @@ +use std::{ + collections::{BTreeSet, HashMap}, + fmt, + sync::Arc, + time::{Duration, Instant}, +}; + +use crate::{consumer::negative_ack_backoff::NegativeAckBackoff, message::proto::MessageIdData}; + +pub(crate) const DEFAULT_NACK_REDELIVERY_DELAY: Duration = Duration::from_secs(60); +pub(crate) const NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL: Duration = Duration::from_millis(500); + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum NegativeAckSchedule { + Immediate, + Scheduled, + DuplicateEarlier, + RetryPending, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum NegativeAckDispatchState { + Pending, + InFlight, +} + +pub(crate) struct NegativeAckTracker { + fixed_delay: Duration, + backoff: Option>, + entries: HashMap, +} + +impl fmt::Debug for NegativeAckTracker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NegativeAckTracker") + .field("fixed_delay", &self.fixed_delay) + .field("entries", &self.entries) + .field("has_backoff", &self.backoff.is_some()) + .finish() + } +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct NormalizedMessageId { + ledger_id: u64, + entry_id: u64, + partition: Option, +} + +#[derive(Debug)] +struct PendingNegativeAck { + due_at: Instant, + state: NegativeAckDispatchState, + scope: PendingBatchScope, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum PendingBatchScope { + FullEntry, + BatchIndexes(BTreeSet), +} + +impl NormalizedMessageId { + fn from_message_id(message_id: &MessageIdData) -> Self { + Self { + ledger_id: message_id.ledger_id, + entry_id: message_id.entry_id, + partition: message_id.partition, + } + } + + fn into_redelivery_message_id(self) -> MessageIdData { + MessageIdData { + ledger_id: self.ledger_id, + entry_id: self.entry_id, + partition: self.partition, + batch_index: None, + ack_set: vec![], + batch_size: None, + first_chunk_message_id: None, + } + } +} + +impl PendingBatchScope { + fn from_message_id(message_id: &MessageIdData) -> Self { + match message_id.batch_index { + Some(batch_index) => Self::BatchIndexes(BTreeSet::from([batch_index])), + None => Self::FullEntry, + } + } + + fn merge(&mut self, incoming: PendingBatchScope) -> bool { + match (&mut *self, incoming) { + (Self::FullEntry, _) => false, + (scope @ Self::BatchIndexes(_), Self::FullEntry) => { + *scope = Self::FullEntry; + true + } + (Self::BatchIndexes(existing), Self::BatchIndexes(incoming)) => { + let mut changed = false; + for batch_index in incoming { + changed |= existing.insert(batch_index); + } + changed + } + } + } + + fn cancel_batch_index(&mut self, batch_index: i32) -> bool { + match self { + Self::FullEntry => false, + Self::BatchIndexes(indexes) => { + indexes.remove(&batch_index); + indexes.is_empty() + } + } + } + + fn cancel_batch_indexes_through(&mut self, batch_index: i32) -> bool { + match self { + Self::FullEntry => false, + Self::BatchIndexes(indexes) => { + indexes.retain(|pending| *pending > batch_index); + indexes.is_empty() + } + } + } +} + +impl NegativeAckTracker { + pub(crate) fn new( + fixed_delay: Option, + backoff: Option>, + ) -> Self { + Self { + fixed_delay: fixed_delay.unwrap_or(DEFAULT_NACK_REDELIVERY_DELAY), + backoff, + entries: HashMap::new(), + } + } + + pub(crate) fn select_delay(&self, redelivery_count: Option) -> Duration { + match (redelivery_count, &self.backoff) { + (Some(count), Some(backoff)) => backoff.next(count), + _ => self.fixed_delay, + } + } + + pub(crate) fn redelivery_message_id(message_id: &MessageIdData) -> MessageIdData { + NormalizedMessageId::from_message_id(message_id).into_redelivery_message_id() + } + + pub(crate) fn is_same_redelivery_entry(left: &MessageIdData, right: &MessageIdData) -> bool { + NormalizedMessageId::from_message_id(left) == NormalizedMessageId::from_message_id(right) + } + + pub(crate) fn schedule( + &mut self, + message_id: MessageIdData, + redelivery_count: Option, + now: Instant, + ) -> NegativeAckSchedule { + let delay = self.select_delay(redelivery_count); + let key = NormalizedMessageId::from_message_id(&message_id); + let incoming_scope = PendingBatchScope::from_message_id(&message_id); + + if delay == Duration::ZERO { + match self.entries.get_mut(&key) { + None => { + self.entries.insert( + key, + PendingNegativeAck { + due_at: now, + state: NegativeAckDispatchState::InFlight, + scope: incoming_scope, + }, + ); + return NegativeAckSchedule::Immediate; + } + Some(existing) if existing.state == NegativeAckDispatchState::InFlight => { + return NegativeAckSchedule::RetryPending; + } + Some(existing) => { + existing.scope.merge(incoming_scope); + existing.due_at = now; + existing.state = NegativeAckDispatchState::InFlight; + return NegativeAckSchedule::Immediate; + } + } + } + + let due_at = self.due_at(now, delay); + + match self.entries.get_mut(&key) { + None => { + self.entries.insert( + key, + PendingNegativeAck { + due_at, + state: NegativeAckDispatchState::Pending, + scope: incoming_scope, + }, + ); + NegativeAckSchedule::Scheduled + } + Some(existing) if existing.state == NegativeAckDispatchState::InFlight => { + NegativeAckSchedule::RetryPending + } + Some(existing) => { + let scope_changed = existing.scope.merge(incoming_scope); + if due_at < existing.due_at { + existing.due_at = due_at; + NegativeAckSchedule::Scheduled + } else if scope_changed { + NegativeAckSchedule::Scheduled + } else { + NegativeAckSchedule::DuplicateEarlier + } + } + } + } + + fn due_at(&self, now: Instant, delay: Duration) -> Instant { + now.checked_add(delay) + .or_else(|| now.checked_add(self.fixed_delay)) + .or_else(|| now.checked_add(DEFAULT_NACK_REDELIVERY_DELAY)) + .unwrap_or(now) + } + + #[cfg(test)] + pub(crate) fn pending_len(&self) -> usize { + self.entries + .values() + .filter(|entry| entry.state == NegativeAckDispatchState::Pending) + .count() + } + + #[cfg(test)] + pub(crate) fn in_flight_len(&self) -> usize { + self.entries + .values() + .filter(|entry| entry.state == NegativeAckDispatchState::InFlight) + .count() + } + + pub(crate) fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + #[cfg(test)] + pub(crate) fn next_due_time(&self) -> Option { + self.entries + .values() + .filter(|entry| entry.state == NegativeAckDispatchState::Pending) + .map(|entry| entry.due_at) + .min() + } + + pub(crate) fn collect_due(&mut self, now: Instant) -> Vec { + let mut due = Vec::new(); + for (key, entry) in &mut self.entries { + if entry.state == NegativeAckDispatchState::Pending && entry.due_at <= now { + entry.state = NegativeAckDispatchState::InFlight; + due.push((*key).into_redelivery_message_id()); + } + } + due + } + + pub(crate) fn mark_dispatched(&mut self, message_ids: &[MessageIdData]) { + for message_id in message_ids { + self.entries + .remove(&NormalizedMessageId::from_message_id(message_id)); + } + } + + pub(crate) fn mark_retry_pending(&mut self, message_ids: &[MessageIdData], retry_at: Instant) { + for message_id in message_ids { + if let Some(entry) = self + .entries + .get_mut(&NormalizedMessageId::from_message_id(message_id)) + { + entry.state = NegativeAckDispatchState::Pending; + entry.due_at = retry_at; + } + } + } + + pub(crate) fn cancel_ack(&mut self, message_id: &MessageIdData) { + let key = NormalizedMessageId::from_message_id(message_id); + let Some(batch_index) = message_id.batch_index else { + self.entries.remove(&key); + return; + }; + if self + .entries + .get_mut(&key) + .is_some_and(|entry| entry.scope.cancel_batch_index(batch_index)) + { + self.entries.remove(&key); + } + } + + pub(crate) fn cancel_cumulative_ack(&mut self, message_id: &MessageIdData) { + let ack_key = NormalizedMessageId::from_message_id(message_id); + let mut remove_keys = Vec::new(); + + for key in self.entries.keys() { + if key.partition != ack_key.partition { + continue; + } + if key.ledger_id < ack_key.ledger_id + || (key.ledger_id == ack_key.ledger_id && key.entry_id < ack_key.entry_id) + { + remove_keys.push(*key); + } + } + + for key in remove_keys { + self.entries.remove(&key); + } + + let Some(batch_index) = message_id.batch_index else { + self.entries.remove(&ack_key); + return; + }; + if self + .entries + .get_mut(&ack_key) + .is_some_and(|entry| entry.scope.cancel_batch_indexes_through(batch_index)) + { + self.entries.remove(&ack_key); + } + } + + pub(crate) fn clear(&mut self) { + self.entries.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct FixedBackoff(Duration); + + impl NegativeAckBackoff for FixedBackoff { + fn next(&self, _redelivery_count: u32) -> Duration { + self.0 + } + } + + fn backoff(delay: Duration) -> Arc { + Arc::new(FixedBackoff(delay)) + } + + fn message_id( + ledger_id: u64, + entry_id: u64, + partition: Option, + batch_index: Option, + ) -> MessageIdData { + MessageIdData { + ledger_id, + entry_id, + partition, + batch_index, + ack_set: vec![], + batch_size: None, + first_chunk_message_id: None, + } + } + + #[test] + fn missing_fixed_delay_defaults_to_sixty_seconds() { + let tracker = NegativeAckTracker::new(None, None); + + assert_eq!(tracker.select_delay(None), DEFAULT_NACK_REDELIVERY_DELAY); + assert_eq!(tracker.select_delay(None), Duration::from_secs(60)); + } + + #[test] + fn explicit_zero_fixed_delay_selects_immediate_redelivery() { + let tracker = NegativeAckTracker::new(Some(Duration::ZERO), None); + + assert_eq!(tracker.select_delay(None), Duration::ZERO); + } + + #[test] + fn message_redelivery_count_uses_configured_backoff() { + let tracker = NegativeAckTracker::new(None, Some(backoff(Duration::from_secs(42)))); + + assert_eq!(tracker.select_delay(Some(7)), Duration::from_secs(42)); + } + + #[test] + fn missing_redelivery_count_uses_fixed_delay_even_when_backoff_exists() { + let tracker = NegativeAckTracker::new( + Some(Duration::from_secs(5)), + Some(backoff(Duration::from_secs(42))), + ); + + assert_eq!(tracker.select_delay(None), Duration::from_secs(5)); + } + + #[test] + fn zero_backoff_delay_selects_immediate_redelivery() { + let tracker = + NegativeAckTracker::new(Some(Duration::from_secs(5)), Some(backoff(Duration::ZERO))); + + assert_eq!(tracker.select_delay(Some(7)), Duration::ZERO); + } + + #[test] + fn immediate_schedule_is_in_flight_until_dispatch_result() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::ZERO), None); + + assert_eq!( + tracker.schedule(message_id(1, 2, None, None), None, now), + NegativeAckSchedule::Immediate + ); + assert_eq!(tracker.pending_len(), 0); + assert_eq!(tracker.in_flight_len(), 1); + + tracker.mark_retry_pending( + &[message_id(1, 2, None, None)], + now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL, + ); + assert_eq!(tracker.pending_len(), 1); + assert_eq!(tracker.in_flight_len(), 0); + assert!(tracker + .collect_due(now + Duration::from_millis(499)) + .is_empty()); + + let due = tracker.collect_due(now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL); + tracker.mark_dispatched(&due); + assert!(tracker.is_empty()); + } + + #[test] + fn oversized_backoff_delay_falls_back_without_panicking() { + let now = Instant::now(); + let mut tracker = + NegativeAckTracker::new(Some(Duration::from_secs(5)), Some(backoff(Duration::MAX))); + + assert_eq!( + tracker.schedule(message_id(1, 2, None, None), Some(7), now), + NegativeAckSchedule::Scheduled + ); + assert_eq!(tracker.next_due_time(), Some(now + Duration::from_secs(5))); + } + + #[test] + fn oversized_backoff_and_fixed_delays_fall_back_to_default_delay() { + let now = Instant::now(); + let mut tracker = + NegativeAckTracker::new(Some(Duration::MAX), Some(backoff(Duration::MAX))); + + assert_eq!( + tracker.schedule(message_id(1, 2, None, None), Some(7), now), + NegativeAckSchedule::Scheduled + ); + assert_eq!( + tracker.next_due_time(), + Some(now + DEFAULT_NACK_REDELIVERY_DELAY) + ); + } + + #[test] + fn redelivery_message_id_normalizes_batch_specific_fields() { + let mut original = message_id(1, 2, Some(0), Some(1)); + original.ack_set = vec![1]; + original.batch_size = Some(3); + original.first_chunk_message_id = Some(Box::new(message_id(1, 1, Some(0), None))); + + let redelivery_id = NegativeAckTracker::redelivery_message_id(&original); + + assert_eq!(redelivery_id.ledger_id, 1); + assert_eq!(redelivery_id.entry_id, 2); + assert_eq!(redelivery_id.partition, Some(0)); + assert_eq!(redelivery_id.batch_index, None); + assert!(redelivery_id.ack_set.is_empty()); + assert_eq!(redelivery_id.batch_size, None); + assert_eq!(redelivery_id.first_chunk_message_id, None); + } + + #[test] + fn same_redelivery_entry_ignores_batch_specific_fields() { + let mut original = message_id(1, 2, Some(0), Some(1)); + original.ack_set = vec![1]; + original.batch_size = Some(3); + + assert!(NegativeAckTracker::is_same_redelivery_entry( + &original, + &message_id(1, 2, Some(0), Some(2)) + )); + assert!(!NegativeAckTracker::is_same_redelivery_entry( + &original, + &message_id(1, 2, Some(1), Some(1)) + )); + assert!(!NegativeAckTracker::is_same_redelivery_entry( + &original, + &message_id(1, 3, Some(0), Some(1)) + )); + } + + #[test] + fn batch_ids_share_entry_key_but_track_pending_indexes_separately() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + assert_eq!( + tracker.schedule(message_id(1, 2, Some(0), Some(1)), None, now), + NegativeAckSchedule::Scheduled + ); + assert_eq!( + tracker.schedule(message_id(1, 2, Some(0), Some(2)), None, now), + NegativeAckSchedule::Scheduled + ); + + assert_eq!(tracker.pending_len(), 1); + tracker.cancel_ack(&message_id(1, 2, Some(0), Some(1))); + assert_eq!(tracker.pending_len(), 1); + tracker.cancel_ack(&message_id(1, 2, Some(0), Some(2))); + assert!(tracker.is_empty()); + } + + #[test] + fn duplicate_schedule_keeps_earliest_due_time() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + assert_eq!( + tracker.schedule(message_id(1, 2, None, Some(1)), None, now), + NegativeAckSchedule::Scheduled + ); + assert_eq!( + tracker.schedule( + message_id(1, 2, None, Some(1)), + None, + now + Duration::from_secs(5) + ), + NegativeAckSchedule::DuplicateEarlier + ); + + assert_eq!(tracker.pending_len(), 1); + assert_eq!(tracker.next_due_time(), Some(now + Duration::from_secs(10))); + } + + #[test] + fn whole_entry_ack_cancels_all_pending_batch_indexes() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 2, None, Some(1)), None, now); + tracker.schedule(message_id(1, 2, None, Some(2)), None, now); + tracker.cancel_ack(&message_id(1, 2, None, None)); + + assert!(tracker.is_empty()); + } + + #[test] + fn ack_set_ack_with_batch_index_preserves_sibling_pending_indexes() { + let now = Instant::now(); + let mut ack = message_id(1, 2, None, Some(1)); + ack.ack_set = vec![0]; + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 2, None, Some(1)), None, now); + tracker.schedule(message_id(1, 2, None, Some(2)), None, now); + tracker.cancel_ack(&ack); + + assert_eq!(tracker.pending_len(), 1); + tracker.cancel_ack(&message_id(1, 2, None, Some(2))); + assert!(tracker.is_empty()); + } + + #[test] + fn ack_set_without_batch_index_cancels_the_whole_pending_entry() { + let now = Instant::now(); + let mut ack = message_id(1, 2, None, None); + ack.ack_set = vec![0]; + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 2, None, Some(1)), None, now); + tracker.schedule(message_id(1, 2, None, Some(2)), None, now); + tracker.cancel_ack(&ack); + + assert!(tracker.is_empty()); + } + + #[test] + fn cumulative_ack_is_partition_scoped_and_handles_same_entry_batches() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 1, Some(0), Some(1)), None, now); + tracker.schedule(message_id(1, 2, Some(0), Some(1)), None, now); + tracker.schedule(message_id(1, 2, Some(0), Some(2)), None, now); + tracker.schedule(message_id(1, 3, Some(0), Some(1)), None, now); + tracker.schedule(message_id(1, 1, Some(1), Some(1)), None, now); + + tracker.cancel_cumulative_ack(&message_id(1, 2, Some(0), Some(1))); + + assert_eq!(tracker.pending_len(), 3); + tracker.cancel_ack(&message_id(1, 2, Some(0), Some(2))); + assert_eq!(tracker.pending_len(), 2); + tracker.cancel_ack(&message_id(1, 3, Some(0), Some(1))); + tracker.cancel_ack(&message_id(1, 1, Some(1), Some(1))); + assert!(tracker.is_empty()); + } + + #[test] + fn cumulative_ack_set_with_batch_index_preserves_later_sibling_indexes() { + let now = Instant::now(); + let mut ack = message_id(1, 2, Some(0), Some(1)); + ack.ack_set = vec![0]; + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 1, Some(0), Some(1)), None, now); + tracker.schedule(message_id(1, 2, Some(0), Some(1)), None, now); + tracker.schedule(message_id(1, 2, Some(0), Some(2)), None, now); + tracker.schedule(message_id(1, 3, Some(0), Some(1)), None, now); + + tracker.cancel_cumulative_ack(&ack); + + assert_eq!(tracker.pending_len(), 2); + tracker.cancel_ack(&message_id(1, 2, Some(0), Some(2))); + tracker.cancel_ack(&message_id(1, 3, Some(0), Some(1))); + assert!(tracker.is_empty()); + } + + #[test] + fn collect_due_marks_entries_in_flight_and_returns_entry_level_ids() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + let mut original = message_id(1, 2, Some(0), Some(1)); + original.ack_set = vec![1]; + original.batch_size = Some(3); + original.first_chunk_message_id = Some(Box::new(message_id(1, 1, Some(0), None))); + + tracker.schedule(original, None, now); + + assert!(tracker.collect_due(now + Duration::from_secs(9)).is_empty()); + let due = tracker.collect_due(now + Duration::from_secs(10)); + assert_eq!(due.len(), 1); + assert_eq!(due[0].ledger_id, 1); + assert_eq!(due[0].entry_id, 2); + assert_eq!(due[0].partition, Some(0)); + assert_eq!(due[0].batch_index, None); + assert!(due[0].ack_set.is_empty()); + assert_eq!(due[0].batch_size, None); + assert_eq!(due[0].first_chunk_message_id, None); + assert_eq!(tracker.pending_len(), 0); + assert_eq!(tracker.in_flight_len(), 1); + assert!(tracker + .collect_due(now + Duration::from_secs(11)) + .is_empty()); + } + + #[test] + fn dispatch_lifecycle_removes_or_retries_in_flight_entries() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 2, None, None), None, now); + let due = tracker.collect_due(now + Duration::from_secs(10)); + tracker.mark_retry_pending(&due, now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL); + assert_eq!(tracker.pending_len(), 1); + assert_eq!(tracker.in_flight_len(), 0); + assert!(tracker + .collect_due(now + Duration::from_millis(499)) + .is_empty()); + let retried = tracker.collect_due(now + NEGATIVE_ACK_REDELIVERY_TICK_INTERVAL); + tracker.mark_dispatched(&retried); + assert!(tracker.is_empty()); + } + + #[test] + fn clear_removes_pending_and_in_flight_entries() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_secs(10)), None); + + tracker.schedule(message_id(1, 2, None, None), None, now); + tracker.schedule(message_id(1, 3, None, None), None, now); + tracker.collect_due(now + Duration::from_secs(10)); + tracker.clear(); + + assert!(tracker.is_empty()); + } + + #[test] + fn scheduled_nack_is_not_due_before_delay_elapses() { + let delay = Duration::from_secs(10); + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(delay), None); + + tracker.schedule(message_id(1, 2, None, None), None, now); + + assert!(tracker + .collect_due(now + delay - Duration::from_millis(1)) + .is_empty()); + } + + #[test] + fn scheduled_nack_is_due_at_or_after_delay_elapses() { + let delay = Duration::from_secs(10); + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(delay), None); + + tracker.schedule(message_id(1, 3, None, None), None, now); + + let due = tracker.collect_due(now + delay); + assert_eq!(due.len(), 1); + } + + #[test] + fn nack_with_id_fixed_delay_fallback_when_backoff_configured() { + let tracker = NegativeAckTracker::new( + Some(Duration::from_secs(5)), + Some(backoff(Duration::from_secs(42))), + ); + + assert_eq!(tracker.select_delay(None), Duration::from_secs(5)); + assert_eq!(tracker.select_delay(Some(1)), Duration::from_secs(42)); + } + + #[test] + fn dlq_regression_batch_normalization_no_duplicate_scheduling() { + let now = Instant::now(); + let mut tracker = NegativeAckTracker::new(Some(Duration::from_millis(1)), None); + + tracker.schedule(message_id(1, 2, None, Some(0)), None, now); + tracker.schedule(message_id(1, 2, None, Some(1)), None, now); + + assert_eq!(tracker.pending_len(), 1); + let due = tracker.collect_due(now + Duration::from_millis(1)); + assert_eq!(due.len(), 1); + } +} diff --git a/src/consumer/topic.rs b/src/consumer/topic.rs index e9e9d917..ebaadce2 100644 --- a/src/consumer/topic.rs +++ b/src/consumer/topic.rs @@ -19,7 +19,9 @@ use crate::{ connection::Connection, consumer::{ config::ConsumerConfig, - data::{DeadLetterPolicy, EngineMessage, MessageData, MessageIdDataReceiver}, + data::{ + DeadLetterPolicy, InternalEngineMessage, InternalMessageIdDataReceiver, MessageData, + }, engine::ConsumerEngine, message::Message, }, @@ -35,8 +37,8 @@ pub struct TopicConsumer { pub(crate) consumer_id: u64, pub(crate) config: ConsumerConfig, topic: String, - messages: Pin>, - engine_tx: mpsc::UnboundedSender>, + messages: Pin>, + engine_tx: mpsc::UnboundedSender>, data_type: PhantomData T::Output>, pub(crate) dead_letter_policy: Option, pub(super) last_message_received: Option>, @@ -62,6 +64,8 @@ impl TopicConsumer { unacked_message_redelivery_delay, options, dead_letter_policy, + nack_redelivery_delay, + negative_ack_backoff, } = config.clone(); let consumer_id = consumer_id.unwrap_or_else(|| CONSUMER_ID_GENERATOR.fetch_add(1, Ordering::SeqCst)); @@ -90,7 +94,7 @@ impl TopicConsumer { let res = client.executor.spawn(Box::pin(async move { while interval.next().await.is_some() { if redelivery_tx - .send(EngineMessage::UnackedRedelivery) + .send(InternalEngineMessage::UnackedRedelivery) .await .is_err() { @@ -118,6 +122,8 @@ impl TopicConsumer { engine_rx, batch_size, unacked_message_redelivery_delay, + nack_redelivery_delay, + negative_ack_backoff, dead_letter_policy.clone(), options.clone(), ); @@ -154,7 +160,7 @@ impl TopicConsumer { pub async fn connection(&mut self) -> Result>, Error> { let (resolver, response) = oneshot::channel(); self.engine_tx - .send(EngineMessage::GetConnection(resolver)) + .send(InternalEngineMessage::GetConnection(resolver)) .await .map_err(|_| ConsumerError::Connection(ConnectionError::Disconnected))?; @@ -183,7 +189,7 @@ impl TopicConsumer { #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn ack(&mut self, msg: &Message) -> Result<(), ConsumerError> { self.engine_tx - .send(EngineMessage::Ack(msg.message_id().clone(), false)) + .send(InternalEngineMessage::Ack(msg.message_id().clone(), false)) .await?; Ok(()) } @@ -191,20 +197,20 @@ impl TopicConsumer { #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn ack_with_id(&mut self, msg_id: MessageIdData) -> Result<(), ConsumerError> { self.engine_tx - .send(EngineMessage::Ack(msg_id, false)) + .send(InternalEngineMessage::Ack(msg_id, false)) .await?; Ok(()) } #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] - pub(crate) fn acker(&self) -> mpsc::UnboundedSender> { + pub(crate) fn acker(&self) -> mpsc::UnboundedSender> { self.engine_tx.clone() } #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn cumulative_ack(&mut self, msg: &Message) -> Result<(), ConsumerError> { self.engine_tx - .send(EngineMessage::Ack(msg.message_id().clone(), true)) + .send(InternalEngineMessage::Ack(msg.message_id().clone(), true)) .await?; Ok(()) } @@ -215,7 +221,7 @@ impl TopicConsumer { msg_id: MessageIdData, ) -> Result<(), ConsumerError> { self.engine_tx - .send(EngineMessage::Ack(msg_id, true)) + .send(InternalEngineMessage::Ack(msg_id, true)) .await?; Ok(()) } @@ -223,14 +229,19 @@ impl TopicConsumer { #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn nack(&mut self, msg: &Message) -> Result<(), ConsumerError> { self.engine_tx - .send(EngineMessage::Nack(msg.message_id().clone())) + .send(InternalEngineMessage::Nack( + msg.message_id().clone(), + Some(msg.redelivery_count()), + )) .await?; Ok(()) } #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn nack_with_id(&mut self, msg_id: MessageIdData) -> Result<(), ConsumerError> { - self.engine_tx.send(EngineMessage::Nack(msg_id)).await?; + self.engine_tx + .send(InternalEngineMessage::Nack(msg_id, None)) + .await?; Ok(()) } @@ -263,6 +274,9 @@ impl TopicConsumer { #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] pub async fn close(&mut self) -> Result<(), Error> { let consumer_id = self.consumer_id; + let _ = self + .engine_tx + .unbounded_send(InternalEngineMessage::ClearNegativeAcks); self.connection() .await? .sender() @@ -295,12 +309,17 @@ impl TopicConsumer { } #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] - fn create_message(&self, message_id: MessageIdData, payload: Payload) -> Message { + fn create_message( + &self, + message_id: MessageIdData, + payload: Payload, + redelivery_count: Option, + ) -> Message { let message_id = MessageData { id: message_id, batch_size: payload.metadata.num_messages_in_batch, }; - Message::new(&self.topic, message_id, payload) + Message::new_with_redelivery_count(&self.topic, message_id, payload, redelivery_count) } #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))] @@ -322,10 +341,10 @@ impl Stream for TopicConsumer { match self.messages.as_mut().poll_next(cx) { Poll::Pending => Poll::Pending, Poll::Ready(None) => Poll::Ready(None), - Poll::Ready(Some(Ok((id, payload)))) => { + Poll::Ready(Some(Ok((id, payload, redelivery_count)))) => { self.last_message_received = Some(Utc::now()); self.messages_received += 1; - Poll::Ready(Some(Ok(self.create_message(id, payload)))) + Poll::Ready(Some(Ok(self.create_message(id, payload, redelivery_count)))) } Poll::Ready(Some(Err(e))) => { error!("we are using in the single-consumer and we got an error, {e}"); @@ -334,3 +353,92 @@ impl Stream for TopicConsumer { } } } + +#[cfg(test)] +mod tests { + use futures::{channel::mpsc, executor::block_on, StreamExt}; + + use super::*; + use crate::{ + executor::TokioExecutor, + message::{proto::MessageIdData, Metadata}, + }; + + fn message_id() -> MessageIdData { + MessageIdData { + ledger_id: 1, + entry_id: 2, + partition: None, + batch_index: None, + ack_set: vec![], + batch_size: None, + first_chunk_message_id: None, + } + } + + fn message(redelivery_count: Option) -> Message> { + Message::new_with_redelivery_count( + "persistent://public/default/topic", + MessageData { + id: message_id(), + batch_size: None, + }, + Payload { + metadata: Metadata::default(), + data: vec![], + }, + redelivery_count, + ) + } + + fn topic_consumer() -> ( + TopicConsumer, TokioExecutor>, + mpsc::UnboundedReceiver>, + ) { + let (_messages_tx, messages_rx) = mpsc::channel(1); + let (engine_tx, engine_rx) = mpsc::unbounded(); + + ( + TopicConsumer { + consumer_id: 1, + config: ConsumerConfig::default(), + topic: "persistent://public/default/topic".to_string(), + messages: Box::pin(messages_rx), + engine_tx, + data_type: PhantomData, + dead_letter_policy: None, + last_message_received: None, + messages_received: 0, + }, + engine_rx, + ) + } + + #[test] + fn nack_first_delivery_uses_zero_redelivery_count_for_backoff() { + let (mut consumer, mut engine_rx) = topic_consumer(); + + block_on(consumer.nack(&message(None))).unwrap(); + + match block_on(engine_rx.next()).expect("nack event") { + InternalEngineMessage::Nack(_, redelivery_count) => { + assert_eq!(redelivery_count, Some(0)); + } + _ => panic!("expected nack engine message"), + } + } + + #[test] + fn nack_redelivery_uses_broker_redelivery_count_for_backoff() { + let (mut consumer, mut engine_rx) = topic_consumer(); + + block_on(consumer.nack(&message(Some(3)))).unwrap(); + + match block_on(engine_rx.next()).expect("nack event") { + InternalEngineMessage::Nack(_, redelivery_count) => { + assert_eq!(redelivery_count, Some(3)); + } + _ => panic!("expected nack engine message"), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 0bbca9f7..1747f510 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -192,7 +192,10 @@ pub use connection::Authentication; pub use connection_manager::{ BrokerAddress, ConnectionRetryOptions, OperationRetryOptions, TlsOptions, }; -pub use consumer::{Consumer, ConsumerBuilder, ConsumerOptions}; +pub use consumer::{ + Consumer, ConsumerBuilder, ConsumerOptions, MultiplierRedeliveryBackoff, + MultiplierRedeliveryBackoffBuilder, NegativeAckBackoff, +}; pub use error::Error; #[cfg(any( feature = "async-std-runtime", diff --git a/src/reader.rs b/src/reader.rs index dba1e4c2..fb763252 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -10,7 +10,9 @@ use url::Url; use crate::{ client::DeserializeMessage, - consumer::{ConsumerOptions, DeadLetterPolicy, EngineMessage, Message, TopicConsumer}, + consumer::{ + data::InternalEngineMessage, ConsumerOptions, DeadLetterPolicy, Message, TopicConsumer, + }, error::Error, executor::Executor, message::proto::{command_subscribe::SubType, MessageIdData}, @@ -55,9 +57,11 @@ impl Stream for Reader { let message_id = msg.message_id().clone(); this.state = Some(State::PollingAck( msg, - Box::pin( - async move { acker.send(EngineMessage::Ack(message_id, false)).await }, - ), + Box::pin(async move { + acker + .send(InternalEngineMessage::Ack(message_id, false)) + .await + }), )); Pin::new(this).poll_next(cx) }