From 37b1cd4c2d9e8cfa8dd13ad863cc404b0d10f6a7 Mon Sep 17 00:00:00 2001 From: colin Date: Fri, 19 Jun 2026 11:23:24 -0700 Subject: [PATCH 1/5] enhancement(mqtt source): support end-to-end acknowledgements The mqtt source now participates in Vector's end-to-end acknowledgement framework. When acknowledgements are enabled, the QoS 1 PUBACK is deferred via an OrderedFinalizer until the events are delivered to all sinks; with the source's existing clean_session=false + QoS AtLeastOnce subscription, unacknowledged messages are redelivered after a reconnect (at-least-once). Closes #21967 --- ...qtt_source_acknowledgements.enhancement.md | 1 + src/common/mqtt.rs | 5 + src/sources/mqtt/config.rs | 67 ++- src/sources/mqtt/integration_tests.rs | 109 ++++- src/sources/mqtt/source.rs | 396 ++++++++++++++++-- .../components/sources/generated/mqtt.cue | 21 + 6 files changed, 558 insertions(+), 41 deletions(-) create mode 100644 changelog.d/21967_mqtt_source_acknowledgements.enhancement.md diff --git a/changelog.d/21967_mqtt_source_acknowledgements.enhancement.md b/changelog.d/21967_mqtt_source_acknowledgements.enhancement.md new file mode 100644 index 0000000000000..3eb341aede3b8 --- /dev/null +++ b/changelog.d/21967_mqtt_source_acknowledgements.enhancement.md @@ -0,0 +1 @@ +The `mqtt` source now supports end-to-end acknowledgements. When `acknowledgements` is enabled, the QoS 1 `PUBACK` for an incoming publish is deferred until the resulting events have been delivered to all connected sinks. Because the source already uses `clean_session = false` and subscribes at QoS `AtLeastOnce`, an unacknowledged message is redelivered by the broker after a crash or reconnect, providing at-least-once delivery when the broker resumes the persistent session. A stable `client_id` must be configured when `acknowledgements` is enabled, so the session (and its unacknowledged messages) can be resumed after a restart. Publishes delivered by the broker with QoS 0 cannot be acknowledged end-to-end and are forwarded with a warning instead of being registered for deferred acknowledgement. The source also warns if the broker starts a new session while acknowledgements are enabled, since unacknowledged messages from any previous session for that client ID will not be redelivered, and re-subscribes so the fresh session continues receiving future messages. diff --git a/src/common/mqtt.rs b/src/common/mqtt.rs index 8d4d377ab6746..d9e982815b36b 100644 --- a/src/common/mqtt.rs +++ b/src/common/mqtt.rs @@ -104,6 +104,11 @@ pub enum ConfigurationError { /// Credentials provided were incomplete #[snafu(display("Username and password must be either both or neither provided."))] IncompleteCredentials, + /// Acknowledgements enabled without a stable client ID + #[snafu(display( + "A stable `client_id` must be set when `acknowledgements` are enabled, so the MQTT session and its unacknowledged messages can be resumed after a restart." + ))] + AcknowledgementsRequireClientId, } #[derive(Clone)] diff --git a/src/sources/mqtt/config.rs b/src/sources/mqtt/config.rs index 92b8cd78f6a5b..86b845d7b7770 100644 --- a/src/sources/mqtt/config.rs +++ b/src/sources/mqtt/config.rs @@ -19,8 +19,8 @@ use crate::{ ConfigurationError, ConfigurationSnafu, MqttCommonConfig, MqttConnector, MqttError, TlsSnafu, }, - config::{SourceConfig, SourceContext, SourceOutput}, - serde::{OneOrMany, default_decoding, default_framing_message_based}, + config::{SourceAcknowledgementsConfig, SourceConfig, SourceContext, SourceOutput}, + serde::{OneOrMany, bool_or_struct, default_decoding, default_framing_message_based}, }; /// Configuration for the `mqtt` source. @@ -61,6 +61,10 @@ pub struct MqttSourceConfig { #[serde(default = "default_topic_key")] #[configurable(metadata(docs::examples = "topic"))] pub topic_key: OptionalValuePath, + + #[configurable(derived)] + #[serde(default, deserialize_with = "bool_or_struct")] + pub acknowledgements: SourceAcknowledgementsConfig, } fn default_topic() -> OneOrMany { @@ -77,14 +81,22 @@ impl SourceConfig for MqttSourceConfig { async fn build(&self, cx: SourceContext) -> crate::Result { let log_namespace = cx.log_namespace(self.log_namespace); - let connector = self.build_connector()?; + let acknowledgements = cx.do_acknowledgements(self.acknowledgements); + + let connector = self.build_connector(acknowledgements)?; let decoder = DecodingConfig::new(self.framing.clone(), self.decoding.clone(), log_namespace) .build()?; - let sink = MqttSource::new(connector.clone(), decoder, log_namespace, self.clone())?; - Ok(Box::pin(sink.run(cx.out, cx.shutdown))) + let source = MqttSource::new( + connector.clone(), + decoder, + log_namespace, + self.clone(), + acknowledgements, + )?; + Ok(Box::pin(source.run(cx.out, cx.shutdown))) } fn outputs(&self, global_log_namespace: LogNamespace) -> Vec { @@ -107,12 +119,22 @@ impl SourceConfig for MqttSourceConfig { } fn can_acknowledge(&self) -> bool { - false + true } } impl MqttSourceConfig { - fn build_connector(&self) -> Result { + fn build_connector(&self, acknowledgements: bool) -> Result { + // End-to-end acknowledgements rely on resuming the MQTT session (and its + // unacknowledged in-flight messages) after a restart, which is keyed by the + // client ID. A generated/random client ID would start a fresh session and + // orphan those messages, silently breaking at-least-once — so require an + // explicit, stable client ID when acknowledgements are enabled. + if acknowledgements && self.common.client_id.is_none() { + return Err(ConfigurationError::AcknowledgementsRequireClientId) + .context(ConfigurationSnafu); + } + let client_id = self.common.client_id.clone().unwrap_or_else(|| { let hash = rand::rng() .sample_iter(&rand_distr::Alphanumeric) @@ -133,6 +155,16 @@ impl MqttSourceConfig { options.set_max_packet_size(self.common.max_packet_size, self.common.max_packet_size); options.set_clean_session(false); + + // With end-to-end acknowledgements enabled, defer the QoS-1 PUBACK until + // the event has been delivered to all sinks. rumqttc then requires every + // incoming publish to be acked explicitly via `client.ack(&publish)`. + // Combined with `clean_session(false)` and QoS `AtLeastOnce`, an unacked + // message is redelivered by the broker after a crash/reconnect. + if acknowledgements { + options.set_manual_acks(true); + } + match (&self.common.user, &self.common.password) { (Some(user), Some(password)) => { options.set_credentials(user, password); @@ -171,4 +203,25 @@ mod test { fn generate_config() { crate::test_util::test_generate_config::(); } + + #[test] + fn acknowledgements_require_a_stable_client_id() { + // Without acks, a client ID is auto-generated — fine. + let default_config = MqttSourceConfig::default(); + assert!(default_config.build_connector(false).is_ok()); + + // With acks and no explicit client ID, building must fail (a generated ID + // would orphan the session's unacknowledged messages after a restart). + assert!(default_config.build_connector(true).is_err()); + + // With acks and an explicit client ID, building succeeds. + let with_client_id = MqttSourceConfig { + common: MqttCommonConfig { + client_id: Some("stable-id".to_owned()), + ..Default::default() + }, + ..Default::default() + }; + assert!(with_client_id.build_connector(true).is_ok()); + } } diff --git a/src/sources/mqtt/integration_tests.rs b/src/sources/mqtt/integration_tests.rs index ed1013f4bef0f..45662af4c77ab 100644 --- a/src/sources/mqtt/integration_tests.rs +++ b/src/sources/mqtt/integration_tests.rs @@ -11,7 +11,7 @@ use crate::{ SourceSender, common::mqtt::MqttCommonConfig, config::{SourceConfig, SourceContext, log_schema}, - event::Event, + event::{Event, EventStatus}, serde::OneOrMany, sources::mqtt::MqttSourceConfig, test_util::{ @@ -40,9 +40,21 @@ async fn send_test_events(client: &AsyncClient, topic: &str, messages: &Vec String { + event + .as_log() + .get(log_schema().message_key_target_path().unwrap()) + .unwrap() + .to_string_lossy() + .into_owned() +} + async fn get_mqtt_client() -> AsyncClient { + // Unique client ID per producer: brokers that strictly enforce client-ID + // uniqueness (e.g. RabbitMQ) otherwise kick a previous connection when tests + // run concurrently, which manifests as spurious publish timeouts. let mut mqtt_options = MqttOptions::new( - "integration-test-producer", + format!("integration-test-producer-{}", random_string(6)), mqtt_broker_address(), mqtt_broker_port(), ); @@ -118,6 +130,99 @@ async fn mqtt_one_topic_happy() { .await; } +/// With end-to-end acknowledgements enabled, a message that is received but not +/// successfully delivered (the sink rejects it) must not be acknowledged to the +/// broker, so the broker redelivers it. This proves the at-least-once guarantee +/// added by the `acknowledgements` option: no data is lost when a downstream +/// failure or crash occurs before the write is confirmed. +#[tokio::test] +async fn mqtt_redelivers_unacknowledged_messages() { + trace_init(); + + let topic = "source-redelivery-test"; + // A stable client ID so the second connection resumes the same persistent + // session (the source sets `clean_session = false`); the broker then + // redelivers any in-flight QoS 1 message that was never acknowledged. + let client_id = format!("sourceAckTest{}", random_string(6)); + let message = random_string(32); + + let make_config = || MqttSourceConfig { + common: MqttCommonConfig { + host: mqtt_broker_address(), + port: mqtt_broker_port(), + client_id: Some(client_id.clone()), + ..Default::default() + }, + topic: OneOrMany::One(topic.to_owned()), + acknowledgements: true.into(), + ..MqttSourceConfig::default() + }; + + // Phase 1: the first instance subscribes (creating the persistent session) + // and receives the message, but its sink rejects every event, so the source + // never sends a PUBACK. + let (tx1, mut rx1) = SourceSender::new_test_finalize(EventStatus::Rejected); + let config1 = make_config(); + let source1 = tokio::spawn(async move { + config1 + .build(SourceContext::new_test(tx1, None)) + .await + .unwrap() + .await + .unwrap() + }); + + // Wait for the subscription to be established before publishing. + tokio::time::sleep(Duration::from_millis(500)).await; + + let producer = get_mqtt_client().await; + producer + .publish(topic, QoS::AtLeastOnce, false, message.as_bytes()) + .await + .unwrap(); + + // The first instance must actually receive (and reject) the message so that + // it remains unacknowledged in the broker. + let first = timeout(Duration::from_secs(5), rx1.next()) + .await + .expect("timed out waiting for first delivery") + .expect("source stream ended unexpectedly"); + assert_eq!(message_body(&first), message); + drop(first); + + // Give the source a moment to observe the rejected status (and therefore + // skip the ack), then drop the connection without acknowledging. + tokio::time::sleep(Duration::from_millis(200)).await; + source1.abort(); + drop(source1.await); + + // Phase 2: a new instance resumes the same session; the broker must + // redeliver the unacknowledged message. + let (tx2, mut rx2) = SourceSender::new_test(); + let config2 = make_config(); + let source2 = tokio::spawn(async move { + config2 + .build(SourceContext::new_test(tx2, None)) + .await + .unwrap() + .await + .unwrap() + }); + + let redelivered = timeout(Duration::from_secs(10), rx2.next()) + .await + .expect("timed out waiting for redelivery: the message was lost") + .expect("source stream ended unexpectedly"); + assert_eq!( + message_body(&redelivered), + message, + "redelivered message did not match the original" + ); + + source2.abort(); + drop(source2.await); +} + #[tokio::test] async fn mqtt_many_topics_happy() { trace_init(); diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index d0f3efea5b0f7..d5ed5b13cfcae 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1,8 +1,10 @@ +use futures::StreamExt; use itertools::Itertools; use rumqttc::{Event as MqttEvent, Incoming, Publish, QoS, SubscribeFilter}; use vector_lib::{ codecs::Decoder, config::{LegacyKey, LogNamespace}, + finalizer::UnorderedFinalizer, internal_event::EventsReceived, lookup::path, }; @@ -10,18 +12,116 @@ use vector_lib::{ use crate::{ SourceSender, common::mqtt::MqttConnector, - event::{BatchNotifier, Event}, + event::{BatchNotifier, BatchStatus, Event}, internal_events::{EndpointBytesReceived, StreamClosedError}, serde::OneOrMany, shutdown::ShutdownSignal, sources::{mqtt::MqttSourceConfig, util}, }; +/// Identifies an in-flight publish so its QoS-1 PUBACK can be sent once the +/// downstream sinks confirm delivery. Only the packet id (carried by `Publish`) +/// is needed to ack; the payload is cleared before the entry is retained so +/// pending acks don't pin payloads in memory under backpressure. +#[derive(Clone, Debug)] +struct FinalizerEntry { + publish: Publish, + connection_generation: u64, +} + +fn should_ack_finalized_publish( + status: BatchStatus, + entry_generation: u64, + connection_generation: u64, +) -> bool { + status == BatchStatus::Delivered && entry_generation == connection_generation +} + +fn publish_supports_end_to_end_acknowledgements(qos: QoS) -> bool { + qos != QoS::AtMostOnce +} + +fn should_defer_publish_ack(acknowledgements: bool, qos: QoS) -> bool { + acknowledgements && publish_supports_end_to_end_acknowledgements(qos) +} + +fn warn_unsupported_acknowledgement_qos(qos: QoS, topic: &str) { + warn!( + message = "MQTT acknowledgements require publishes with QoS 1 or greater; forwarding message without end-to-end acknowledgement guarantee.", + ?qos, + topic, + internal_log_rate_limit = true, + ); +} + +const fn should_warn_session_not_resumed(acknowledgements: bool, session_present: bool) -> bool { + acknowledgements && !session_present +} + +const fn should_resubscribe_after_connack(has_connected: bool, session_present: bool) -> bool { + has_connected && !session_present +} + +fn warn_session_not_resumed() { + warn!( + message = "MQTT broker started a new session while acknowledgements are enabled; unacknowledged messages from any previous session for this client ID will not be redelivered.", + internal_log_rate_limit = true, + ); +} + +fn warn_resubscribe_failed() { + warn!( + message = "Failed to queue MQTT re-subscribe request after reconnect; will retry while connected.", + internal_log_rate_limit = true, + ); +} + +#[derive(Default)] +struct PendingAcks { + publishes: Vec, +} + +impl PendingAcks { + fn push(&mut self, publish: Publish) { + self.publishes.push(publish); + } + + fn clear(&mut self) { + self.publishes.clear(); + } + + fn retry(&mut self, client: &rumqttc::AsyncClient) { + self.retry_with(|publish| client.try_ack(publish).is_ok()); + } + + fn try_ack(&mut self, connected: bool, publish: Publish, client: &rumqttc::AsyncClient) { + self.try_ack_with(connected, publish, |publish| { + client.try_ack(publish).is_ok() + }); + } + + fn try_ack_with( + &mut self, + connected: bool, + publish: Publish, + mut try_ack: impl FnMut(&Publish) -> bool, + ) { + if connected && !try_ack(&publish) { + self.push(publish); + } + } + + fn retry_with(&mut self, mut try_ack: impl FnMut(&Publish) -> bool) { + self.publishes.retain(|publish| !try_ack(publish)); + } +} + pub struct MqttSource { connector: MqttConnector, decoder: Decoder, log_namespace: LogNamespace, config: MqttSourceConfig, + acknowledgements: bool, } impl MqttSource { @@ -30,55 +130,145 @@ impl MqttSource { decoder: Decoder, log_namespace: LogNamespace, config: MqttSourceConfig, + acknowledgements: bool, ) -> crate::Result { Ok(Self { connector, decoder, log_namespace, config, + acknowledgements, }) } pub async fn run(self, mut out: SourceSender, shutdown: ShutdownSignal) -> Result<(), ()> { let (client, mut connection) = self.connector.connect(); - match &self.config.topic { - OneOrMany::One(topic) => { - client - .subscribe(topic, QoS::AtLeastOnce) - .await - .map_err(|_| ())?; + self.subscribe(&client)?; + + // Finalizer drives end-to-end acknowledgements: each in-flight publish is + // registered with its batch-status receiver, and we send the QoS-1 PUBACK + // only once the sinks report `Delivered`. Unused when acknowledgements are + // disabled (rumqttc auto-acks in that mode). MQTT PUBACKs are independent + // per packet id (unlike Kafka offsets), so finalization is unordered — a + // slow/stuck batch must not hold back acks for already-delivered publishes. + let (finalizer, mut ack_stream) = UnorderedFinalizer::::maybe_new( + self.acknowledgements, + Some(shutdown.clone()), + ); + + // PUBACKs that rumqttc's bounded request channel was too full to accept, + // retained for retry rather than dropped. Dropping a PUBACK for an already + // delivered message would pin it in the broker's in-flight window until the + // next reconnect. This is bounded in practice by that in-flight window (the + // broker stops delivering once it fills), and the event loop below drains the + // request channel, so entries flush on subsequent iterations. + let mut pending_acks = PendingAcks::default(); + let mut connected = false; + let mut has_connected = false; + let mut pending_resubscribe = false; + let mut connection_generation = 0; + + loop { + if connected && pending_resubscribe { + pending_resubscribe = !self.try_subscribe(&client); } - OneOrMany::Many(topics) => { - client - .subscribe_many( - topics - .iter() - .cloned() - .map(|topic| SubscribeFilter::new(topic, QoS::AtLeastOnce)), - ) - .await - .map_err(|_| ())?; + + // Retry deferred PUBACKs while connected (the event loop below drains the + // request channel). Skipped while disconnected: a publish's packet id is + // only valid on the connection it arrived on, so stale PUBACKs must not be + // replayed across a reconnect. + if connected { + pending_acks.retry(&client); } - } - loop { tokio::select! { _ = shutdown.clone() => return Ok(()), + entry = ack_stream.next() => { + // Only PUBACK delivered events. On Errored/Rejected we skip the + // ack so the broker redelivers after reconnect (QoS-1 + + // clean_session=false), giving at-least-once delivery. Use the + // non-blocking `try_ack` — awaiting `ack` could deadlock, since + // this same task polls the event loop that drains rumqttc's request + // channel. If that channel is full, retain the PUBACK for retry + // (above) instead of dropping it. + if let Some((status, entry)) = entry + && should_ack_finalized_publish( + status, + entry.connection_generation, + connection_generation, + ) + { + pending_acks.try_ack(connected, entry.publish, &client); + } + }, mqtt_event = connection.poll() => { - // If an error is returned here there is currently no way to tie this back - // to the event that was posted which means we can't accurately provide - // delivery guarantees. - // We need this issue resolved first: - // https://github.com/bytebeamio/rumqtt/issues/349 + // Providing at-least-once here does not require correlating a + // connection/poll error back to a specific in-flight publish. + // rumqtt#349 (no packet id for *outbound* publishes) concerns the + // publish/sink direction and does not apply to a subscribe-only + // source: each incoming Publish already carries its packet id, and + // we withhold its QoS-1 PUBACK until the event is delivered + // end-to-end. Anything left unacked when the connection drops is + // redelivered by the broker on reconnect (clean_session=false + QoS + // AtLeastOnce). match mqtt_event { Ok(MqttEvent::Incoming(Incoming::Publish(publish))) => { - self.process_message(publish, &mut out).await; + self.process_message( + publish, + &mut out, + finalizer.as_ref(), + connection_generation, + ).await; + } + Ok(MqttEvent::Incoming(Incoming::SubAck(suback))) + if self.acknowledgements => + { + for return_code in suback.return_codes { + if let rumqttc::SubscribeReasonCode::Success(qos) = return_code + && !publish_supports_end_to_end_acknowledgements(qos) + { + warn!( + message = "MQTT broker granted a subscription QoS below the level required for end-to-end acknowledgements.", + ?qos, + internal_log_rate_limit = true, + ); + } + } + } + // A (re)connected session resumes here; the broker will + // redeliver any unacknowledged publishes, so drop deferred + // PUBACKs whose packet ids came from the previous connection. + Ok(MqttEvent::Incoming(Incoming::ConnAck(connack))) => { + let needs_resubscribe = should_resubscribe_after_connack( + has_connected, + connack.session_present, + ); + if should_warn_session_not_resumed( + self.acknowledgements, + connack.session_present, + ) { + warn_session_not_resumed(); + } + connected = true; + has_connected = true; + connection_generation += 1; + pending_acks.clear(); + if let Some(finalizer) = &finalizer { + finalizer.flush(); + } + if needs_resubscribe && !self.try_subscribe(&client) { + pending_resubscribe = true; + } } - Ok(MqttEvent::Incoming( - Incoming::PubAck(_) | Incoming::PubRec(_) | Incoming::PubComp(_), - )) => { - // TODO Handle acknowledgement - https://github.com/vectordotdev/vector/issues/21967 + // Connection lost: same stale-packet-id reasoning, and rumqttc + // drops its own queued acks while reconnecting. + Err(_) => { + connected = false; + pending_acks.clear(); + if let Some(finalizer) = &finalizer { + finalizer.flush(); + } } _ => {} } @@ -87,7 +277,39 @@ impl MqttSource { } } - async fn process_message(&self, publish: Publish, out: &mut SourceSender) { + fn try_subscribe(&self, client: &rumqttc::AsyncClient) -> bool { + match self.subscribe(client) { + Ok(()) => true, + Err(()) => { + warn_resubscribe_failed(); + false + } + } + } + + fn subscribe(&self, client: &rumqttc::AsyncClient) -> Result<(), ()> { + match &self.config.topic { + OneOrMany::One(topic) => client + .try_subscribe(topic, QoS::AtLeastOnce) + .map_err(|_| ()), + OneOrMany::Many(topics) => client + .try_subscribe_many( + topics + .iter() + .cloned() + .map(|topic| SubscribeFilter::new(topic, QoS::AtLeastOnce)), + ) + .map_err(|_| ()), + } + } + + async fn process_message( + &self, + mut publish: Publish, + out: &mut SourceSender, + finalizer: Option<&UnorderedFinalizer>, + connection_generation: u64, + ) { emit!(EndpointBytesReceived { byte_size: publish.payload.len(), protocol: "mqtt", @@ -95,7 +317,14 @@ impl MqttSource { }); let events_received = register!(EventsReceived); - let (batch, _batch_receiver) = BatchNotifier::maybe_new_with_receiver(false); + let use_end_to_end_acknowledgements = + should_defer_publish_ack(finalizer.is_some(), publish.qos); + if finalizer.is_some() && !use_end_to_end_acknowledgements { + warn_unsupported_acknowledgement_qos(publish.qos, &publish.topic); + } + + let (batch, batch_receiver) = + BatchNotifier::maybe_new_with_receiver(use_end_to_end_acknowledgements); // Error is logged by `vector_lib::codecs::Decoder`, no further handling // is needed here. let decoded = util::decode_message( @@ -116,7 +345,23 @@ impl MqttSource { let count = decoded.len(); match out.send_batch(decoded).await { - Ok(()) => {} + Ok(()) => { + // Register the publish for deferred PUBACK once the batch is + // delivered. Without acknowledgements `batch_receiver` is None and + // rumqttc has already auto-acked. The payload is no longer needed + // (ack only uses the packet id), so clear it before retaining the + // entry to avoid pinning payloads in memory while sinks process. + if let Some((finalizer, receiver)) = finalizer.zip(batch_receiver) { + publish.payload = Default::default(); + finalizer.add( + FinalizerEntry { + publish, + connection_generation, + }, + receiver, + ); + } + } Err(_) => emit!(StreamClosedError { count }), } } @@ -137,3 +382,90 @@ impl MqttSource { } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn publish(pkid: u16) -> Publish { + let mut publish = Publish::new("topic", QoS::AtLeastOnce, vec![1, 2, 3]); + publish.pkid = pkid; + publish + } + + #[test] + fn pending_acks_keeps_failed_retries() { + let mut pending_acks = PendingAcks::default(); + pending_acks.push(publish(1)); + pending_acks.push(publish(2)); + pending_acks.push(publish(3)); + + let mut attempted = Vec::new(); + pending_acks.retry_with(|publish| { + attempted.push(publish.pkid); + publish.pkid != 2 + }); + + assert_eq!(attempted, vec![1, 2, 3]); + assert_eq!(pending_acks.publishes.len(), 1); + assert_eq!(pending_acks.publishes[0].pkid, 2); + + pending_acks.retry_with(|_| true); + assert!(pending_acks.publishes.is_empty()); + } + + #[test] + fn pending_acks_clear_drops_stale_packet_ids() { + let mut pending_acks = PendingAcks::default(); + pending_acks.push(publish(1)); + pending_acks.push(publish(2)); + + pending_acks.clear(); + + assert!(pending_acks.publishes.is_empty()); + } + + #[test] + fn pending_acks_drops_finalized_publish_while_disconnected() { + let mut pending_acks = PendingAcks::default(); + let mut attempted = false; + + pending_acks.try_ack_with(false, publish(1), |_| { + attempted = true; + true + }); + + assert!(!attempted); + assert!(pending_acks.publishes.is_empty()); + } + + #[test] + fn finalized_publish_must_match_current_connection_generation() { + assert!(should_ack_finalized_publish(BatchStatus::Delivered, 2, 2)); + assert!(!should_ack_finalized_publish(BatchStatus::Delivered, 1, 2)); + assert!(!should_ack_finalized_publish(BatchStatus::Rejected, 2, 2)); + } + + #[test] + fn qos0_publish_does_not_defer_acknowledgement() { + assert!(!should_defer_publish_ack(true, QoS::AtMostOnce)); + assert!(should_defer_publish_ack(true, QoS::AtLeastOnce)); + assert!(should_defer_publish_ack(true, QoS::ExactlyOnce)); + assert!(!should_defer_publish_ack(false, QoS::AtLeastOnce)); + } + + #[test] + fn warns_when_acknowledgement_session_is_not_resumed() { + assert!(should_warn_session_not_resumed(true, false)); + assert!(!should_warn_session_not_resumed(true, true)); + assert!(!should_warn_session_not_resumed(false, false)); + } + + #[test] + fn resubscribes_when_reconnect_starts_fresh_session() { + assert!(!should_resubscribe_after_connack(false, false)); + assert!(!should_resubscribe_after_connack(false, true)); + assert!(!should_resubscribe_after_connack(true, true)); + assert!(should_resubscribe_after_connack(true, false)); + } +} diff --git a/website/cue/reference/components/sources/generated/mqtt.cue b/website/cue/reference/components/sources/generated/mqtt.cue index 4a699a4e6951a..75f556b472399 100644 --- a/website/cue/reference/components/sources/generated/mqtt.cue +++ b/website/cue/reference/components/sources/generated/mqtt.cue @@ -1,6 +1,27 @@ package metadata generated: components: sources: mqtt: configuration: { + acknowledgements: { + deprecated: true + description: """ + Controls how acknowledgements are handled by this source. + + This setting is **deprecated** in favor of enabling `acknowledgements` at the [global][global_acks] or sink level. + + Enabling or disabling acknowledgements at the source level has **no effect** on acknowledgement behavior. + + See [End-to-end Acknowledgements][e2e_acks] for more information on how event acknowledgement is handled. + + [global_acks]: https://vector.dev/docs/reference/configuration/global-options/#acknowledgements + [e2e_acks]: https://vector.dev/docs/architecture/end-to-end-acknowledgements/ + """ + required: false + type: object: options: enabled: { + description: "Whether or not end-to-end acknowledgements are enabled for this source." + required: false + type: bool: {} + } + } client_id: { description: "MQTT client ID." required: false From 98217901df63092a2b223416fe567541c3390c81 Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 1 Jul 2026 13:42:06 -0700 Subject: [PATCH 2/5] add more tests to all possible state transitions trying to get ahead of lots of smaller comments on how this works by moving to more state machine like helpers into more dedicated structs --- src/sources/mqtt/source.rs | 383 +++++++++++++++++++++++++++++-------- 1 file changed, 308 insertions(+), 75 deletions(-) diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index d5ed5b13cfcae..616cb54914af4 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -19,6 +19,8 @@ use crate::{ sources::{mqtt::MqttSourceConfig, util}, }; +const SUBSCRIPTION_QOS: QoS = QoS::AtLeastOnce; + /// Identifies an in-flight publish so its QoS-1 PUBACK can be sent once the /// downstream sinks confirm delivery. Only the packet id (carried by `Publish`) /// is needed to ack; the payload is cleared before the entry is retained so @@ -29,20 +31,94 @@ struct FinalizerEntry { connection_generation: u64, } -fn should_ack_finalized_publish( - status: BatchStatus, - entry_generation: u64, +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct ProtocolState { + connected: bool, + pending_resubscribe: bool, connection_generation: u64, -) -> bool { - status == BatchStatus::Delivered && entry_generation == connection_generation +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct LoopActions { + retry_pending_acks: bool, + retry_resubscribe: bool, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct ConnAckActions { + warn_session_not_resumed: bool, + clear_pending_acks: bool, + flush_finalizer: bool, + resubscribe: bool, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct DisconnectActions { + clear_pending_acks: bool, + flush_finalizer: bool, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct PublishAckDecision { + defer_ack: bool, + warn_unsupported_qos: bool, +} + +impl ProtocolState { + const fn loop_actions(&self) -> LoopActions { + LoopActions { + retry_pending_acks: self.connected, + retry_resubscribe: self.connected && self.pending_resubscribe, + } + } + + const fn on_connack( + &mut self, + acknowledgements: bool, + session_present: bool, + ) -> ConnAckActions { + let actions = ConnAckActions { + warn_session_not_resumed: acknowledgements && !session_present, + clear_pending_acks: true, + flush_finalizer: true, + resubscribe: self.connection_generation > 0 && !session_present, + }; + + self.connected = true; + self.connection_generation += 1; + + actions + } + + const fn on_disconnect(&mut self) -> DisconnectActions { + self.connected = false; + + DisconnectActions { + clear_pending_acks: true, + flush_finalizer: true, + } + } + + const fn on_resubscribe_result(&mut self, success: bool) { + self.pending_resubscribe = !success; + } + + fn should_ack_finalized_publish(&self, status: BatchStatus, entry_generation: u64) -> bool { + status == BatchStatus::Delivered && entry_generation == self.connection_generation + } } fn publish_supports_end_to_end_acknowledgements(qos: QoS) -> bool { qos != QoS::AtMostOnce } -fn should_defer_publish_ack(acknowledgements: bool, qos: QoS) -> bool { - acknowledgements && publish_supports_end_to_end_acknowledgements(qos) +fn publish_ack_decision(acknowledgements: bool, qos: QoS) -> PublishAckDecision { + let defer_ack = acknowledgements && publish_supports_end_to_end_acknowledgements(qos); + + PublishAckDecision { + defer_ack, + warn_unsupported_qos: acknowledgements && !defer_ack, + } } fn warn_unsupported_acknowledgement_qos(qos: QoS, topic: &str) { @@ -54,14 +130,6 @@ fn warn_unsupported_acknowledgement_qos(qos: QoS, topic: &str) { ); } -const fn should_warn_session_not_resumed(acknowledgements: bool, session_present: bool) -> bool { - acknowledgements && !session_present -} - -const fn should_resubscribe_after_connack(has_connected: bool, session_present: bool) -> bool { - has_connected && !session_present -} - fn warn_session_not_resumed() { warn!( message = "MQTT broker started a new session while acknowledgements are enabled; unacknowledged messages from any previous session for this client ID will not be redelivered.", @@ -163,22 +231,20 @@ impl MqttSource { // next reconnect. This is bounded in practice by that in-flight window (the // broker stops delivering once it fills), and the event loop below drains the // request channel, so entries flush on subsequent iterations. + let mut protocol_state = ProtocolState::default(); let mut pending_acks = PendingAcks::default(); - let mut connected = false; - let mut has_connected = false; - let mut pending_resubscribe = false; - let mut connection_generation = 0; loop { - if connected && pending_resubscribe { - pending_resubscribe = !self.try_subscribe(&client); + let actions = protocol_state.loop_actions(); + if actions.retry_resubscribe { + protocol_state.on_resubscribe_result(self.try_subscribe(&client)); } // Retry deferred PUBACKs while connected (the event loop below drains the // request channel). Skipped while disconnected: a publish's packet id is // only valid on the connection it arrived on, so stale PUBACKs must not be // replayed across a reconnect. - if connected { + if actions.retry_pending_acks { pending_acks.retry(&client); } @@ -193,13 +259,12 @@ impl MqttSource { // channel. If that channel is full, retain the PUBACK for retry // (above) instead of dropping it. if let Some((status, entry)) = entry - && should_ack_finalized_publish( + && protocol_state.should_ack_finalized_publish( status, entry.connection_generation, - connection_generation, ) { - pending_acks.try_ack(connected, entry.publish, &client); + pending_acks.try_ack(protocol_state.connected, entry.publish, &client); } }, mqtt_event = connection.poll() => { @@ -218,7 +283,7 @@ impl MqttSource { publish, &mut out, finalizer.as_ref(), - connection_generation, + protocol_state.connection_generation, ).await; } Ok(MqttEvent::Incoming(Incoming::SubAck(suback))) @@ -240,33 +305,35 @@ impl MqttSource { // redeliver any unacknowledged publishes, so drop deferred // PUBACKs whose packet ids came from the previous connection. Ok(MqttEvent::Incoming(Incoming::ConnAck(connack))) => { - let needs_resubscribe = should_resubscribe_after_connack( - has_connected, - connack.session_present, - ); - if should_warn_session_not_resumed( + let actions = protocol_state.on_connack( self.acknowledgements, connack.session_present, - ) { + ); + if actions.warn_session_not_resumed { warn_session_not_resumed(); } - connected = true; - has_connected = true; - connection_generation += 1; - pending_acks.clear(); - if let Some(finalizer) = &finalizer { + if actions.clear_pending_acks { + pending_acks.clear(); + } + if actions.flush_finalizer + && let Some(finalizer) = &finalizer + { finalizer.flush(); } - if needs_resubscribe && !self.try_subscribe(&client) { - pending_resubscribe = true; + if actions.resubscribe { + protocol_state.on_resubscribe_result(self.try_subscribe(&client)); } } // Connection lost: same stale-packet-id reasoning, and rumqttc // drops its own queued acks while reconnecting. Err(_) => { - connected = false; - pending_acks.clear(); - if let Some(finalizer) = &finalizer { + let actions = protocol_state.on_disconnect(); + if actions.clear_pending_acks { + pending_acks.clear(); + } + if actions.flush_finalizer + && let Some(finalizer) = &finalizer + { finalizer.flush(); } } @@ -290,14 +357,14 @@ impl MqttSource { fn subscribe(&self, client: &rumqttc::AsyncClient) -> Result<(), ()> { match &self.config.topic { OneOrMany::One(topic) => client - .try_subscribe(topic, QoS::AtLeastOnce) + .try_subscribe(topic, SUBSCRIPTION_QOS) .map_err(|_| ()), OneOrMany::Many(topics) => client .try_subscribe_many( topics .iter() .cloned() - .map(|topic| SubscribeFilter::new(topic, QoS::AtLeastOnce)), + .map(|topic| SubscribeFilter::new(topic, SUBSCRIPTION_QOS)), ) .map_err(|_| ()), } @@ -317,14 +384,13 @@ impl MqttSource { }); let events_received = register!(EventsReceived); - let use_end_to_end_acknowledgements = - should_defer_publish_ack(finalizer.is_some(), publish.qos); - if finalizer.is_some() && !use_end_to_end_acknowledgements { + let ack_decision = publish_ack_decision(finalizer.is_some(), publish.qos); + if ack_decision.warn_unsupported_qos { warn_unsupported_acknowledgement_qos(publish.qos, &publish.topic); } let (batch, batch_receiver) = - BatchNotifier::maybe_new_with_receiver(use_end_to_end_acknowledgements); + BatchNotifier::maybe_new_with_receiver(ack_decision.defer_ack); // Error is logged by `vector_lib::codecs::Decoder`, no further handling // is needed here. let decoded = util::decode_message( @@ -426,46 +492,213 @@ mod tests { } #[test] - fn pending_acks_drops_finalized_publish_while_disconnected() { - let mut pending_acks = PendingAcks::default(); - let mut attempted = false; + fn pending_acks_backpressure_matrix() { + for (connected, try_ack_succeeds, expected_attempted, expected_queued) in [ + (false, true, false, false), + (true, true, true, false), + (true, false, true, true), + ] { + let mut pending_acks = PendingAcks::default(); + let mut attempted = false; + + pending_acks.try_ack_with(connected, publish(1), |_| { + attempted = true; + try_ack_succeeds + }); + + assert_eq!(attempted, expected_attempted); + assert_eq!(!pending_acks.publishes.is_empty(), expected_queued); + } + } - pending_acks.try_ack_with(false, publish(1), |_| { - attempted = true; - true - }); + #[test] + fn protocol_contract_matrix_for_requested_and_granted_qos() { + assert!(publish_supports_end_to_end_acknowledgements( + SUBSCRIPTION_QOS + )); + + for (granted_qos, supports_acknowledgements) in [ + (QoS::AtMostOnce, false), + (QoS::AtLeastOnce, true), + (QoS::ExactlyOnce, true), + ] { + assert_eq!( + publish_supports_end_to_end_acknowledgements(granted_qos), + supports_acknowledgements + ); + } + } - assert!(!attempted); - assert!(pending_acks.publishes.is_empty()); + #[test] + fn protocol_contract_matrix_for_publisher_qos() { + for (acknowledgements, publisher_qos, expected) in [ + ( + false, + QoS::AtMostOnce, + PublishAckDecision { + defer_ack: false, + warn_unsupported_qos: false, + }, + ), + ( + false, + QoS::AtLeastOnce, + PublishAckDecision { + defer_ack: false, + warn_unsupported_qos: false, + }, + ), + ( + true, + QoS::AtMostOnce, + PublishAckDecision { + defer_ack: false, + warn_unsupported_qos: true, + }, + ), + ( + true, + QoS::AtLeastOnce, + PublishAckDecision { + defer_ack: true, + warn_unsupported_qos: false, + }, + ), + ( + true, + QoS::ExactlyOnce, + PublishAckDecision { + defer_ack: true, + warn_unsupported_qos: false, + }, + ), + ] { + assert_eq!( + publish_ack_decision(acknowledgements, publisher_qos), + expected + ); + } } #[test] - fn finalized_publish_must_match_current_connection_generation() { - assert!(should_ack_finalized_publish(BatchStatus::Delivered, 2, 2)); - assert!(!should_ack_finalized_publish(BatchStatus::Delivered, 1, 2)); - assert!(!should_ack_finalized_publish(BatchStatus::Rejected, 2, 2)); + fn protocol_contract_matrix_for_session_reset_and_connection_generation() { + for (acknowledgements, session_present, expected_warn) in [ + (false, false, false), + (true, false, true), + (true, true, false), + ] { + let mut state = ProtocolState::default(); + let actions = state.on_connack(acknowledgements, session_present); + + assert_eq!(actions.warn_session_not_resumed, expected_warn); + assert!(actions.clear_pending_acks); + assert!(actions.flush_finalizer); + assert!(!actions.resubscribe); + assert!(state.connected); + assert_eq!(state.connection_generation, 1); + } + + let mut resumed_session = ProtocolState::default(); + resumed_session.on_connack(true, true); + let actions = resumed_session.on_connack(true, true); + assert_eq!( + actions, + ConnAckActions { + warn_session_not_resumed: false, + clear_pending_acks: true, + flush_finalizer: true, + resubscribe: false, + } + ); + assert_eq!(resumed_session.connection_generation, 2); + + let mut fresh_session = ProtocolState::default(); + fresh_session.on_connack(true, true); + let actions = fresh_session.on_connack(true, false); + assert_eq!( + actions, + ConnAckActions { + warn_session_not_resumed: true, + clear_pending_acks: true, + flush_finalizer: true, + resubscribe: true, + } + ); + assert_eq!(fresh_session.connection_generation, 2); } #[test] - fn qos0_publish_does_not_defer_acknowledgement() { - assert!(!should_defer_publish_ack(true, QoS::AtMostOnce)); - assert!(should_defer_publish_ack(true, QoS::AtLeastOnce)); - assert!(should_defer_publish_ack(true, QoS::ExactlyOnce)); - assert!(!should_defer_publish_ack(false, QoS::AtLeastOnce)); + fn protocol_contract_matrix_for_pending_resubscribe() { + let mut state = ProtocolState::default(); + state.on_resubscribe_result(false); + assert_eq!( + state.loop_actions(), + LoopActions { + retry_pending_acks: false, + retry_resubscribe: false, + } + ); + + state.on_connack(true, true); + assert_eq!( + state.loop_actions(), + LoopActions { + retry_pending_acks: true, + retry_resubscribe: true, + } + ); + + state.on_resubscribe_result(true); + assert_eq!( + state.loop_actions(), + LoopActions { + retry_pending_acks: true, + retry_resubscribe: false, + } + ); } #[test] - fn warns_when_acknowledgement_session_is_not_resumed() { - assert!(should_warn_session_not_resumed(true, false)); - assert!(!should_warn_session_not_resumed(true, true)); - assert!(!should_warn_session_not_resumed(false, false)); + fn protocol_contract_matrix_for_disconnect() { + let mut state = ProtocolState::default(); + state.on_connack(true, true); + + let actions = state.on_disconnect(); + + assert_eq!( + actions, + DisconnectActions { + clear_pending_acks: true, + flush_finalizer: true, + } + ); + assert!(!state.connected); + assert_eq!(state.connection_generation, 1); + assert_eq!( + state.loop_actions(), + LoopActions { + retry_pending_acks: false, + retry_resubscribe: false, + } + ); } #[test] - fn resubscribes_when_reconnect_starts_fresh_session() { - assert!(!should_resubscribe_after_connack(false, false)); - assert!(!should_resubscribe_after_connack(false, true)); - assert!(!should_resubscribe_after_connack(true, true)); - assert!(should_resubscribe_after_connack(true, false)); + fn protocol_contract_matrix_for_finalization_statuses() { + let mut state = ProtocolState::default(); + state.on_connack(true, true); + state.on_connack(true, true); + + for (status, entry_generation, should_ack) in [ + (BatchStatus::Delivered, 2, true), + (BatchStatus::Delivered, 1, false), + (BatchStatus::Errored, 2, false), + (BatchStatus::Rejected, 2, false), + ] { + assert_eq!( + state.should_ack_finalized_publish(status, entry_generation), + should_ack + ); + } } } From 41011d144da107f2fde77ae684d6fe5e583417eb Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 1 Jul 2026 14:33:15 -0700 Subject: [PATCH 3/5] new design that adds a new task this removes much of the tracing by storing state within the task --- src/sources/mqtt/source.rs | 292 +++++++++---------------------------- 1 file changed, 70 insertions(+), 222 deletions(-) diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index 616cb54914af4..d94e2d703e4ea 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1,6 +1,9 @@ use futures::StreamExt; use itertools::Itertools; -use rumqttc::{Event as MqttEvent, Incoming, Publish, QoS, SubscribeFilter}; +use rumqttc::{ + AsyncClient, Event as MqttEvent, EventLoop, Incoming, Publish, QoS, SubscribeFilter, +}; +use tokio::sync::mpsc; use vector_lib::{ codecs::Decoder, config::{LegacyKey, LogNamespace}, @@ -34,27 +37,18 @@ struct FinalizerEntry { #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] struct ProtocolState { connected: bool, - pending_resubscribe: bool, connection_generation: u64, } -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] -struct LoopActions { - retry_pending_acks: bool, - retry_resubscribe: bool, -} - #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] struct ConnAckActions { warn_session_not_resumed: bool, - clear_pending_acks: bool, flush_finalizer: bool, resubscribe: bool, } #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] struct DisconnectActions { - clear_pending_acks: bool, flush_finalizer: bool, } @@ -64,14 +58,12 @@ struct PublishAckDecision { warn_unsupported_qos: bool, } -impl ProtocolState { - const fn loop_actions(&self) -> LoopActions { - LoopActions { - retry_pending_acks: self.connected, - retry_resubscribe: self.connected && self.pending_resubscribe, - } - } +enum PolledMqttEvent { + Event(MqttEvent), + Disconnect, +} +impl ProtocolState { const fn on_connack( &mut self, acknowledgements: bool, @@ -79,7 +71,6 @@ impl ProtocolState { ) -> ConnAckActions { let actions = ConnAckActions { warn_session_not_resumed: acknowledgements && !session_present, - clear_pending_acks: true, flush_finalizer: true, resubscribe: self.connection_generation > 0 && !session_present, }; @@ -94,17 +85,14 @@ impl ProtocolState { self.connected = false; DisconnectActions { - clear_pending_acks: true, flush_finalizer: true, } } - const fn on_resubscribe_result(&mut self, success: bool) { - self.pending_resubscribe = !success; - } - fn should_ack_finalized_publish(&self, status: BatchStatus, entry_generation: u64) -> bool { - status == BatchStatus::Delivered && entry_generation == self.connection_generation + self.connected + && status == BatchStatus::Delivered + && entry_generation == self.connection_generation } } @@ -137,51 +125,31 @@ fn warn_session_not_resumed() { ); } -fn warn_resubscribe_failed() { +fn warn_subscribe_failed() { warn!( - message = "Failed to queue MQTT re-subscribe request after reconnect; will retry while connected.", + message = "Failed to queue MQTT subscribe request.", internal_log_rate_limit = true, ); } -#[derive(Default)] -struct PendingAcks { - publishes: Vec, -} - -impl PendingAcks { - fn push(&mut self, publish: Publish) { - self.publishes.push(publish); - } - - fn clear(&mut self) { - self.publishes.clear(); - } - - fn retry(&mut self, client: &rumqttc::AsyncClient) { - self.retry_with(|publish| client.try_ack(publish).is_ok()); - } - - fn try_ack(&mut self, connected: bool, publish: Publish, client: &rumqttc::AsyncClient) { - self.try_ack_with(connected, publish, |publish| { - client.try_ack(publish).is_ok() - }); - } +async fn poll_mqtt_connection( + mut connection: EventLoop, + events_tx: mpsc::UnboundedSender, + shutdown: ShutdownSignal, +) { + loop { + let event = tokio::select! { + _ = shutdown.clone() => break, + event = connection.poll() => match event { + Ok(event) => PolledMqttEvent::Event(event), + Err(_) => PolledMqttEvent::Disconnect, + }, + }; - fn try_ack_with( - &mut self, - connected: bool, - publish: Publish, - mut try_ack: impl FnMut(&Publish) -> bool, - ) { - if connected && !try_ack(&publish) { - self.push(publish); + if events_tx.send(event).is_err() { + break; } } - - fn retry_with(&mut self, mut try_ack: impl FnMut(&Publish) -> bool) { - self.publishes.retain(|publish| !try_ack(publish)); - } } pub struct MqttSource { @@ -210,9 +178,15 @@ impl MqttSource { } pub async fn run(self, mut out: SourceSender, shutdown: ShutdownSignal) -> Result<(), ()> { - let (client, mut connection) = self.connector.connect(); + let (client, connection) = self.connector.connect(); + let (mqtt_events_tx, mut mqtt_events_rx) = mpsc::unbounded_channel(); + tokio::spawn(poll_mqtt_connection( + connection, + mqtt_events_tx, + shutdown.clone(), + )); - self.subscribe(&client)?; + self.subscribe(&client).await?; // Finalizer drives end-to-end acknowledgements: each in-flight publish is // registered with its batch-status receiver, and we send the QoS-1 PUBACK @@ -225,49 +199,29 @@ impl MqttSource { Some(shutdown.clone()), ); - // PUBACKs that rumqttc's bounded request channel was too full to accept, - // retained for retry rather than dropped. Dropping a PUBACK for an already - // delivered message would pin it in the broker's in-flight window until the - // next reconnect. This is bounded in practice by that in-flight window (the - // broker stops delivering once it fills), and the event loop below drains the - // request channel, so entries flush on subsequent iterations. let mut protocol_state = ProtocolState::default(); - let mut pending_acks = PendingAcks::default(); loop { - let actions = protocol_state.loop_actions(); - if actions.retry_resubscribe { - protocol_state.on_resubscribe_result(self.try_subscribe(&client)); - } - - // Retry deferred PUBACKs while connected (the event loop below drains the - // request channel). Skipped while disconnected: a publish's packet id is - // only valid on the connection it arrived on, so stale PUBACKs must not be - // replayed across a reconnect. - if actions.retry_pending_acks { - pending_acks.retry(&client); - } - tokio::select! { _ = shutdown.clone() => return Ok(()), entry = ack_stream.next() => { // Only PUBACK delivered events. On Errored/Rejected we skip the // ack so the broker redelivers after reconnect (QoS-1 + - // clean_session=false), giving at-least-once delivery. Use the - // non-blocking `try_ack` — awaiting `ack` could deadlock, since - // this same task polls the event loop that drains rumqttc's request - // channel. If that channel is full, retain the PUBACK for retry - // (above) instead of dropping it. + // clean_session=false), giving at-least-once delivery. The MQTT + // event loop is polled by a separate task, so awaiting `ack` here + // does not block the task responsible for draining rumqttc's + // request channel. Stale finalizers from prior connections are + // ignored because packet IDs are only valid on their connection. if let Some((status, entry)) = entry && protocol_state.should_ack_finalized_publish( status, entry.connection_generation, ) { - pending_acks.try_ack(protocol_state.connected, entry.publish, &client); + client.ack(&entry.publish).await.map_err(|_| ())?; } }, - mqtt_event = connection.poll() => { + mqtt_event = mqtt_events_rx.recv() => { // Providing at-least-once here does not require correlating a // connection/poll error back to a specific in-flight publish. // rumqtt#349 (no packet id for *outbound* publishes) concerns the @@ -278,7 +232,7 @@ impl MqttSource { // redelivered by the broker on reconnect (clean_session=false + QoS // AtLeastOnce). match mqtt_event { - Ok(MqttEvent::Incoming(Incoming::Publish(publish))) => { + Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(publish)))) => { self.process_message( publish, &mut out, @@ -286,7 +240,7 @@ impl MqttSource { protocol_state.connection_generation, ).await; } - Ok(MqttEvent::Incoming(Incoming::SubAck(suback))) + Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::SubAck(suback)))) if self.acknowledgements => { for return_code in suback.return_codes { @@ -304,7 +258,7 @@ impl MqttSource { // A (re)connected session resumes here; the broker will // redeliver any unacknowledged publishes, so drop deferred // PUBACKs whose packet ids came from the previous connection. - Ok(MqttEvent::Incoming(Incoming::ConnAck(connack))) => { + Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::ConnAck(connack)))) => { let actions = protocol_state.on_connack( self.acknowledgements, connack.session_present, @@ -312,31 +266,26 @@ impl MqttSource { if actions.warn_session_not_resumed { warn_session_not_resumed(); } - if actions.clear_pending_acks { - pending_acks.clear(); - } if actions.flush_finalizer && let Some(finalizer) = &finalizer { finalizer.flush(); } if actions.resubscribe { - protocol_state.on_resubscribe_result(self.try_subscribe(&client)); + self.subscribe(&client).await?; } } // Connection lost: same stale-packet-id reasoning, and rumqttc // drops its own queued acks while reconnecting. - Err(_) => { + Some(PolledMqttEvent::Disconnect) => { let actions = protocol_state.on_disconnect(); - if actions.clear_pending_acks { - pending_acks.clear(); - } if actions.flush_finalizer && let Some(finalizer) = &finalizer { finalizer.flush(); } } + None => return Ok(()), _ => {} } } @@ -344,30 +293,26 @@ impl MqttSource { } } - fn try_subscribe(&self, client: &rumqttc::AsyncClient) -> bool { - match self.subscribe(client) { - Ok(()) => true, - Err(()) => { - warn_resubscribe_failed(); - false + async fn subscribe(&self, client: &AsyncClient) -> Result<(), ()> { + let result = match &self.config.topic { + OneOrMany::One(topic) => client.subscribe(topic, SUBSCRIPTION_QOS).await, + OneOrMany::Many(topics) => { + client + .subscribe_many( + topics + .iter() + .cloned() + .map(|topic| SubscribeFilter::new(topic, SUBSCRIPTION_QOS)), + ) + .await } - } - } + }; - fn subscribe(&self, client: &rumqttc::AsyncClient) -> Result<(), ()> { - match &self.config.topic { - OneOrMany::One(topic) => client - .try_subscribe(topic, SUBSCRIPTION_QOS) - .map_err(|_| ()), - OneOrMany::Many(topics) => client - .try_subscribe_many( - topics - .iter() - .cloned() - .map(|topic| SubscribeFilter::new(topic, SUBSCRIPTION_QOS)), - ) - .map_err(|_| ()), + if result.is_err() { + warn_subscribe_failed(); } + + result.map_err(|_| ()) } async fn process_message( @@ -453,64 +398,6 @@ impl MqttSource { mod tests { use super::*; - fn publish(pkid: u16) -> Publish { - let mut publish = Publish::new("topic", QoS::AtLeastOnce, vec![1, 2, 3]); - publish.pkid = pkid; - publish - } - - #[test] - fn pending_acks_keeps_failed_retries() { - let mut pending_acks = PendingAcks::default(); - pending_acks.push(publish(1)); - pending_acks.push(publish(2)); - pending_acks.push(publish(3)); - - let mut attempted = Vec::new(); - pending_acks.retry_with(|publish| { - attempted.push(publish.pkid); - publish.pkid != 2 - }); - - assert_eq!(attempted, vec![1, 2, 3]); - assert_eq!(pending_acks.publishes.len(), 1); - assert_eq!(pending_acks.publishes[0].pkid, 2); - - pending_acks.retry_with(|_| true); - assert!(pending_acks.publishes.is_empty()); - } - - #[test] - fn pending_acks_clear_drops_stale_packet_ids() { - let mut pending_acks = PendingAcks::default(); - pending_acks.push(publish(1)); - pending_acks.push(publish(2)); - - pending_acks.clear(); - - assert!(pending_acks.publishes.is_empty()); - } - - #[test] - fn pending_acks_backpressure_matrix() { - for (connected, try_ack_succeeds, expected_attempted, expected_queued) in [ - (false, true, false, false), - (true, true, true, false), - (true, false, true, true), - ] { - let mut pending_acks = PendingAcks::default(); - let mut attempted = false; - - pending_acks.try_ack_with(connected, publish(1), |_| { - attempted = true; - try_ack_succeeds - }); - - assert_eq!(attempted, expected_attempted); - assert_eq!(!pending_acks.publishes.is_empty(), expected_queued); - } - } - #[test] fn protocol_contract_matrix_for_requested_and_granted_qos() { assert!(publish_supports_end_to_end_acknowledgements( @@ -591,7 +478,6 @@ mod tests { let actions = state.on_connack(acknowledgements, session_present); assert_eq!(actions.warn_session_not_resumed, expected_warn); - assert!(actions.clear_pending_acks); assert!(actions.flush_finalizer); assert!(!actions.resubscribe); assert!(state.connected); @@ -605,7 +491,6 @@ mod tests { actions, ConnAckActions { warn_session_not_resumed: false, - clear_pending_acks: true, flush_finalizer: true, resubscribe: false, } @@ -619,7 +504,6 @@ mod tests { actions, ConnAckActions { warn_session_not_resumed: true, - clear_pending_acks: true, flush_finalizer: true, resubscribe: true, } @@ -627,37 +511,6 @@ mod tests { assert_eq!(fresh_session.connection_generation, 2); } - #[test] - fn protocol_contract_matrix_for_pending_resubscribe() { - let mut state = ProtocolState::default(); - state.on_resubscribe_result(false); - assert_eq!( - state.loop_actions(), - LoopActions { - retry_pending_acks: false, - retry_resubscribe: false, - } - ); - - state.on_connack(true, true); - assert_eq!( - state.loop_actions(), - LoopActions { - retry_pending_acks: true, - retry_resubscribe: true, - } - ); - - state.on_resubscribe_result(true); - assert_eq!( - state.loop_actions(), - LoopActions { - retry_pending_acks: true, - retry_resubscribe: false, - } - ); - } - #[test] fn protocol_contract_matrix_for_disconnect() { let mut state = ProtocolState::default(); @@ -668,19 +521,11 @@ mod tests { assert_eq!( actions, DisconnectActions { - clear_pending_acks: true, flush_finalizer: true, } ); assert!(!state.connected); assert_eq!(state.connection_generation, 1); - assert_eq!( - state.loop_actions(), - LoopActions { - retry_pending_acks: false, - retry_resubscribe: false, - } - ); } #[test] @@ -700,5 +545,8 @@ mod tests { should_ack ); } + + state.on_disconnect(); + assert!(!state.should_ack_finalized_publish(BatchStatus::Delivered, 2)); } } From e18dc1000a3a148e27185a9a9e9c2bc2036c84e8 Mon Sep 17 00:00:00 2001 From: colin Date: Thu, 2 Jul 2026 15:23:05 -0700 Subject: [PATCH 4/5] enhancement(mqtt source): bound the poller-to-main event backlog to avoid unbounded growth The poller task forwarded events to the main task over an unbounded channel, so a stalled downstream sink let the backlog grow without limit. Switch to a bounded channel and drop events on saturation instead: a dropped QoS 1/2 publish was never acknowledged, so the broker's own retry timer redelivers it independently of any reconnect. Forcing a reconnect on saturation was tried first but causes a reconnect storm, since the broker immediately refloods the freshly reconnected client from the same backlog. Also raises the mqtt integration test broker's max_mqueue_len/max_inflight and shortens retry_interval so the new saturation-recovery test can actually reach and exercise this path against a real broker. --- src/sources/mqtt/integration_tests.rs | 104 +++++++- src/sources/mqtt/source.rs | 266 +++++++++++++++------ tests/integration/mqtt/config/compose.yaml | 14 ++ 3 files changed, 313 insertions(+), 71 deletions(-) diff --git a/src/sources/mqtt/integration_tests.rs b/src/sources/mqtt/integration_tests.rs index 45662af4c77ab..70f86b4632787 100644 --- a/src/sources/mqtt/integration_tests.rs +++ b/src/sources/mqtt/integration_tests.rs @@ -49,7 +49,7 @@ fn message_body(event: &Event) -> String { .into_owned() } -async fn get_mqtt_client() -> AsyncClient { +async fn get_mqtt_client_with_options(configure: impl FnOnce(&mut MqttOptions)) -> AsyncClient { // Unique client ID per producer: brokers that strictly enforce client-ID // uniqueness (e.g. RabbitMQ) otherwise kick a previous connection when tests // run concurrently, which manifests as spurious publish timeouts. @@ -59,6 +59,7 @@ async fn get_mqtt_client() -> AsyncClient { mqtt_broker_port(), ); mqtt_options.set_keep_alive(Duration::from_secs(5)); + configure(&mut mqtt_options); let (client, mut eventloop) = AsyncClient::new(mqtt_options, 10); @@ -71,6 +72,10 @@ async fn get_mqtt_client() -> AsyncClient { client } +async fn get_mqtt_client() -> AsyncClient { + get_mqtt_client_with_options(|_| {}).await +} + #[tokio::test] async fn mqtt_one_topic_happy() { trace_init(); @@ -223,6 +228,103 @@ async fn mqtt_redelivers_unacknowledged_messages() { drop(source2.await); } +/// A downstream that can't keep up must not deadlock the source or grow memory +/// without bound. `poll_mqtt_connection` forwards polled events to the main task +/// over a channel bounded by `EVENTS_CHANNEL_CAPACITY`; once that backlog +/// saturates, the event is dropped rather than buffered without bound or blocking +/// (blocking would risk deadlocking against rumqttc's own outgoing request +/// channel, since the poller is also the only thing that drains it). This test +/// proves that path actually recovers instead of hanging, and that every message +/// is still eventually delivered afterward: a dropped QoS 1 publish was never +/// acknowledged, so the broker's own retry timer redelivers it independently of +/// the connection ever reconnecting. +#[tokio::test] +async fn mqtt_recovers_from_downstream_saturation() { + trace_init(); + + let topic = format!("source-saturation-test-{}", random_string(6)); + let client_id = format!("sourceSaturationTest{}", random_string(6)); + // Comfortably more than the poller's internal event backlog capacity, so the + // backlog saturates and starts dropping messages before the sink below starts + // draining. + let num_events = super::source::EVENTS_CHANNEL_CAPACITY * 3; + let (input, ..) = random_lines_with_stream(16, num_events, None); + + let config = MqttSourceConfig { + common: MqttCommonConfig { + host: mqtt_broker_address(), + port: mqtt_broker_port(), + client_id: Some(client_id), + ..Default::default() + }, + topic: OneOrMany::One(topic.clone()), + acknowledgements: true.into(), + ..MqttSourceConfig::default() + }; + + // A small, undrained buffer stands in for a stalled/slow sink: `out.send_batch` + // inside the source blocks almost immediately once it fills. + let (tx, rx) = SourceSender::new_test_finalize(EventStatus::Delivered); + let source = tokio::spawn(async move { + config + .build(SourceContext::new_test(tx, None)) + .await + .unwrap() + .await + .unwrap() + }); + + tokio::time::sleep(Duration::from_millis(500)).await; + + // rumqttc's default outgoing QoS-1 inflight window (100) would otherwise pace + // the producer to the broker's PUBACK round-trip time, making it too slow to + // ever build up a backlog at the subscriber. Raise it so publishing is only + // limited by the network, not by our own producer's flow control. + let producer = get_mqtt_client_with_options(|opts| { + opts.set_inflight(u16::MAX); + }) + .await; + send_test_events(&producer, &topic, &input).await; + + // Give the poller time to saturate its backlog and start dropping messages + // while nothing is draining the sink. + tokio::time::sleep(Duration::from_secs(2)).await; + + // Now let the sink drain. Duplicates are expected (a dropped message is + // redelivered by the broker's retry timer, independently of whatever else is + // still flowing through), so drain until every distinct message has been seen + // rather than taking a fixed count. If the poller had deadlocked instead of + // recovering, this times out. + let mut expected_messages: HashSet<_> = input.into_iter().collect(); + let mut total_received = 0usize; + let drain = async { + tokio::pin!(rx); + while !expected_messages.is_empty() { + let event = rx.next().await.expect("source stream ended unexpectedly"); + expected_messages.remove(&message_body(&event)); + total_received += 1; + } + }; + timeout(Duration::from_secs(60), drain).await.expect( + "timed out waiting for all messages after downstream recovered: the source likely deadlocked", + ); + + // If nothing was ever actually dropped (e.g. a regression reverted the + // backlog channel to unbounded), every message would still eventually + // arrive via simple buffering and the check above would pass for the wrong + // reason. Requiring more received events than distinct inputs pins the + // pass on the drop-and-redeliver path having actually run. + assert!( + total_received > num_events, + "expected at least one duplicate from broker redelivery of a dropped message, \ + got {total_received} events for {num_events} distinct inputs -- the backlog \ + may never have actually saturated" + ); + + source.abort(); + drop(source.await); +} + #[tokio::test] async fn mqtt_many_topics_happy() { trace_init(); diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index d94e2d703e4ea..69fb335f0ebee 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1,3 +1,5 @@ +use std::ops::ControlFlow; + use futures::StreamExt; use itertools::Itertools; use rumqttc::{ @@ -24,6 +26,17 @@ use crate::{ const SUBSCRIPTION_QOS: QoS = QoS::AtLeastOnce; +/// Bound on the number of polled MQTT events buffered between the poller task and +/// the main task, for events it's safe to drop under backlog (see +/// `safe_to_drop_under_backlog`). Matches the capacity of rumqttc's own +/// outgoing-request channel (see `MqttConnector::connect`) as a reasonable +/// starting point; the two aren't required to match, but there's no reason to +/// buffer more here than rumqttc itself buffers for outgoing acks/subscribes. +/// +/// `pub(super)` so the integration test that exercises the saturation path can +/// publish comfortably more than this many messages without hard-coding the value. +pub(super) const EVENTS_CHANNEL_CAPACITY: usize = 1024; + /// Identifies an in-flight publish so its QoS-1 PUBACK can be sent once the /// downstream sinks confirm delivery. Only the packet id (carried by `Publish`) /// is needed to ack; the payload is cleared before the entry is retained so @@ -132,9 +145,56 @@ fn warn_subscribe_failed() { ); } +fn warn_event_backlog_saturated() { + warn!( + message = "MQTT event backlog is saturated because downstream is not keeping up; dropping this message. It was not yet acknowledged, so the broker's own retry timer will redeliver it independently of any reconnect.", + internal_log_rate_limit = true, + ); +} + +/// Whether dropping this specific polled event under a saturated backlog is +/// protocol-safe. Only a QoS 1/2 publish that will actually be deferred-acked +/// (acknowledgements enabled, see `publish_ack_decision`) is safe to drop here: +/// only then does the broker still consider it unacknowledged and worth +/// retrying on its own. Everything else must never be dropped this way: +/// connection lifecycle events (ConnAck/SubAck/Disconnect) drive `ProtocolState` +/// and losing one desyncs it from the real connection (e.g. permanently skipping +/// a post-reconnect resubscribe); a publish rumqttc has already auto-acked +/// (acknowledgements disabled) or that has no acknowledgement at all (QoS 0) has +/// no redelivery guarantee, so dropping it would be true, unrecoverable data +/// loss rather than a redelivered retry. +fn safe_to_drop_under_backlog(event: &PolledMqttEvent, acknowledgements: bool) -> bool { + match event { + PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(publish))) => { + acknowledgements && publish_supports_end_to_end_acknowledgements(publish.qos) + } + _ => false, + } +} + +// Polls the MQTT event loop on a dedicated task so the main task (see `run`) can +// use the blocking `client.ack`/`client.subscribe` without risking a deadlock: +// `EventLoop::poll` is the only thing that drains rumqttc's outgoing request +// channel (used by those calls), so this task must never block on anything other +// than `poll` itself. +// +// Polled events are split across two channels by `safe_to_drop_under_backlog`: +// - Events safe to drop go through a bounded channel with a non-blocking +// `try_send`; if the main task falls behind and it fills, the event is +// dropped (with a warning) rather than buffered without bound or blocking. +// An earlier version of this forced a reconnect on saturation instead, but +// that reintroduces the underlying problem immediately: the broker refloods +// a freshly (re)connected client from the same backlog, so under sustained +// backpressure it degenerates into a reconnect storm instead of actually +// recovering. +// - Everything else is forwarded over an unbounded channel and never dropped. +// `send` on an unbounded channel never blocks, so this doesn't reintroduce +// the deadlock this task is designed to avoid. async fn poll_mqtt_connection( mut connection: EventLoop, - events_tx: mpsc::UnboundedSender, + acknowledgements: bool, + droppable_tx: mpsc::Sender, + guaranteed_tx: mpsc::UnboundedSender, shutdown: ShutdownSignal, ) { loop { @@ -146,7 +206,13 @@ async fn poll_mqtt_connection( }, }; - if events_tx.send(event).is_err() { + if safe_to_drop_under_backlog(&event, acknowledgements) { + match droppable_tx.try_send(event) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => warn_event_backlog_saturated(), + Err(mpsc::error::TrySendError::Closed(_)) => break, + } + } else if guaranteed_tx.send(event).is_err() { break; } } @@ -179,10 +245,13 @@ impl MqttSource { pub async fn run(self, mut out: SourceSender, shutdown: ShutdownSignal) -> Result<(), ()> { let (client, connection) = self.connector.connect(); - let (mqtt_events_tx, mut mqtt_events_rx) = mpsc::unbounded_channel(); + let (droppable_tx, mut droppable_rx) = mpsc::channel(EVENTS_CHANNEL_CAPACITY); + let (guaranteed_tx, mut guaranteed_rx) = mpsc::unbounded_channel(); tokio::spawn(poll_mqtt_connection( connection, - mqtt_events_tx, + self.acknowledgements, + droppable_tx, + guaranteed_tx, shutdown.clone(), )); @@ -221,76 +290,102 @@ impl MqttSource { client.ack(&entry.publish).await.map_err(|_| ())?; } }, - mqtt_event = mqtt_events_rx.recv() => { - // Providing at-least-once here does not require correlating a - // connection/poll error back to a specific in-flight publish. - // rumqtt#349 (no packet id for *outbound* publishes) concerns the - // publish/sink direction and does not apply to a subscribe-only - // source: each incoming Publish already carries its packet id, and - // we withhold its QoS-1 PUBACK until the event is delivered - // end-to-end. Anything left unacked when the connection drops is - // redelivered by the broker on reconnect (clean_session=false + QoS - // AtLeastOnce). - match mqtt_event { - Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(publish)))) => { - self.process_message( - publish, - &mut out, - finalizer.as_ref(), - protocol_state.connection_generation, - ).await; - } - Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::SubAck(suback)))) - if self.acknowledgements => - { - for return_code in suback.return_codes { - if let rumqttc::SubscribeReasonCode::Success(qos) = return_code - && !publish_supports_end_to_end_acknowledgements(qos) - { - warn!( - message = "MQTT broker granted a subscription QoS below the level required for end-to-end acknowledgements.", - ?qos, - internal_log_rate_limit = true, - ); - } - } - } - // A (re)connected session resumes here; the broker will - // redeliver any unacknowledged publishes, so drop deferred - // PUBACKs whose packet ids came from the previous connection. - Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::ConnAck(connack)))) => { - let actions = protocol_state.on_connack( - self.acknowledgements, - connack.session_present, - ); - if actions.warn_session_not_resumed { - warn_session_not_resumed(); - } - if actions.flush_finalizer - && let Some(finalizer) = &finalizer - { - finalizer.flush(); - } - if actions.resubscribe { - self.subscribe(&client).await?; - } - } - // Connection lost: same stale-packet-id reasoning, and rumqttc - // drops its own queued acks while reconnecting. - Some(PolledMqttEvent::Disconnect) => { - let actions = protocol_state.on_disconnect(); - if actions.flush_finalizer - && let Some(finalizer) = &finalizer - { - finalizer.flush(); - } - } - None => return Ok(()), - _ => {} + mqtt_event = droppable_rx.recv() => { + if let ControlFlow::Break(result) = self.handle_mqtt_event( + mqtt_event, &mut out, &finalizer, &mut protocol_state, &client, + ).await { + return result; + } + }, + mqtt_event = guaranteed_rx.recv() => { + if let ControlFlow::Break(result) = self.handle_mqtt_event( + mqtt_event, &mut out, &finalizer, &mut protocol_state, &client, + ).await { + return result; + } + }, + } + } + } + + // Providing at-least-once here does not require correlating a + // connection/poll error back to a specific in-flight publish. rumqtt#349 (no + // packet id for *outbound* publishes) concerns the publish/sink direction + // and does not apply to a subscribe-only source: each incoming Publish + // already carries its packet id, and we withhold its QoS-1 PUBACK until the + // event is delivered end-to-end. Anything left unacked when the connection + // drops is redelivered by the broker on reconnect (clean_session=false + QoS + // AtLeastOnce). + // + // Returns `ControlFlow::Break` with the result `run` should return, or + // `ControlFlow::Continue` to keep looping. + async fn handle_mqtt_event( + &self, + mqtt_event: Option, + out: &mut SourceSender, + finalizer: &Option>, + protocol_state: &mut ProtocolState, + client: &AsyncClient, + ) -> ControlFlow> { + match mqtt_event { + Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(publish)))) => { + self.process_message( + publish, + out, + finalizer.as_ref(), + protocol_state.connection_generation, + ) + .await; + } + Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::SubAck(suback)))) + if self.acknowledgements => + { + for return_code in suback.return_codes { + if let rumqttc::SubscribeReasonCode::Success(qos) = return_code + && !publish_supports_end_to_end_acknowledgements(qos) + { + warn!( + message = "MQTT broker granted a subscription QoS below the level required for end-to-end acknowledgements.", + ?qos, + internal_log_rate_limit = true, + ); } } } + // A (re)connected session resumes here; the broker will redeliver + // any unacknowledged publishes, so drop deferred PUBACKs whose + // packet ids came from the previous connection. + Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::ConnAck(connack)))) => { + let actions = + protocol_state.on_connack(self.acknowledgements, connack.session_present); + if actions.warn_session_not_resumed { + warn_session_not_resumed(); + } + if actions.flush_finalizer + && let Some(finalizer) = finalizer + { + finalizer.flush(); + } + if actions.resubscribe + && let Err(()) = self.subscribe(client).await + { + return ControlFlow::Break(Err(())); + } + } + // Connection lost: same stale-packet-id reasoning, and rumqttc drops + // its own queued acks while reconnecting. + Some(PolledMqttEvent::Disconnect) => { + let actions = protocol_state.on_disconnect(); + if actions.flush_finalizer + && let Some(finalizer) = finalizer + { + finalizer.flush(); + } + } + None => return ControlFlow::Break(Ok(())), + _ => {} } + ControlFlow::Continue(()) } async fn subscribe(&self, client: &AsyncClient) -> Result<(), ()> { @@ -549,4 +644,35 @@ mod tests { state.on_disconnect(); assert!(!state.should_ack_finalized_publish(BatchStatus::Delivered, 2)); } + + #[test] + fn protocol_contract_matrix_for_backlog_drop_safety() { + fn publish_event(qos: QoS) -> PolledMqttEvent { + PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(Publish::new( + "topic", + qos, + vec![1, 2, 3], + )))) + } + + for (acknowledgements, qos, safe_to_drop) in [ + (false, QoS::AtMostOnce, false), + (false, QoS::AtLeastOnce, false), + (false, QoS::ExactlyOnce, false), + (true, QoS::AtMostOnce, false), + (true, QoS::AtLeastOnce, true), + (true, QoS::ExactlyOnce, true), + ] { + assert_eq!( + safe_to_drop_under_backlog(&publish_event(qos), acknowledgements), + safe_to_drop, + "acknowledgements={acknowledgements}, qos={qos:?}" + ); + } + + assert!(!safe_to_drop_under_backlog( + &PolledMqttEvent::Disconnect, + true + )); + } } diff --git a/tests/integration/mqtt/config/compose.yaml b/tests/integration/mqtt/config/compose.yaml index 9ce016da98cb1..3fb82e22d1b29 100644 --- a/tests/integration/mqtt/config/compose.yaml +++ b/tests/integration/mqtt/config/compose.yaml @@ -5,6 +5,20 @@ services: image: docker.io/emqx:${CONFIG_VERSION} ports: - 1883:1883 + environment: + # Default is 1000, too small to exercise the source's own + # backlog/backpressure handling (which buffers well beyond that) without the + # broker's own per-client queue dropping messages first. + - EMQX_MQTT__MAX_MQUEUE_LEN=100000 + # Default is 32: the broker won't have more than this many QoS 1/2 messages + # unacknowledged in flight to a client at once, which is far too small to + # ever reach the source's own backlog capacity regardless of test volume. + - EMQX_MQTT__MAX_INFLIGHT=10000 + # Default is 30s: how long the broker waits for a PUBACK before retrying a + # QoS 1/2 publish. Shortened so a test that deliberately drops messages to + # saturate the source's backlog doesn't have to wait 30s per retry round to + # observe redelivery. + - EMQX_MQTT__RETRY_INTERVAL=2s networks: default: From 295993ef7ed28cb0b6b9593f8dd1b0841d50a229 Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 8 Jul 2026 14:58:56 -0700 Subject: [PATCH 5/5] fix(mqtt source): preserve event order and PUBACK ordering, drop the reconnect escalation Addressed PR review findings: - Merge the droppable/guaranteed event channels back into one, preserving the order events were polled in. Splitting into two channels (from the previous commit) broke this: tokio::select! doesn't guarantee relative order between two channels when both have items ready, so a backlogged publish could reach the main task before a ConnAck that was actually polled first, letting ProtocolState be acted on while stale. - Switch UnorderedFinalizer to OrderedFinalizer: PUBACKs must be sent in the order their publishes were received (MQTT-4.6.0-2), same as the kafka source's ordered offset commits. - Drop `biased` on the main select: it fixed a narrow stale-ack race but starves the ack-stream branch entirely under sustained publish volume, since the events channel being ready again next iteration is the common case, not the exception. Fair random selection keeps that race narrow instead of guaranteeing worse behavior under load. Also attempted, then reverted, a fix for the remaining gap (dropped messages rely on broker-specific retry-while-connected behavior, which isn't spec- guaranteed): forcing a reconnect once the backlog stays saturated long enough. Verified against a live EMQX broker that this triggers a broker-side `unexpected_sock_close`/`einval` error after which the session stops delivering anything further at all -- worse than the problem it was meant to solve. Reverted to plain dropping and documented the limitation in code, since forcing a *correct* reconnect needs more investigation than fits here. --- src/sources/mqtt/integration_tests.rs | 5 +- src/sources/mqtt/source.rs | 160 +++++++++++++++++--------- 2 files changed, 106 insertions(+), 59 deletions(-) diff --git a/src/sources/mqtt/integration_tests.rs b/src/sources/mqtt/integration_tests.rs index 70f86b4632787..13362b8d339f9 100644 --- a/src/sources/mqtt/integration_tests.rs +++ b/src/sources/mqtt/integration_tests.rs @@ -236,8 +236,9 @@ async fn mqtt_redelivers_unacknowledged_messages() { /// channel, since the poller is also the only thing that drains it). This test /// proves that path actually recovers instead of hanging, and that every message /// is still eventually delivered afterward: a dropped QoS 1 publish was never -/// acknowledged, so the broker's own retry timer redelivers it independently of -/// the connection ever reconnecting. +/// acknowledged, so the broker's own retry timer redelivers it (see +/// `compose.yaml`'s shortened `retry_interval`) independently of whatever else +/// is still flowing through. #[tokio::test] async fn mqtt_recovers_from_downstream_saturation() { trace_init(); diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index 69fb335f0ebee..faa887ea2891f 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1,4 +1,10 @@ -use std::ops::ControlFlow; +use std::{ + ops::ControlFlow, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; use futures::StreamExt; use itertools::Itertools; @@ -9,7 +15,7 @@ use tokio::sync::mpsc; use vector_lib::{ codecs::Decoder, config::{LegacyKey, LogNamespace}, - finalizer::UnorderedFinalizer, + finalizer::OrderedFinalizer, internal_event::EventsReceived, lookup::path, }; @@ -147,22 +153,23 @@ fn warn_subscribe_failed() { fn warn_event_backlog_saturated() { warn!( - message = "MQTT event backlog is saturated because downstream is not keeping up; dropping this message. It was not yet acknowledged, so the broker's own retry timer will redeliver it independently of any reconnect.", + message = "MQTT event backlog is saturated because downstream is not keeping up; dropping this message. It was not yet acknowledged, so the broker still considers it unacknowledged; whether and when it's redelivered depends on the broker (some retry independently of a reconnect, others only redeliver after this connection reconnects for some other reason).", internal_log_rate_limit = true, ); } /// Whether dropping this specific polled event under a saturated backlog is /// protocol-safe. Only a QoS 1/2 publish that will actually be deferred-acked -/// (acknowledgements enabled, see `publish_ack_decision`) is safe to drop here: -/// only then does the broker still consider it unacknowledged and worth -/// retrying on its own. Everything else must never be dropped this way: -/// connection lifecycle events (ConnAck/SubAck/Disconnect) drive `ProtocolState` -/// and losing one desyncs it from the real connection (e.g. permanently skipping -/// a post-reconnect resubscribe); a publish rumqttc has already auto-acked -/// (acknowledgements disabled) or that has no acknowledgement at all (QoS 0) has -/// no redelivery guarantee, so dropping it would be true, unrecoverable data -/// loss rather than a redelivered retry. +/// (acknowledgements enabled, see `publish_ack_decision`) is safe to drop +/// here: only then does the broker still consider it unacknowledged, and +/// therefore a candidate for redelivery (see `warn_event_backlog_saturated` +/// for the caveat on when that actually happens). Everything else must never +/// be dropped this way: connection lifecycle events (ConnAck/SubAck/ +/// Disconnect) drive `ProtocolState` and losing one desyncs it from the real +/// connection (e.g. permanently skipping a post-reconnect resubscribe); a +/// publish rumqttc has already auto-acked (acknowledgements disabled) or that +/// has no acknowledgement at all (QoS 0) has no redelivery guarantee at all, +/// so dropping it would be true, unrecoverable data loss. fn safe_to_drop_under_backlog(event: &PolledMqttEvent, acknowledgements: bool) -> bool { match event { PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(publish))) => { @@ -175,26 +182,44 @@ fn safe_to_drop_under_backlog(event: &PolledMqttEvent, acknowledgements: bool) - // Polls the MQTT event loop on a dedicated task so the main task (see `run`) can // use the blocking `client.ack`/`client.subscribe` without risking a deadlock: // `EventLoop::poll` is the only thing that drains rumqttc's outgoing request -// channel (used by those calls), so this task must never block on anything other -// than `poll` itself. +// channel (used by those calls), so this task must never block on anything +// other than `poll` itself. +// +// Polled events are forwarded over a single unbounded channel, preserving the +// order they were polled in: a ConnAck/Disconnect must be observed by the main +// task before any publish that arrived after it, or `ProtocolState` can be +// acted on (or acked against) while stale. An earlier version of this used two +// separate channels split by drop-safety, but `tokio::select!` doesn't +// preserve relative order between two channels when both have items ready, so +// a backlogged publish could reach the main task before a ConnAck that was +// actually polled first. +// +// Memory is still bounded despite the channel being unbounded: events safe to +// drop under backlog (see `safe_to_drop_under_backlog`) are only sent while +// `droppable_backlog`'s count of currently-outstanding droppable events is +// under `EVENTS_CHANNEL_CAPACITY`; once it saturates, further ones are dropped +// (with a warning) instead of sent. `send` on an unbounded channel never +// blocks, so none of this risks the deadlock this task is designed to avoid. +// Everything else (connection lifecycle events, and any publish that isn't +// safe to drop) is sent unconditionally and never dropped. // -// Polled events are split across two channels by `safe_to_drop_under_backlog`: -// - Events safe to drop go through a bounded channel with a non-blocking -// `try_send`; if the main task falls behind and it fills, the event is -// dropped (with a warning) rather than buffered without bound or blocking. -// An earlier version of this forced a reconnect on saturation instead, but -// that reintroduces the underlying problem immediately: the broker refloods -// a freshly (re)connected client from the same backlog, so under sustained -// backpressure it degenerates into a reconnect storm instead of actually -// recovering. -// - Everything else is forwarded over an unbounded channel and never dropped. -// `send` on an unbounded channel never blocks, so this doesn't reintroduce -// the deadlock this task is designed to avoid. +// This deliberately does not try to force a reconnect when the backlog stays +// saturated: an earlier version did (via `EventLoop::clean`), reasoning that +// the MQTT spec ties QoS >0 redelivery to reconnection rather than guaranteeing +// an independent retry-while-connected timer. In practice that forced, +// non-graceful reconnect triggered broker-specific error handling (observed +// against EMQX as an `unexpected_sock_close`/`einval` server-side error) after +// which the session stopped delivering anything further at all — worse than +// the problem it was meant to solve. Forcing a *correct* reconnect (e.g. a +// graceful `client.disconnect()` sequenced so the broker actually receives it +// before the socket closes) needs more investigation than fits here, so for +// now this only drops and relies on whatever the broker/connection naturally +// provides, same as before this file started attempting to force anything. async fn poll_mqtt_connection( mut connection: EventLoop, acknowledgements: bool, - droppable_tx: mpsc::Sender, - guaranteed_tx: mpsc::UnboundedSender, + events_tx: mpsc::UnboundedSender, + droppable_backlog: Arc, shutdown: ShutdownSignal, ) { loop { @@ -207,12 +232,14 @@ async fn poll_mqtt_connection( }; if safe_to_drop_under_backlog(&event, acknowledgements) { - match droppable_tx.try_send(event) { - Ok(()) => {} - Err(mpsc::error::TrySendError::Full(_)) => warn_event_backlog_saturated(), - Err(mpsc::error::TrySendError::Closed(_)) => break, + if droppable_backlog.load(Ordering::Acquire) >= EVENTS_CHANNEL_CAPACITY { + warn_event_backlog_saturated(); + continue; } - } else if guaranteed_tx.send(event).is_err() { + droppable_backlog.fetch_add(1, Ordering::AcqRel); + } + + if events_tx.send(event).is_err() { break; } } @@ -245,13 +272,13 @@ impl MqttSource { pub async fn run(self, mut out: SourceSender, shutdown: ShutdownSignal) -> Result<(), ()> { let (client, connection) = self.connector.connect(); - let (droppable_tx, mut droppable_rx) = mpsc::channel(EVENTS_CHANNEL_CAPACITY); - let (guaranteed_tx, mut guaranteed_rx) = mpsc::unbounded_channel(); + let (events_tx, mut events_rx) = mpsc::unbounded_channel(); + let droppable_backlog = Arc::new(AtomicUsize::new(0)); tokio::spawn(poll_mqtt_connection( connection, self.acknowledgements, - droppable_tx, - guaranteed_tx, + events_tx, + Arc::clone(&droppable_backlog), shutdown.clone(), )); @@ -260,10 +287,11 @@ impl MqttSource { // Finalizer drives end-to-end acknowledgements: each in-flight publish is // registered with its batch-status receiver, and we send the QoS-1 PUBACK // only once the sinks report `Delivered`. Unused when acknowledgements are - // disabled (rumqttc auto-acks in that mode). MQTT PUBACKs are independent - // per packet id (unlike Kafka offsets), so finalization is unordered — a - // slow/stuck batch must not hold back acks for already-delivered publishes. - let (finalizer, mut ack_stream) = UnorderedFinalizer::::maybe_new( + // disabled (rumqttc auto-acks in that mode). PUBACKs must be sent in the + // order their publishes were received ([MQTT-4.6.0-2]), so finalization is + // ordered: a slow/stuck earlier batch holds back acks for publishes + // received after it, same as the `kafka` source's ordered offset commits. + let (finalizer, mut ack_stream) = OrderedFinalizer::::maybe_new( self.acknowledgements, Some(shutdown.clone()), ); @@ -272,7 +300,28 @@ impl MqttSource { loop { tokio::select! { + // Deliberately not `biased`: a connection-lifecycle event + // already sitting in `events_rx` when an ack-stream + // completion is also ready should ideally be processed first + // (`should_ack_finalized_publish` needs `protocol_state` to + // reflect it, or it can act on a dead connection), but always + // favoring `events_rx` was tried and starves the ack-stream + // branch outright under sustained publish volume -- exactly + // the saturated-backlog scenario this file already has to + // handle -- since `events_rx` being ready again by the next + // loop iteration is the common case, not the exception, and + // `biased` would then never yield to the other branch. Fair + // random selection keeps that race narrow (both branches + // must be ready in the same instant) instead of guaranteeing + // it under exactly the load this needs to keep working under. _ = shutdown.clone() => return Ok(()), + mqtt_event = events_rx.recv() => { + if let ControlFlow::Break(result) = self.handle_mqtt_event( + mqtt_event, &mut out, &finalizer, &mut protocol_state, &client, &droppable_backlog, + ).await { + return result; + } + }, entry = ack_stream.next() => { // Only PUBACK delivered events. On Errored/Rejected we skip the // ack so the broker redelivers after reconnect (QoS-1 + @@ -290,20 +339,6 @@ impl MqttSource { client.ack(&entry.publish).await.map_err(|_| ())?; } }, - mqtt_event = droppable_rx.recv() => { - if let ControlFlow::Break(result) = self.handle_mqtt_event( - mqtt_event, &mut out, &finalizer, &mut protocol_state, &client, - ).await { - return result; - } - }, - mqtt_event = guaranteed_rx.recv() => { - if let ControlFlow::Break(result) = self.handle_mqtt_event( - mqtt_event, &mut out, &finalizer, &mut protocol_state, &client, - ).await { - return result; - } - }, } } } @@ -323,10 +358,21 @@ impl MqttSource { &self, mqtt_event: Option, out: &mut SourceSender, - finalizer: &Option>, + finalizer: &Option>, protocol_state: &mut ProtocolState, client: &AsyncClient, + droppable_backlog: &AtomicUsize, ) -> ControlFlow> { + // Mirrors the poller's own count of this event against + // `EVENTS_CHANNEL_CAPACITY` (see `poll_mqtt_connection`); recomputing + // the same predicate here, rather than tagging the event itself, keeps + // `PolledMqttEvent` free of backlog-accounting concerns. + if let Some(event) = &mqtt_event + && safe_to_drop_under_backlog(event, self.acknowledgements) + { + droppable_backlog.fetch_sub(1, Ordering::AcqRel); + } + match mqtt_event { Some(PolledMqttEvent::Event(MqttEvent::Incoming(Incoming::Publish(publish)))) => { self.process_message( @@ -414,7 +460,7 @@ impl MqttSource { &self, mut publish: Publish, out: &mut SourceSender, - finalizer: Option<&UnorderedFinalizer>, + finalizer: Option<&OrderedFinalizer>, connection_generation: u64, ) { emit!(EndpointBytesReceived {