From 4c3b4c5a2df9355f6ea2d07a0b09909db9722a83 Mon Sep 17 00:00:00 2001 From: colin Date: Fri, 19 Jun 2026 11:23:24 -0700 Subject: [PATCH 1/8] 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 878e246d434145beb7816e3252fa5451dd98f035 Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 1 Jul 2026 13:42:06 -0700 Subject: [PATCH 2/8] 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 69246f13cc0256e2154b97e41a779f4eddf8f5ae Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 8 Jul 2026 15:16:21 -0700 Subject: [PATCH 3/8] fix(mqtt source): send PUBACKs in the order publishes were received MQTT requires a client to send PUBACK (QoS 1) and PUBREC (QoS 2) packets in the order the corresponding PUBLISH packets were received (MQTT-4.6.0-2). UnorderedFinalizer lets a later batch's ack race ahead of an earlier one still pending delivery, violating this. Switch to OrderedFinalizer, same as the kafka source's ordered offset commits: a slow/stuck earlier batch now holds back acks for publishes received after it, trading a little head-of-line latency for protocol compliance. --- src/sources/mqtt/source.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index 616cb54914af4..88d296b9566f2 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -4,7 +4,7 @@ use rumqttc::{Event as MqttEvent, Incoming, Publish, QoS, SubscribeFilter}; use vector_lib::{ codecs::Decoder, config::{LegacyKey, LogNamespace}, - finalizer::UnorderedFinalizer, + finalizer::OrderedFinalizer, internal_event::EventsReceived, lookup::path, }; @@ -217,10 +217,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()), ); @@ -374,7 +375,7 @@ impl MqttSource { &self, mut publish: Publish, out: &mut SourceSender, - finalizer: Option<&UnorderedFinalizer>, + finalizer: Option<&OrderedFinalizer>, connection_generation: u64, ) { emit!(EndpointBytesReceived { From 630007fd2e1c5ec6acb6a8f7d70c1b16018ff8d5 Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 8 Jul 2026 16:15:59 -0700 Subject: [PATCH 4/8] fix(mqtt source): withhold later acks after a failure, recreate finalizer on reconnect Addressed PR review findings: - Once a publish in the current connection generation finalizes as anything other than Delivered, withhold acks for every later publish in that same generation too, not just the failed one. MQTT requires PUBACKs to be sent in the order their publishes were received (MQTT-4.6.0-2); acking a later packet while an earlier one was never acked would violate that order. The withheld ones are redelivered the same way the one that failed already is. - Replace finalizer.flush() with recreating the finalizer/ack-stream pair on (re)connect. flush() only clears entries already pulled into its internal ordered set; an entry added concurrently (via finalizer.add) that the finalizer's background task hadn't yet picked up survives the flush and gets pushed into the fresh set anyway. Since OrderedFinalizer won't yield anything newer until that stale entry resolves, it could hold back acks for the new generation. Dropping the old finalizer outright removes the race: its background task sees the channel closed and exits, taking any not-yet-picked-up entries with it. --- src/sources/mqtt/source.rs | 106 +++++++++++++++++++++++++++++-------- 1 file changed, 83 insertions(+), 23 deletions(-) diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index 88d296b9566f2..62a52b5fb1c57 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1,4 +1,4 @@ -use futures::StreamExt; +use futures::{StreamExt, stream::BoxStream}; use itertools::Itertools; use rumqttc::{Event as MqttEvent, Incoming, Publish, QoS, SubscribeFilter}; use vector_lib::{ @@ -36,6 +36,16 @@ struct ProtocolState { connected: bool, pending_resubscribe: bool, connection_generation: u64, + /// Set once a publish in the current generation finalizes as anything + /// other than `Delivered`. MQTT requires PUBACKs to be sent in the order + /// their publishes were received ([MQTT-4.6.0-2]); acking a later packet + /// while an earlier one in the same generation was never acked would + /// violate that order, so once set, no further acks are sent until the + /// next reconnect. The withheld ones are redelivered the same way the + /// packet that triggered this already is. + /// + /// [MQTT-4.6.0-2]: https://docs.oasis-open.org/mqtt/mqtt/v3.1.1/os/mqtt-v3.1.1-os.html + ack_suppressed: bool, } #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -48,14 +58,14 @@ struct LoopActions { struct ConnAckActions { warn_session_not_resumed: bool, clear_pending_acks: bool, - flush_finalizer: bool, + recreate_finalizer: bool, resubscribe: bool, } #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] struct DisconnectActions { clear_pending_acks: bool, - flush_finalizer: bool, + recreate_finalizer: bool, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -80,12 +90,13 @@ impl ProtocolState { let actions = ConnAckActions { warn_session_not_resumed: acknowledgements && !session_present, clear_pending_acks: true, - flush_finalizer: true, + recreate_finalizer: true, resubscribe: self.connection_generation > 0 && !session_present, }; self.connected = true; self.connection_generation += 1; + self.ack_suppressed = false; actions } @@ -95,7 +106,7 @@ impl ProtocolState { DisconnectActions { clear_pending_acks: true, - flush_finalizer: true, + recreate_finalizer: true, } } @@ -103,8 +114,21 @@ impl ProtocolState { self.pending_resubscribe = !success; } - fn should_ack_finalized_publish(&self, status: BatchStatus, entry_generation: u64) -> bool { - status == BatchStatus::Delivered && entry_generation == self.connection_generation + /// Whether this finalized publish should be acked. Also updates + /// `ack_suppressed`: see its doc comment for why a non-`Delivered` status + /// in the current generation withholds every ack after it, not just its + /// own. + fn should_ack_finalized_publish(&mut self, status: BatchStatus, entry_generation: u64) -> bool { + if entry_generation != self.connection_generation { + return false; + } + + if status != BatchStatus::Delivered { + self.ack_suppressed = true; + return false; + } + + !self.ack_suppressed } } @@ -221,10 +245,7 @@ impl MqttSource { // 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()), - ); + let (mut finalizer, mut ack_stream) = self.new_finalizer(&shutdown); // PUBACKs that rumqttc's bounded request channel was too full to accept, // retained for retry rather than dropped. Dropping a PUBACK for an already @@ -316,10 +337,8 @@ impl MqttSource { if actions.clear_pending_acks { pending_acks.clear(); } - if actions.flush_finalizer - && let Some(finalizer) = &finalizer - { - finalizer.flush(); + if actions.recreate_finalizer { + (finalizer, ack_stream) = self.new_finalizer(&shutdown); } if actions.resubscribe { protocol_state.on_resubscribe_result(self.try_subscribe(&client)); @@ -332,10 +351,8 @@ impl MqttSource { if actions.clear_pending_acks { pending_acks.clear(); } - if actions.flush_finalizer - && let Some(finalizer) = &finalizer - { - finalizer.flush(); + if actions.recreate_finalizer { + (finalizer, ack_stream) = self.new_finalizer(&shutdown); } } _ => {} @@ -345,6 +362,28 @@ impl MqttSource { } } + // Builds a fresh finalizer/ack-stream pair, discarding whatever the + // previous one held. Used instead of `FinalizerSet::flush` on + // (re)connect: `flush` only clears entries already pulled into its + // internal ordered set, not ones a concurrent `finalizer.add` call sent + // but that the finalizer's background task hadn't yet picked up: those + // survive the flush and get pushed into the set anyway. Because + // `OrderedFinalizer` won't yield anything newer until that stale entry + // resolves, it can hold back acks for the new connection generation even + // though `should_ack_finalized_publish` would correctly skip it once + // yielded. Dropping the old `finalizer` (and so its sender) instead lets + // its background task see the channel closed and exit outright, taking + // any not-yet-picked-up entries with it — no race window. + fn new_finalizer( + &self, + shutdown: &ShutdownSignal, + ) -> ( + Option>, + BoxStream<'static, (BatchStatus, FinalizerEntry)>, + ) { + OrderedFinalizer::::maybe_new(self.acknowledgements, Some(shutdown.clone())) + } + fn try_subscribe(&self, client: &rumqttc::AsyncClient) -> bool { match self.subscribe(client) { Ok(()) => true, @@ -593,7 +632,7 @@ mod tests { assert_eq!(actions.warn_session_not_resumed, expected_warn); assert!(actions.clear_pending_acks); - assert!(actions.flush_finalizer); + assert!(actions.recreate_finalizer); assert!(!actions.resubscribe); assert!(state.connected); assert_eq!(state.connection_generation, 1); @@ -607,7 +646,7 @@ mod tests { ConnAckActions { warn_session_not_resumed: false, clear_pending_acks: true, - flush_finalizer: true, + recreate_finalizer: true, resubscribe: false, } ); @@ -621,7 +660,7 @@ mod tests { ConnAckActions { warn_session_not_resumed: true, clear_pending_acks: true, - flush_finalizer: true, + recreate_finalizer: true, resubscribe: true, } ); @@ -670,7 +709,7 @@ mod tests { actions, DisconnectActions { clear_pending_acks: true, - flush_finalizer: true, + recreate_finalizer: true, } ); assert!(!state.connected); @@ -702,4 +741,25 @@ mod tests { ); } } + + // [MQTT-4.6.0-2] requires PUBACKs to be sent in the order their publishes + // were received: once an earlier packet in a generation goes unacked + // (Errored/Rejected), a later packet in that same generation must not be + // acked either, even if it was Delivered -- or the ack order would be + // violated. Only a reconnect (new generation) should resume acking. + #[test] + fn finalization_suppresses_later_acks_in_generation_after_a_failure() { + let mut state = ProtocolState::default(); + state.on_connack(true, true); + + assert!(state.should_ack_finalized_publish(BatchStatus::Delivered, 1)); + + assert!(!state.should_ack_finalized_publish(BatchStatus::Rejected, 1)); + assert!(!state.should_ack_finalized_publish(BatchStatus::Delivered, 1)); + assert!(!state.should_ack_finalized_publish(BatchStatus::Delivered, 1)); + + // A reconnect starts a fresh generation, resuming normal acking. + state.on_connack(true, false); + assert!(state.should_ack_finalized_publish(BatchStatus::Delivered, 2)); + } } From 5ecb5bb996e665f1a7033b8fa22c768242d3bbc7 Mon Sep 17 00:00:00 2001 From: colin Date: Thu, 9 Jul 2026 16:19:52 -0700 Subject: [PATCH 5/8] fix(mqtt source): drop stale acks, confirm resubscribes via SUBACK, reconnect on withheld ack Three fixes to reconnect/ack handling, all found by review on #25785: - On disconnect, rumqttc's own `EventLoop::clean()` moves any PUBACK/PUBREC that `try_ack` queued but hadn't reached the network yet into `connection.pending` for automatic replay on the next connection. Those reference packet ids from the connection that just died, so filter them out alongside clearing Vector's own `pending_acks`, or the broker could receive a stale ack for the wrong publish once reconnected. - A (re)subscribe was considered done as soon as `try_subscribe` queued the request locally, not once the broker actually confirmed it via SUBACK. If that queued request is lost (a disconnect racing the send), a later reconnect can report `session_present = true` for a session that never got the subscription, leaving the source connected but receiving nothing. `pending_resubscribe` now only clears on an actual SUBACK, and survives a `session_present = true` reconnect if the previous attempt was never confirmed. - Withholding a PUBACK for ordering correctness left the connection open indefinitely. Brokers that only redeliver in-flight QoS 1 messages on reconnect (not on an independent timer) would then just fill their in-flight window and stop delivering once withheld, with no way to recover. Force a reconnect the first time an ack is withheld in a generation, rate-limited (FORCED_RECONNECT_COOLDOWN) so a persistent downstream failure can't thrash the connection. Uses a graceful `try_disconnect()` rather than dropping the network directly: verified against a live EMQX broker that the latter triggers a broker-side `unexpected_sock_close` error that stops delivery entirely, whereas an explicit `Disconnect` is treated as a normal close. Verified against a live EMQX broker (including a new mqtt_forces_reconnect_after_withheld_ack integration test), rerun 4x stable. --- src/sources/mqtt/integration_tests.rs | 71 ++++++ src/sources/mqtt/source.rs | 333 ++++++++++++++++++++++---- 2 files changed, 363 insertions(+), 41 deletions(-) diff --git a/src/sources/mqtt/integration_tests.rs b/src/sources/mqtt/integration_tests.rs index 45662af4c77ab..0701fab7a128a 100644 --- a/src/sources/mqtt/integration_tests.rs +++ b/src/sources/mqtt/integration_tests.rs @@ -223,6 +223,77 @@ async fn mqtt_redelivers_unacknowledged_messages() { drop(source2.await); } +/// A withheld ack (the sink rejects the message, so its PUBACK is never sent) +/// must force a reconnect on its own -- without an external trigger like a +/// process restart or manual disconnect -- so the broker redelivers the +/// still-unacknowledged publish. Otherwise a broker that only redelivers +/// QoS 1 messages on reconnect (rather than on an independent timer) would +/// leave the source connected but stuck, never receiving it again. +#[tokio::test] +async fn mqtt_forces_reconnect_after_withheld_ack() { + trace_init(); + + let topic = "source-forced-reconnect-test"; + let client_id = format!("sourceForcedReconnectTest{}", random_string(6)); + let message = random_string(32); + + let config = MqttSourceConfig { + common: MqttCommonConfig { + host: mqtt_broker_address(), + port: mqtt_broker_port(), + client_id: Some(client_id), + ..Default::default() + }, + topic: OneOrMany::One(topic.to_owned()), + acknowledgements: true.into(), + ..MqttSourceConfig::default() + }; + + // The first delivery attempt for any publish is rejected (withholding its + // ack); every attempt after that succeeds. So the only way a second + // delivery of the same message shows up is a broker redelivery following + // a reconnect the source must have forced itself. + let (tx, mut rx) = SourceSender::new_test_errors(|count| count == 0); + 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; + + let producer = get_mqtt_client().await; + producer + .publish(topic, QoS::AtLeastOnce, false, message.as_bytes()) + .await + .unwrap(); + + let first = timeout(Duration::from_secs(5), rx.next()) + .await + .expect("timed out waiting for first delivery") + .expect("source stream ended unexpectedly"); + assert_eq!(message_body(&first), message); + + let redelivered = timeout(Duration::from_secs(10), rx.next()) + .await + .expect( + "timed out waiting for the forced reconnect to trigger redelivery: \ + the source never reconnected on its own", + ) + .expect("source stream ended unexpectedly"); + assert_eq!( + message_body(&redelivered), + message, + "redelivered message did not match the original" + ); + + 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 62a52b5fb1c57..c392bacc1e40d 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use futures::{StreamExt, stream::BoxStream}; use itertools::Itertools; use rumqttc::{Event as MqttEvent, Incoming, Publish, QoS, SubscribeFilter}; @@ -21,6 +23,16 @@ use crate::{ const SUBSCRIPTION_QOS: QoS = QoS::AtLeastOnce; +/// Minimum time between forced reconnects triggered by a withheld ack (see +/// `ack_suppressed`). Brokers that only redeliver in-flight QoS 1 messages on +/// reconnect would otherwise leave a withheld publish (and everything +/// suppressed after it) stuck until some unrelated reconnect happens, so a +/// forced reconnect gives it another chance. But if the downstream failure +/// causing the suppression is persistent, the redelivered publish will likely +/// fail again immediately, and reconnecting on every failure would thrash the +/// connection; this cooldown caps that to at most once per interval. +const FORCED_RECONNECT_COOLDOWN: Duration = Duration::from_secs(30); + /// 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 @@ -34,7 +46,13 @@ struct FinalizerEntry { #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] struct ProtocolState { connected: bool, + /// A (re)subscribe is needed and has not yet been confirmed by a SUBACK. pending_resubscribe: bool, + /// A (re)subscribe request has been queued and we're waiting to see + /// whether it results in a SUBACK, so `loop_actions` doesn't keep + /// resending it every iteration in the meantime. Cleared by `on_suback` + /// (success) or by queueing failing again (so the next iteration retries). + resubscribe_awaiting_suback: bool, connection_generation: u64, /// Set once a publish in the current generation finalizes as anything /// other than `Delivered`. MQTT requires PUBACKs to be sent in the order @@ -74,11 +92,19 @@ struct PublishAckDecision { warn_unsupported_qos: bool, } +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct FinalizedPublishDecision { + should_ack: bool, + just_suppressed: bool, +} + impl ProtocolState { const fn loop_actions(&self) -> LoopActions { LoopActions { retry_pending_acks: self.connected, - retry_resubscribe: self.connected && self.pending_resubscribe, + retry_resubscribe: self.connected + && self.pending_resubscribe + && !self.resubscribe_awaiting_suback, } } @@ -87,16 +113,32 @@ impl ProtocolState { acknowledgements: bool, session_present: bool, ) -> ConnAckActions { + // `session_present` alone can't be trusted when a previous (re)subscribe + // was never confirmed by a SUBACK: the broker's persistent session can + // exist (session_present = true) while missing that subscription, + // because the SUBSCRIBE that would have added it was lost (e.g. a + // disconnect racing the send) before this reconnect. So once a + // resubscribe is needed, it stays needed across reconnects regardless + // of what this connection's `session_present` says, until a SUBACK + // actually confirms it (`on_suback`). + let resubscribe = + self.pending_resubscribe || (self.connection_generation > 0 && !session_present); + let actions = ConnAckActions { warn_session_not_resumed: acknowledgements && !session_present, clear_pending_acks: true, recreate_finalizer: true, - resubscribe: self.connection_generation > 0 && !session_present, + resubscribe, }; self.connected = true; self.connection_generation += 1; self.ack_suppressed = false; + self.pending_resubscribe = resubscribe; + // A fresh connection means any previous queue attempt no longer + // applies (its SUBACK, if it ever arrives, would be for a subscribe + // on a connection that's gone) -- retry immediately on this one. + self.resubscribe_awaiting_suback = false; actions } @@ -110,25 +152,54 @@ impl ProtocolState { } } - const fn on_resubscribe_result(&mut self, success: bool) { - self.pending_resubscribe = !success; + /// Records whether a (re)subscribe request was successfully queued. + /// `pending_resubscribe` is intentionally left set either way: queueing + /// only means rumqttc accepted the request locally, not that the broker + /// received or processed it, so it's only cleared once a SUBACK actually + /// confirms success (`on_suback`). A request that's lost (e.g. a + /// disconnect racing the send) is retried on the next reconnect. + const fn on_resubscribe_result(&mut self, queued: bool) { + self.resubscribe_awaiting_suback = queued; + } + + /// A SUBACK confirms the broker processed the (re)subscribe, so stop + /// retrying it. + const fn on_suback(&mut self) { + self.pending_resubscribe = false; + self.resubscribe_awaiting_suback = false; } - /// Whether this finalized publish should be acked. Also updates - /// `ack_suppressed`: see its doc comment for why a non-`Delivered` status - /// in the current generation withholds every ack after it, not just its - /// own. - fn should_ack_finalized_publish(&mut self, status: BatchStatus, entry_generation: u64) -> bool { + /// Decides whether a finalized publish should be acked, and whether this + /// is the moment its generation's `ack_suppressed` first became set. See + /// `ack_suppressed`'s doc comment for why a non-`Delivered` status + /// withholds every ack after it, not just its own; `just_suppressed` is + /// how the caller learns it needs to force a reconnect so the broker + /// redelivers the withheld publish (see `FORCED_RECONNECT_COOLDOWN`). + fn finalize_publish( + &mut self, + status: BatchStatus, + entry_generation: u64, + ) -> FinalizedPublishDecision { if entry_generation != self.connection_generation { - return false; + return FinalizedPublishDecision { + should_ack: false, + just_suppressed: false, + }; } if status != BatchStatus::Delivered { + let just_suppressed = !self.ack_suppressed; self.ack_suppressed = true; - return false; + return FinalizedPublishDecision { + should_ack: false, + just_suppressed, + }; } - !self.ack_suppressed + FinalizedPublishDecision { + should_ack: !self.ack_suppressed, + just_suppressed: false, + } } } @@ -168,6 +239,13 @@ fn warn_resubscribe_failed() { ); } +fn warn_forcing_reconnect_after_suppressed_ack() { + warn!( + message = "An MQTT publish was not delivered and its acknowledgement was withheld; forcing a reconnect so the broker redelivers it.", + internal_log_rate_limit = true, + ); +} + #[derive(Default)] struct PendingAcks { publishes: Vec, @@ -256,6 +334,14 @@ impl MqttSource { let mut protocol_state = ProtocolState::default(); let mut pending_acks = PendingAcks::default(); + // Set when a graceful MQTT `Disconnect` (forced by a withheld ack, see + // `FORCED_RECONNECT_COOLDOWN`) couldn't be queued immediately (rumqttc's + // request channel was full); retried below until it succeeds. Cleared + // on the next disconnect, since by then the reconnect it was asking + // for has already happened. + let mut forced_reconnect_pending = false; + let mut forced_reconnect_cooldown_until: Option = None; + loop { let actions = protocol_state.loop_actions(); if actions.retry_resubscribe { @@ -270,6 +356,10 @@ impl MqttSource { pending_acks.retry(&client); } + if forced_reconnect_pending { + forced_reconnect_pending = client.try_disconnect().is_err(); + } + tokio::select! { _ = shutdown.clone() => return Ok(()), entry = ack_stream.next() => { @@ -280,13 +370,36 @@ impl MqttSource { // 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 - && protocol_state.should_ack_finalized_publish( + if let Some((status, entry)) = entry { + let decision = protocol_state.finalize_publish( status, entry.connection_generation, - ) - { - pending_acks.try_ack(protocol_state.connected, entry.publish, &client); + ); + if decision.should_ack { + pending_acks.try_ack(protocol_state.connected, entry.publish, &client); + } + // A withheld ack leaves the publish (and everything + // suppressed after it) stuck until the connection + // reconnects, so force one -- rate-limited so a + // persistent downstream failure (the redelivered + // publish failing again immediately) can't thrash the + // connection. Graceful (`try_disconnect`, an MQTT + // `Disconnect` packet) rather than dropping the + // network directly: an earlier attempt at the latter + // (`EventLoop::clean()`) was verified against a live + // EMQX broker to trigger a broker-side + // `unexpected_sock_close` error that stopped delivery + // entirely, whereas an explicit `Disconnect` is a + // normal, expected client-initiated close. + if decision.just_suppressed { + let now = tokio::time::Instant::now(); + if forced_reconnect_cooldown_until.is_none_or(|until| now >= until) { + forced_reconnect_cooldown_until = + Some(now + FORCED_RECONNECT_COOLDOWN); + warn_forcing_reconnect_after_suppressed_ack(); + forced_reconnect_pending = client.try_disconnect().is_err(); + } + } } }, mqtt_event = connection.poll() => { @@ -308,18 +421,23 @@ impl MqttSource { protocol_state.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, - ); + Ok(MqttEvent::Incoming(Incoming::SubAck(suback))) => { + // A SUBACK confirms the broker processed our (re)subscribe; + // stop retrying it regardless of whether acknowledgements + // are enabled (retry tracking isn't ack-specific). + protocol_state.on_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, + ); + } } } } @@ -344,16 +462,34 @@ impl MqttSource { 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. + // Connection lost: same stale-packet-id reasoning. `poll()` + // already called `EventLoop::clean()` internally before + // returning this error, which moves any PUBACK/PUBREC that + // `try_ack` had queued but that hadn't reached the network yet + // into `connection.pending` for automatic replay on the next + // connection. Those reference packet ids from the connection + // that just died, so drop them here too, or the broker could + // receive a stale ack for the wrong publish (or be sent one at + // all when it started a fresh session) once reconnected. Err(_) => { let actions = protocol_state.on_disconnect(); if actions.clear_pending_acks { pending_acks.clear(); + connection.pending.retain(|request| { + !matches!( + request, + rumqttc::Request::PubAck(_) | rumqttc::Request::PubRec(_) + ) + }); } if actions.recreate_finalizer { (finalizer, ack_stream) = self.new_finalizer(&shutdown); } + // The reconnect a pending forced-disconnect was + // asking for has now happened (by this error, if + // not by the disconnect itself); no need to send + // another once reconnected. + forced_reconnect_pending = false; } _ => {} } @@ -370,7 +506,7 @@ impl MqttSource { // survive the flush and get pushed into the set anyway. Because // `OrderedFinalizer` won't yield anything newer until that stale entry // resolves, it can hold back acks for the new connection generation even - // though `should_ack_finalized_publish` would correctly skip it once + // though `finalize_publish` would correctly skip it once // yielded. Dropping the old `finalizer` (and so its sender) instead lets // its background task see the channel closed and exit outright, taking // any not-yet-picked-up entries with it — no race window. @@ -667,19 +803,45 @@ mod tests { assert_eq!(fresh_session.connection_generation, 2); } + // Queueing a (re)subscribe request only means rumqttc accepted it + // locally, not that the broker processed it: a disconnect can lose it + // between queueing and actually reaching the broker, and the broker's + // next SUBACK-less reconnect can report `session_present = true` for a + // session that never got the subscription. So `pending_resubscribe` may + // only be cleared by an actual SUBACK (`on_suback`), and a successfully + // queued request must stop the loop from resending it every iteration + // (`resubscribe_awaiting_suback`) without yet considering it done. #[test] fn protocol_contract_matrix_for_pending_resubscribe() { let mut state = ProtocolState::default(); - state.on_resubscribe_result(false); + state.on_connack(true, true); + state.on_disconnect(); + let actions = state.on_connack(true, false); + assert!(actions.resubscribe); assert_eq!( state.loop_actions(), LoopActions { - retry_pending_acks: false, + retry_pending_acks: true, + retry_resubscribe: true, + } + ); + + // Successfully queueing it stops the per-iteration resend, but does + // not yet count as done: only a SUBACK can do that. + state.on_resubscribe_result(true); + assert_eq!( + state.loop_actions(), + LoopActions { + retry_pending_acks: true, retry_resubscribe: false, } ); - state.on_connack(true, true); + // The queued request never reaches the broker (lost in a race with + // a disconnect); failing to queue it again resumes the per-iteration + // retry rather than leaving it stuck waiting for a SUBACK that will + // never come. + state.on_resubscribe_result(false); assert_eq!( state.loop_actions(), LoopActions { @@ -688,7 +850,9 @@ mod tests { } ); + // Queued again, and this time a SUBACK actually confirms it. state.on_resubscribe_result(true); + state.on_suback(); assert_eq!( state.loop_actions(), LoopActions { @@ -698,6 +862,40 @@ mod tests { ); } + // A fresh-session reconnect whose queued SUBSCRIBE is lost (dropped + // before reaching the broker) must not be considered resubscribed just + // because a later reconnect happens to report `session_present = true` + // for that now-subscription-less session: only a SUBACK may clear + // `pending_resubscribe`. + #[test] + fn resubscribe_stays_pending_across_a_session_present_reconnect_without_suback() { + let mut state = ProtocolState::default(); + state.on_connack(true, true); + state.on_disconnect(); + state.on_connack(true, false); + state.on_resubscribe_result(true); + assert!(state.pending_resubscribe); + + // The connection drops again before a SUBACK arrives, and the next + // reconnect reports the (empty) session as present. + state.on_disconnect(); + let actions = state.on_connack(true, true); + + assert!( + actions.resubscribe, + "must still resubscribe: the previous attempt was never confirmed by a SUBACK, \ + even though this reconnect reports session_present = true" + ); + assert!(state.pending_resubscribe); + assert_eq!( + state.loop_actions(), + LoopActions { + retry_pending_acks: true, + retry_resubscribe: true, + } + ); + } + #[test] fn protocol_contract_matrix_for_disconnect() { let mut state = ProtocolState::default(); @@ -736,7 +934,7 @@ mod tests { (BatchStatus::Rejected, 2, false), ] { assert_eq!( - state.should_ack_finalized_publish(status, entry_generation), + state.finalize_publish(status, entry_generation).should_ack, should_ack ); } @@ -752,14 +950,67 @@ mod tests { let mut state = ProtocolState::default(); state.on_connack(true, true); - assert!(state.should_ack_finalized_publish(BatchStatus::Delivered, 1)); + assert!( + state + .finalize_publish(BatchStatus::Delivered, 1) + .should_ack + ); - assert!(!state.should_ack_finalized_publish(BatchStatus::Rejected, 1)); - assert!(!state.should_ack_finalized_publish(BatchStatus::Delivered, 1)); - assert!(!state.should_ack_finalized_publish(BatchStatus::Delivered, 1)); + assert!( + !state + .finalize_publish(BatchStatus::Rejected, 1) + .should_ack + ); + assert!( + !state + .finalize_publish(BatchStatus::Delivered, 1) + .should_ack + ); + assert!( + !state + .finalize_publish(BatchStatus::Delivered, 1) + .should_ack + ); // A reconnect starts a fresh generation, resuming normal acking. state.on_connack(true, false); - assert!(state.should_ack_finalized_publish(BatchStatus::Delivered, 2)); + assert!( + state + .finalize_publish(BatchStatus::Delivered, 2) + .should_ack + ); + } + + // A withheld ack must trigger a forced reconnect (so the broker + // redelivers the withheld publish) exactly once per suppression episode, + // not on every subsequent non-`Delivered` finalization in the same + // generation -- the caller only needs one signal to act on, and a fresh + // generation (after that reconnect) can trigger it again if it too fails. + #[test] + fn finalize_publish_signals_just_suppressed_once_per_episode() { + let mut state = ProtocolState::default(); + state.on_connack(true, true); + + let first = state.finalize_publish(BatchStatus::Rejected, 1); + assert!(!first.should_ack); + assert!(first.just_suppressed); + + let second = state.finalize_publish(BatchStatus::Errored, 1); + assert!(!second.should_ack); + assert!( + !second.just_suppressed, + "already suppressed in this generation; must not re-signal" + ); + + // A new generation can independently signal suppression again. + state.on_connack(true, false); + let after_reconnect = state.finalize_publish(BatchStatus::Rejected, 2); + assert!(after_reconnect.just_suppressed); + + // An out-of-generation (stale) entry never signals, even if its + // status would otherwise suppress. + let stale = state.finalize_publish(BatchStatus::Rejected, 1); + assert!(!stale.should_ack); + assert!(!stale.just_suppressed); } } From ae27de3b835f531286817596fc02a75043249dfa Mon Sep 17 00:00:00 2001 From: colin Date: Fri, 10 Jul 2026 10:32:52 -0700 Subject: [PATCH 6/8] fix(mqtt source): harden reconnect/ack state machine against liveness holes and cross-layer staleness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One fix for the latest review finding, plus six more from a deeper audit of the same failure classes (states whose required action has no wake-up path, Vector state diverging from rumqttc's internals, and locally-queued mistaken for peer-confirmed): - A withheld-ack forced reconnect landing inside the cooldown window is now scheduled for cooldown expiry instead of skipped: `just_suppressed` fires only once per generation, so a skipped reconnect never got another trigger and the source stayed connected (and stuck) indefinitely. - On disconnect, clear rumqttc's buffered `state.events`: `clean()` doesn't, so packets batched from the dead connection were yielded after the next ConnAck — a stale Publish got tagged with the new generation (its stale packet id then acked on the new connection), and a stale SubAck falsely confirmed the new connection's subscribe. Also drop a queued `Request::Disconnect` from `connection.pending`, which would otherwise be replayed onto (and immediately kill) the next connection. - SUBACK failure return codes (0x80, e.g. an ACL denial) no longer count as subscribed: the rejected topic is logged at error level (return codes map to topics by order, [MQTT-3.9.3-1]) and the (re)subscribe is retried on a 30s pace instead of being silently abandoned. - The initial SUBSCRIBE now goes through the same SUBACK-confirmation tracking as re-subscribes (issued on the first ConnAck, driven by `pending_resubscribe` starting true). Previously a lost initial SUBSCRIBE followed by a session_present reconnect left the source permanently subscribed to nothing. - The forced reconnect no longer assumes the broker closes the socket after our DISCONNECT (rumqttc never closes its side; the server only SHOULD): if no disconnect is observed within 10s, drop the connection locally, which [MQTT-3.14.4-2] requires of the client anyway. While a DISCONNECT is unanswered, nothing further is written to the connection ([MQTT-3.14.4-1]). - A newly finalized PUBACK now queues behind earlier acks still parked for retry instead of potentially jumping ahead of them ([MQTT-4.6.0-2]). - `keep_alive = 0` is rejected when acknowledgements are enabled: with the heartbeat disabled, a silently dead connection is never detected on a quiet topic, so unacked messages would never be redelivered. Verified against a live EMQX broker: all 4 integration tests (including the forced-reconnect redelivery test) rerun 4x stable; 20 unit tests cover the new state transitions. --- src/common/mqtt.rs | 5 + src/sources/mqtt/config.rs | 27 ++ src/sources/mqtt/source.rs | 554 +++++++++++++++++++++++++++++++------ 3 files changed, 500 insertions(+), 86 deletions(-) diff --git a/src/common/mqtt.rs b/src/common/mqtt.rs index d9e982815b36b..98e845b414a8a 100644 --- a/src/common/mqtt.rs +++ b/src/common/mqtt.rs @@ -109,6 +109,11 @@ pub enum ConfigurationError { "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, + /// Acknowledgements enabled with keep-alive disabled + #[snafu(display( + "`keep_alive` must be at least 1 second when `acknowledgements` are enabled: with keep-alive disabled, a silently dead connection is never detected on a quiet topic, and unacknowledged messages would never be redelivered." + ))] + AcknowledgementsRequireKeepAlive, } #[derive(Clone)] diff --git a/src/sources/mqtt/config.rs b/src/sources/mqtt/config.rs index 86b845d7b7770..58519dc03a8b1 100644 --- a/src/sources/mqtt/config.rs +++ b/src/sources/mqtt/config.rs @@ -135,6 +135,15 @@ impl MqttSourceConfig { .context(ConfigurationSnafu); } + // The ack machinery's liveness (detecting dead connections so unacked + // messages get redelivered, and its retry timers making progress on a + // quiet topic) depends on keep-alive traffic; `keep_alive = 0` + // disables it entirely in rumqttc. + if acknowledgements && self.common.keep_alive == 0 { + return Err(ConfigurationError::AcknowledgementsRequireKeepAlive) + .context(ConfigurationSnafu); + } + let client_id = self.common.client_id.clone().unwrap_or_else(|| { let hash = rand::rng() .sample_iter(&rand_distr::Alphanumeric) @@ -224,4 +233,22 @@ mod test { }; assert!(with_client_id.build_connector(true).is_ok()); } + + // With keep-alive disabled a silently dead connection is never detected + // on a quiet topic, so unacknowledged messages would never be redelivered + // and the ack machinery's retry timers would never make progress. + #[test] + fn acknowledgements_require_keep_alive() { + let config = MqttSourceConfig { + common: MqttCommonConfig { + client_id: Some("stable-id".to_owned()), + keep_alive: 0, + ..Default::default() + }, + ..Default::default() + }; + assert!(config.build_connector(true).is_err()); + // Without acknowledgements, disabling keep-alive stays allowed. + assert!(config.build_connector(false).is_ok()); + } } diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index c392bacc1e40d..48805ade9d104 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -33,6 +33,27 @@ const SUBSCRIPTION_QOS: QoS = QoS::AtLeastOnce; /// connection; this cooldown caps that to at most once per interval. const FORCED_RECONNECT_COOLDOWN: Duration = Duration::from_secs(30); +/// How long to wait, after sending a graceful `Disconnect`, for the broker to +/// close the connection before dropping it locally. [MQTT-3.14.4-2] puts the +/// close obligation on the client anyway; the server only SHOULD close. A +/// broker that keeps the socket open would otherwise leave `ack_suppressed` +/// pinned for the rest of the (never-ending) generation, stalling delivery +/// with nothing left to trigger recovery. +const FORCED_RECONNECT_ESCALATION_TIMEOUT: Duration = Duration::from_secs(10); + +/// How long to wait before re-attempting to queue a forced-reconnect +/// `Disconnect` after rumqttc's request channel was too full to accept it. +/// Non-zero so the retry timer can't monopolize the select loop while the +/// event loop drains the channel. +const FORCED_RECONNECT_RETRY_DELAY: Duration = Duration::from_millis(100); + +/// How long to wait before resending a (re)subscribe after the broker +/// rejected it (SUBACK failure return code, e.g. an ACL denial). Immediate +/// retries would flood the broker with SUBSCRIBE packets at network round-trip +/// speed; never retrying would leave the source silently subscribed to +/// nothing on an otherwise healthy connection if the denial is later lifted. +const RESUBSCRIBE_RETRY_DELAY: Duration = Duration::from_secs(30); + /// 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 @@ -47,6 +68,9 @@ struct FinalizerEntry { struct ProtocolState { connected: bool, /// A (re)subscribe is needed and has not yet been confirmed by a SUBACK. + /// Starts `true`: the initial subscribe needs SUBACK confirmation exactly + /// like a re-subscribe does (a lost initial SUBSCRIBE followed by a + /// `session_present = true` reconnect would otherwise never be retried). pending_resubscribe: bool, /// A (re)subscribe request has been queued and we're waiting to see /// whether it results in a SUBACK, so `loop_actions` doesn't keep @@ -98,6 +122,72 @@ struct FinalizedPublishDecision { just_suppressed: bool, } +/// Schedules the graceful forced reconnects triggered by withheld acks (see +/// `FORCED_RECONNECT_COOLDOWN`). A suppression inside the cooldown window +/// must schedule the reconnect for when the cooldown expires rather than skip +/// it: `just_suppressed` fires only once per connection generation, so a +/// skipped reconnect would never get another trigger and the source would +/// stay connected (and stuck) indefinitely. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct ForcedReconnectSchedule { + due_at: Option, + cooldown_until: Option, + escalate_at: Option, +} + +impl ForcedReconnectSchedule { + /// An ack was just withheld: schedule a reconnect, immediately if outside + /// the cooldown window, at cooldown expiry otherwise. + fn schedule(&mut self, now: tokio::time::Instant) { + let at = self.cooldown_until.map_or(now, |until| until.max(now)); + self.due_at = Some(self.due_at.map_or(at, |already_due| already_due.min(at))); + } + + /// When the scheduled reconnect should fire, if one is scheduled. + const fn due_at(&self) -> Option { + self.due_at + } + + /// The disconnect was queued: clear the schedule, start the cooldown, and + /// arm the escalation deadline in case the broker never closes the + /// connection in response (see `FORCED_RECONNECT_ESCALATION_TIMEOUT`). + fn sent(&mut self, now: tokio::time::Instant) { + self.due_at = None; + self.cooldown_until = Some(now + FORCED_RECONNECT_COOLDOWN); + self.escalate_at = Some(now + FORCED_RECONNECT_ESCALATION_TIMEOUT); + } + + /// The disconnect couldn't be queued (rumqttc's request channel was + /// full): retry shortly, once the event loop has drained the channel. + fn retry(&mut self, now: tokio::time::Instant) { + self.due_at = Some(now + FORCED_RECONNECT_RETRY_DELAY); + } + + /// When to drop the connection locally because the broker never closed it + /// after our `Disconnect`, if one was sent and is still unanswered. + const fn escalate_at(&self) -> Option { + self.escalate_at + } + + /// A `Disconnect` has been sent and the connection hasn't dropped yet, so + /// nothing further should be written to it ([MQTT-3.14.4-1]). + const fn disconnect_sent(&self) -> bool { + self.escalate_at.is_some() + } + + /// The escalation fired: the local teardown is happening now. + const fn escalated(&mut self) { + self.escalate_at = None; + } + + /// The connection dropped: whatever reconnect or escalation was scheduled + /// has been achieved by that disconnect. + const fn cancel(&mut self) { + self.due_at = None; + self.escalate_at = None; + } +} + impl ProtocolState { const fn loop_actions(&self) -> LoopActions { LoopActions { @@ -162,8 +252,10 @@ impl ProtocolState { self.resubscribe_awaiting_suback = queued; } - /// A SUBACK confirms the broker processed the (re)subscribe, so stop - /// retrying it. + /// A SUBACK granting every requested topic filter confirms the + /// (re)subscribe, so stop retrying it. Only called when no return code + /// was a failure; a rejected filter keeps `pending_resubscribe` set and + /// the retry is paced by `RESUBSCRIBE_RETRY_DELAY` instead. const fn on_suback(&mut self) { self.pending_resubscribe = false; self.resubscribe_awaiting_suback = false; @@ -276,7 +368,16 @@ impl PendingAcks { publish: Publish, mut try_ack: impl FnMut(&Publish) -> bool, ) { - if connected && !try_ack(&publish) { + if !connected { + return; + } + // Earlier acks are still queued (the loop-top retry couldn't flush + // them this iteration): this newer ack must queue behind them, not + // jump ahead -- PUBACKs must go out in publish-receipt order + // ([MQTT-4.6.0-2]), and rumqttc's channel can free up between the + // failed retry and now (the intervening `connection.poll()` may have + // partially drained it before being cancelled by this branch). + if !self.publishes.is_empty() || !try_ack(&publish) { self.push(publish); } } @@ -314,8 +415,6 @@ impl MqttSource { pub async fn run(self, mut out: SourceSender, shutdown: ShutdownSignal) -> Result<(), ()> { let (client, mut connection) = self.connector.connect(); - 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 @@ -331,20 +430,31 @@ 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(); - // Set when a graceful MQTT `Disconnect` (forced by a withheld ack, see - // `FORCED_RECONNECT_COOLDOWN`) couldn't be queued immediately (rumqttc's - // request channel was full); retried below until it succeeds. Cleared - // on the next disconnect, since by then the reconnect it was asking - // for has already happened. - let mut forced_reconnect_pending = false; - let mut forced_reconnect_cooldown_until: Option = None; + // The initial subscribe is issued on the first ConnAck, driven by + // `pending_resubscribe` starting `true`, so it goes through the same + // SUBACK-confirmation tracking as every re-subscribe. + let mut protocol_state = ProtocolState { + pending_resubscribe: true, + ..ProtocolState::default() + }; + + let mut forced_reconnect = ForcedReconnectSchedule::default(); + + // When the broker rejects a (re)subscribe (SUBACK failure return + // code), the retry is paced by this timer instead of resent + // immediately; see `RESUBSCRIBE_RETRY_DELAY`. + let mut resubscribe_retry_at: Option = None; loop { + // Once a graceful `Disconnect` has been queued, nothing else may + // be written to the connection ([MQTT-3.14.4-1]); the retries + // resume after the reconnect it causes. + let draining = forced_reconnect.disconnect_sent(); + let actions = protocol_state.loop_actions(); - if actions.retry_resubscribe { + if actions.retry_resubscribe && !draining { protocol_state.on_resubscribe_result(self.try_subscribe(&client)); } @@ -352,14 +462,10 @@ impl MqttSource { // 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 { + if actions.retry_pending_acks && !draining { pending_acks.retry(&client); } - if forced_reconnect_pending { - forced_reconnect_pending = client.try_disconnect().is_err(); - } - tokio::select! { _ = shutdown.clone() => return Ok(()), entry = ack_stream.next() => { @@ -380,28 +486,73 @@ impl MqttSource { } // A withheld ack leaves the publish (and everything // suppressed after it) stuck until the connection - // reconnects, so force one -- rate-limited so a - // persistent downstream failure (the redelivered - // publish failing again immediately) can't thrash the - // connection. Graceful (`try_disconnect`, an MQTT - // `Disconnect` packet) rather than dropping the - // network directly: an earlier attempt at the latter - // (`EventLoop::clean()`) was verified against a live - // EMQX broker to trigger a broker-side - // `unexpected_sock_close` error that stopped delivery - // entirely, whereas an explicit `Disconnect` is a - // normal, expected client-initiated close. + // reconnects, so schedule a forced reconnect: sent by + // the timer branch below, immediately if outside the + // cooldown window, at cooldown expiry otherwise (so a + // persistent downstream failure -- the redelivered + // publish failing again right away -- reconnects at a + // capped rate instead of thrashing the connection or, + // worse, never again). if decision.just_suppressed { - let now = tokio::time::Instant::now(); - if forced_reconnect_cooldown_until.is_none_or(|until| now >= until) { - forced_reconnect_cooldown_until = - Some(now + FORCED_RECONNECT_COOLDOWN); - warn_forcing_reconnect_after_suppressed_ack(); - forced_reconnect_pending = client.try_disconnect().is_err(); - } + forced_reconnect.schedule(tokio::time::Instant::now()); } } }, + // A scheduled forced reconnect is due: gracefully disconnect + // (`try_disconnect` queues an MQTT `Disconnect` packet) so the + // broker redelivers the withheld publish once reconnected. + // Graceful rather than dropping the network directly: an + // earlier attempt at the latter (`EventLoop::clean()`) was + // verified against a live EMQX broker to trigger a broker-side + // `unexpected_sock_close` error that stopped delivery + // entirely, whereas an explicit `Disconnect` is a normal, + // expected client-initiated close. + _ = tokio::time::sleep_until( + forced_reconnect.due_at().unwrap_or_else(tokio::time::Instant::now) + ), if forced_reconnect.due_at().is_some() && protocol_state.connected => { + let now = tokio::time::Instant::now(); + if client.try_disconnect().is_ok() { + warn_forcing_reconnect_after_suppressed_ack(); + forced_reconnect.sent(now); + } else { + // rumqttc's request channel is full; the event loop + // branch below drains it, so retrying stays live. + forced_reconnect.retry(now); + } + }, + // The broker never closed the connection in response to our + // `Disconnect`: close it locally, as [MQTT-3.14.4-2] requires + // of the client anyway. Safe to be abrupt here -- the session + // was already ended gracefully by the `Disconnect` packet, so + // this is not the mid-session socket drop that upset brokers + // when tried previously (see the comment on the branch above). + _ = tokio::time::sleep_until( + forced_reconnect.escalate_at().unwrap_or_else(tokio::time::Instant::now) + ), if forced_reconnect.escalate_at().is_some() => { + forced_reconnect.escalated(); + connection.clean(); + (finalizer, ack_stream) = self.handle_connection_lost( + &mut connection, + &mut protocol_state, + &mut pending_acks, + &mut forced_reconnect, + &mut resubscribe_retry_at, + &shutdown, + ); + }, + // The broker rejected a (re)subscribe earlier; the pacing + // delay has passed, so try again. + _ = tokio::time::sleep_until( + resubscribe_retry_at.unwrap_or_else(tokio::time::Instant::now) + ), if resubscribe_retry_at.is_some() => { + resubscribe_retry_at = None; + if protocol_state.connected + && protocol_state.pending_resubscribe + && !forced_reconnect.disconnect_sent() + { + protocol_state.on_resubscribe_result(self.try_subscribe(&client)); + } + }, mqtt_event = connection.poll() => { // Providing at-least-once here does not require correlating a // connection/poll error back to a specific in-flight publish. @@ -422,24 +573,48 @@ impl MqttSource { ).await; } Ok(MqttEvent::Incoming(Incoming::SubAck(suback))) => { - // A SUBACK confirms the broker processed our (re)subscribe; - // stop retrying it regardless of whether acknowledgements - // are enabled (retry tracking isn't ack-specific). - protocol_state.on_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, + // A SUBACK only confirms the broker *processed* the + // (re)subscribe; each return code says whether the + // matching topic filter was actually granted, in + // SUBSCRIBE order ([MQTT-3.9.3-1]). A rejected + // filter (failure code 0x80, e.g. an ACL denial) + // must not count as subscribed -- keep retrying, + // paced by `RESUBSCRIBE_RETRY_DELAY`. + let mut any_rejected = false; + for (topic, return_code) in + self.subscribed_topics().zip(suback.return_codes.iter()) + { + match return_code { + rumqttc::SubscribeReasonCode::Success(qos) => { + if self.acknowledgements + && !publish_supports_end_to_end_acknowledgements(*qos) + { + warn!( + message = "MQTT broker granted a subscription QoS below the level required for end-to-end acknowledgements.", + topic, + ?qos, + internal_log_rate_limit = true, + ); + } + } + rumqttc::SubscribeReasonCode::Failure => { + any_rejected = true; + error!( + message = "MQTT broker rejected the subscription; will retry.", + topic, internal_log_rate_limit = true, ); } } } + if any_rejected { + resubscribe_retry_at = Some( + tokio::time::Instant::now() + RESUBSCRIBE_RETRY_DELAY, + ); + } else { + protocol_state.on_suback(); + resubscribe_retry_at = None; + } } // A (re)connected session resumes here; the broker will // redeliver any unacknowledged publishes, so drop deferred @@ -461,35 +636,24 @@ impl MqttSource { if actions.resubscribe { protocol_state.on_resubscribe_result(self.try_subscribe(&client)); } + // A retry the previous connection had scheduled no + // longer applies; this connection's subscribe was + // just issued above (or is already confirmed). + resubscribe_retry_at = None; } - // Connection lost: same stale-packet-id reasoning. `poll()` - // already called `EventLoop::clean()` internally before - // returning this error, which moves any PUBACK/PUBREC that - // `try_ack` had queued but that hadn't reached the network yet - // into `connection.pending` for automatic replay on the next - // connection. Those reference packet ids from the connection - // that just died, so drop them here too, or the broker could - // receive a stale ack for the wrong publish (or be sent one at - // all when it started a fresh session) once reconnected. + // Connection lost: `poll()` already called + // `EventLoop::clean()` internally before returning this + // error; scrub everything from the dead connection that + // rumqttc would otherwise carry over into the next one. Err(_) => { - let actions = protocol_state.on_disconnect(); - if actions.clear_pending_acks { - pending_acks.clear(); - connection.pending.retain(|request| { - !matches!( - request, - rumqttc::Request::PubAck(_) | rumqttc::Request::PubRec(_) - ) - }); - } - if actions.recreate_finalizer { - (finalizer, ack_stream) = self.new_finalizer(&shutdown); - } - // The reconnect a pending forced-disconnect was - // asking for has now happened (by this error, if - // not by the disconnect itself); no need to send - // another once reconnected. - forced_reconnect_pending = false; + (finalizer, ack_stream) = self.handle_connection_lost( + &mut connection, + &mut protocol_state, + &mut pending_acks, + &mut forced_reconnect, + &mut resubscribe_retry_at, + &shutdown, + ); } _ => {} } @@ -498,18 +662,74 @@ impl MqttSource { } } + /// The connection is gone (either `poll()` returned an error, or we tore + /// it down locally after an unanswered `Disconnect`): reset all per- + /// connection state, and scrub everything rumqttc would otherwise replay + /// from the dead connection onto the next one. + fn handle_connection_lost( + &self, + connection: &mut rumqttc::EventLoop, + protocol_state: &mut ProtocolState, + pending_acks: &mut PendingAcks, + forced_reconnect: &mut ForcedReconnectSchedule, + resubscribe_retry_at: &mut Option, + shutdown: &ShutdownSignal, + ) -> ( + Option>, + BoxStream<'static, (BatchStatus, FinalizerEntry)>, + ) { + let actions = protocol_state.on_disconnect(); + if actions.clear_pending_acks { + pending_acks.clear(); + // `EventLoop::clean()` moved any queued-but-unsent requests into + // `connection.pending` for automatic replay on the next + // connection. A PUBACK/PUBREC references a packet id that was + // only valid on the connection that just died (replaying it could + // ack the wrong publish), and a replayed `Disconnect` would kill + // the fresh connection the moment it comes up -- drop both. + // (`Subscribe` requests are left in: replaying one is a harmless, + // idempotent re-subscribe.) + connection.pending.retain(|request| { + !matches!( + request, + rumqttc::Request::PubAck(_) + | rumqttc::Request::PubRec(_) + | rumqttc::Request::Disconnect(_) + ) + }); + // rumqttc buffers batches of incoming packets in `state.events` + // and yields them one per poll; `clean()` does NOT clear that + // buffer, so events from the dead connection would otherwise be + // yielded *after* the next ConnAck and be indistinguishable from + // new-connection traffic: a stale Publish would be tagged with + // the new generation (and its stale packet id acked on the new + // connection), and a stale SubAck would falsely confirm the new + // connection's subscribe. Everything buffered here belongs to the + // dead connection -- the reconnect ConnAck is returned directly + // by `poll()`, never through this buffer -- so drop it all. Any + // dropped publish is unacked and will be redelivered. + connection.state.events.clear(); + } + // The reconnect any scheduled forced-disconnect was asking for has + // now happened; no need to send another once reconnected. + forced_reconnect.cancel(); + *resubscribe_retry_at = None; + debug_assert!(actions.recreate_finalizer); + self.new_finalizer(shutdown) + } + // Builds a fresh finalizer/ack-stream pair, discarding whatever the // previous one held. Used instead of `FinalizerSet::flush` on - // (re)connect: `flush` only clears entries already pulled into its - // internal ordered set, not ones a concurrent `finalizer.add` call sent - // but that the finalizer's background task hadn't yet picked up: those - // survive the flush and get pushed into the set anyway. Because + // (re)connect: `flush` only clears entries already pulled into the + // stream's internal ordered set, not ones an earlier `finalizer.add` + // call sent that are still sitting in its channel; those survive the + // flush and get pulled into the set afterwards anyway. Because // `OrderedFinalizer` won't yield anything newer until that stale entry - // resolves, it can hold back acks for the new connection generation even - // though `finalize_publish` would correctly skip it once - // yielded. Dropping the old `finalizer` (and so its sender) instead lets - // its background task see the channel closed and exit outright, taking - // any not-yet-picked-up entries with it — no race window. + // resolves, it could hold back acks for the new connection generation + // even though `finalize_publish` would correctly skip it once yielded. + // Dropping BOTH halves of the old pair instead destroys the channel and + // the ordered set wholesale (the stream is polled inline as `ack_stream`, + // not by a spawned task), so nothing stale can survive into the new pair. fn new_finalizer( &self, shutdown: &ShutdownSignal, @@ -546,6 +766,17 @@ impl MqttSource { } } + /// The configured topic filters in the order `subscribe` sends them, + /// which per [MQTT-3.9.3-1] is also the order of a SUBACK's return codes. + fn subscribed_topics(&self) -> impl Iterator { + match &self.config.topic { + OneOrMany::One(topic) => std::slice::from_ref(topic), + OneOrMany::Many(topics) => topics.as_slice(), + } + .iter() + .map(String::as_str) + } + async fn process_message( &self, mut publish: Publish, @@ -687,6 +918,36 @@ mod tests { } } + // A newer ack must never jump ahead of earlier acks still queued for + // retry, even if the channel has room for it now -- PUBACKs go out in + // publish-receipt order ([MQTT-4.6.0-2]). + #[test] + fn pending_acks_new_ack_queues_behind_earlier_failures() { + let mut pending_acks = PendingAcks::default(); + + // An earlier ack failed and is parked for retry. + pending_acks.try_ack_with(true, publish(1), |_| false); + assert_eq!(pending_acks.publishes.len(), 1); + + // The channel now has room, but the newer ack must still queue + // behind the parked one rather than being sent. + let mut attempted = false; + pending_acks.try_ack_with(true, publish(2), |_| { + attempted = true; + true + }); + assert!(!attempted); + assert_eq!( + pending_acks + .publishes + .iter() + .map(|publish| publish.pkid) + .collect::>(), + vec![1, 2], + "retry order must match receipt order" + ); + } + #[test] fn protocol_contract_matrix_for_requested_and_granted_qos() { assert!(publish_supports_end_to_end_acknowledgements( @@ -1013,4 +1274,125 @@ mod tests { assert!(!stale.should_ack); assert!(!stale.just_suppressed); } + + // A suppression that lands inside the cooldown window must reschedule the + // forced reconnect for cooldown expiry, not skip it: `just_suppressed` + // fires only once per connection generation, so a skipped reconnect never + // gets another trigger and the source would stay connected (and stuck) + // indefinitely. + #[test] + fn forced_reconnect_inside_cooldown_is_deferred_not_skipped() { + let mut schedule = ForcedReconnectSchedule::default(); + let t0 = tokio::time::Instant::now(); + + // First suppression, no cooldown yet: due immediately. + schedule.schedule(t0); + assert_eq!(schedule.due_at(), Some(t0)); + schedule.sent(t0); + assert_eq!(schedule.due_at(), None); + + // The redelivered publish is rejected again 2s later, well within + // the cooldown: due at cooldown expiry, not dropped. + schedule.schedule(t0 + Duration::from_secs(2)); + assert_eq!(schedule.due_at(), Some(t0 + FORCED_RECONNECT_COOLDOWN)); + + let t1 = t0 + FORCED_RECONNECT_COOLDOWN; + schedule.sent(t1); + assert_eq!(schedule.due_at(), None); + + // A suppression after the cooldown has fully expired is due + // immediately again. + let t2 = t1 + FORCED_RECONNECT_COOLDOWN + Duration::from_secs(1); + schedule.schedule(t2); + assert_eq!(schedule.due_at(), Some(t2)); + } + + #[test] + fn forced_reconnect_schedule_retry_and_cancel() { + let mut schedule = ForcedReconnectSchedule::default(); + let t0 = tokio::time::Instant::now(); + + // Queueing the disconnect failed: retried shortly (not immediately, + // which would let the timer branch monopolize the select loop). + schedule.schedule(t0); + let t1 = t0 + Duration::from_millis(10); + schedule.retry(t1); + assert_eq!(schedule.due_at(), Some(t1 + FORCED_RECONNECT_RETRY_DELAY)); + + // A natural disconnect supersedes the scheduled one. + schedule.cancel(); + assert_eq!(schedule.due_at(), None); + } + + // rumqttc never closes its own side of the connection after sending + // DISCONNECT, and the broker is only required to SHOULD-close: if it + // doesn't, nothing else would ever end the generation whose acks are + // suppressed. Sending the disconnect must therefore arm an escalation + // deadline for closing the connection locally, cleared either by + // escalating or by the disconnect actually happening. + #[test] + fn forced_reconnect_escalates_when_broker_never_closes() { + let mut schedule = ForcedReconnectSchedule::default(); + let t0 = tokio::time::Instant::now(); + + assert!(!schedule.disconnect_sent()); + schedule.schedule(t0); + assert!(!schedule.disconnect_sent()); + + schedule.sent(t0); + assert!(schedule.disconnect_sent()); + assert_eq!( + schedule.escalate_at(), + Some(t0 + FORCED_RECONNECT_ESCALATION_TIMEOUT) + ); + + // The deadline fires: local teardown happens, nothing further armed. + schedule.escalated(); + assert!(!schedule.disconnect_sent()); + assert_eq!(schedule.escalate_at(), None); + + // Alternate path: the broker closes in time, the resulting poll error + // cancels both the schedule and the escalation. + schedule.sent(t0 + Duration::from_secs(60)); + assert!(schedule.disconnect_sent()); + schedule.cancel(); + assert!(!schedule.disconnect_sent()); + assert_eq!(schedule.escalate_at(), None); + } + + // The initial subscribe must be confirmed by a SUBACK exactly like a + // re-subscribe: if the first SUBSCRIBE is lost after the broker created + // the persistent session, the next reconnect reports + // `session_present = true` for a session with no subscription, and + // without this the source would never subscribe again. + #[test] + fn initial_subscribe_requires_suback_confirmation() { + // As `run()` initializes it. + let mut state = ProtocolState { + pending_resubscribe: true, + ..ProtocolState::default() + }; + + // First ConnAck must ask for the (initial) subscribe. + let actions = state.on_connack(true, false); + assert!(actions.resubscribe); + state.on_resubscribe_result(true); + + // The SUBSCRIBE is lost; the connection drops; the broker resumes + // the (subscription-less) session. + state.on_disconnect(); + let actions = state.on_connack(true, true); + assert!( + actions.resubscribe, + "initial subscribe was never SUBACK-confirmed, so a \ + session_present reconnect must still resubscribe" + ); + + // Once a SUBACK confirms it, later resumed sessions don't resend. + state.on_resubscribe_result(true); + state.on_suback(); + state.on_disconnect(); + let actions = state.on_connack(true, true); + assert!(!actions.resubscribe); + } } From 520593a19d072f4fa05edbc1446b40fb54ca2d9d Mon Sep 17 00:00:00 2001 From: colin Date: Fri, 10 Jul 2026 11:12:24 -0700 Subject: [PATCH 7/8] fix(mqtt source): keep auto-acked buffered publishes on disconnect, reject invalid topic filters at build Two fixes from review on #25785: - The disconnect-time scrub of rumqttc's buffered `state.events` was unconditional, but in auto-ack mode (acknowledgements off, the default) rumqttc queues the PUBACK the moment it buffers a publish -- it may already be on the wire, and the broker is then allowed to never redeliver. Dropping such a publish would lose it outright. The scrub is now mode-dependent: manual-ack mode still drops everything (nothing buffered has been acked, so the broker redelivers it all), while auto-ack mode keeps buffered publishes for processing and drops only the non-publish events (so e.g. a stale SubAck still can't falsely confirm the new connection's subscribe). - An invalid topic filter (e.g. `foo/#/bar`) is a permanent configuration error, but rumqttc queues the SUBSCRIBE without validating it, so it only surfaced at runtime as a source that receives nothing while the broker rejects the subscription over and over. Topic filters are now validated with `rumqttc::valid_filter` when the source is built, failing fast with an error naming the bad filter. --- src/common/mqtt.rs | 8 ++++++ src/sources/mqtt/config.rs | 52 ++++++++++++++++++++++++++++++++++++++ src/sources/mqtt/source.rs | 35 +++++++++++++++++++------ 3 files changed, 87 insertions(+), 8 deletions(-) diff --git a/src/common/mqtt.rs b/src/common/mqtt.rs index 98e845b414a8a..9beb5a165a7b8 100644 --- a/src/common/mqtt.rs +++ b/src/common/mqtt.rs @@ -114,6 +114,14 @@ pub enum ConfigurationError { "`keep_alive` must be at least 1 second when `acknowledgements` are enabled: with keep-alive disabled, a silently dead connection is never detected on a quiet topic, and unacknowledged messages would never be redelivered." ))] AcknowledgementsRequireKeepAlive, + /// Invalid topic filter provided + #[snafu(display( + "`{topic}` is not a valid MQTT topic filter: `#` must be the last, standalone level and `+` must occupy a whole level." + ))] + InvalidTopicFilter { + /// The invalid topic filter. + topic: String, + }, } #[derive(Clone)] diff --git a/src/sources/mqtt/config.rs b/src/sources/mqtt/config.rs index 58519dc03a8b1..853d8babc81ec 100644 --- a/src/sources/mqtt/config.rs +++ b/src/sources/mqtt/config.rs @@ -144,6 +144,21 @@ impl MqttSourceConfig { .context(ConfigurationSnafu); } + // An invalid filter is a permanent configuration error: rumqttc + // queues the SUBSCRIBE without validating it, so it would otherwise + // only surface at runtime as the broker rejecting (or dropping) the + // subscription over and over -- a running source that receives + // nothing. Fail here instead. + if let Some(topic) = self + .topic + .clone() + .to_vec() + .into_iter() + .find(|topic| !rumqttc::valid_filter(topic)) + { + return Err(ConfigurationError::InvalidTopicFilter { topic }).context(ConfigurationSnafu); + } + let client_id = self.common.client_id.clone().unwrap_or_else(|| { let hash = rand::rng() .sample_iter(&rand_distr::Alphanumeric) @@ -251,4 +266,41 @@ mod test { // Without acknowledgements, disabling keep-alive stays allowed. assert!(config.build_connector(false).is_ok()); } + + // An invalid topic filter is a permanent configuration error and must + // fail at build time: rumqttc queues the SUBSCRIBE without validating it, + // so it would otherwise only surface as a running source that receives + // nothing while the broker rejects the subscription over and over. + #[test] + fn invalid_topic_filters_fail_at_build() { + let config_with_topic = |topic: OneOrMany| MqttSourceConfig { + topic, + ..Default::default() + }; + + for invalid in ["foo/#/bar", "foo/bar#", "foo/b+r", ""] { + assert!( + config_with_topic(OneOrMany::One(invalid.into())) + .build_connector(false) + .is_err(), + "{invalid:?} must be rejected" + ); + } + + // One invalid filter among valid ones still fails. + assert!( + config_with_topic(OneOrMany::Many(vec!["ok/#".into(), "foo/#/bar".into()])) + .build_connector(false) + .is_err() + ); + + for valid in ["foo/#", "foo/+/bar", "+/tele/#", "plain/topic"] { + assert!( + config_with_topic(OneOrMany::One(valid.into())) + .build_connector(false) + .is_ok(), + "{valid:?} must be accepted" + ); + } + } } diff --git a/src/sources/mqtt/source.rs b/src/sources/mqtt/source.rs index 48805ade9d104..a9b3f5090c74f 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -701,14 +701,33 @@ impl MqttSource { // and yields them one per poll; `clean()` does NOT clear that // buffer, so events from the dead connection would otherwise be // yielded *after* the next ConnAck and be indistinguishable from - // new-connection traffic: a stale Publish would be tagged with - // the new generation (and its stale packet id acked on the new - // connection), and a stale SubAck would falsely confirm the new - // connection's subscribe. Everything buffered here belongs to the - // dead connection -- the reconnect ConnAck is returned directly - // by `poll()`, never through this buffer -- so drop it all. Any - // dropped publish is unacked and will be redelivered. - connection.state.events.clear(); + // new-connection traffic. Everything buffered here belongs to the + // dead connection (the reconnect ConnAck is returned directly by + // `poll()`, never through this buffer), but what's safe to drop + // depends on the ack mode: + if self.acknowledgements { + // Manual-ack mode: nothing buffered has been acked (acks are + // only sent after end-to-end finalization, and these publishes + // were never even yielded), so the broker redelivers all of + // it. Dropping everything prevents a stale Publish from being + // tagged with the new generation (its stale packet id then + // acked on the new connection) and a stale SubAck from + // falsely confirming the new connection's subscribe. + connection.state.events.clear(); + } else { + // Auto-ack mode: rumqttc queues the PUBACK the moment it + // buffers a publish, so it may already be on the wire and the + // broker is then allowed to never redeliver -- dropping a + // buffered publish here would lose it outright. Keep the + // publishes for processing (there's no generation-sensitive + // ack machinery in this mode to confuse); drop only the + // stale non-publish events, so e.g. a stale SubAck can't + // falsely confirm the new connection's subscribe. + connection + .state + .events + .retain(|event| matches!(event, MqttEvent::Incoming(Incoming::Publish(_)))); + } } // The reconnect any scheduled forced-disconnect was asking for has // now happened; no need to send another once reconnected. From 69760e688957474acd6eb55464062cd2253a7442 Mon Sep 17 00:00:00 2001 From: colin Date: Fri, 10 Jul 2026 13:06:03 -0700 Subject: [PATCH 8/8] test(mqtt source): exhaustively model-check the protocol state machine, fault-inject against a live broker Review on #25785 kept finding bugs in the gaps between hand-picked test cases: event sequences the unit tests didn't sample, and broker-side failure paths the integration tests never exercised (every existing test used acks-on, client-side aborts only). Close both gaps structurally instead of sampling harder: - `exhaustive_protocol_state_model_check` BFS-explores every reachable abstract state of `ProtocolState` + `ForcedReconnectSchedule` + the resubscribe retry timer under every legal event interleaving (timer state abstracted to armed/not-armed, so the space is finite), asserting on every transition the invariants past findings violated: a suppressed generation always has a recovery action armed, per-connection schedules die with the connection, an unconfirmed subscription always has a live retry path, `pending_resubscribe` is only cleared by a fully granted SUBACK, no write after an unanswered Disconnect. Mutation-verified: reintroducing either the cooldown-skip bug or the session_present resubscribe-clearing bug makes it fail with a minimal counterexample trace. - Three fault-injection integration tests force *server-side* disconnects via MQTT client-ID takeover ([MQTT-3.1.4-2], portable across brokers, no admin API): recovery with the session preserved (offline publishes delivered from the session queue), resubscribing after the broker discards the session entirely (`session_present = false` path, SUBACK-confirmed), and no message loss across a kick in the default auto-ack mode -- the mode and the disconnect direction that previously had zero coverage. All 29 tests pass against a live EMQX broker, rerun 5x stable. --- src/sources/mqtt/integration_tests.rs | 260 +++++++++++++++++- src/sources/mqtt/source.rs | 372 ++++++++++++++++++++++++++ 2 files changed, 630 insertions(+), 2 deletions(-) diff --git a/src/sources/mqtt/integration_tests.rs b/src/sources/mqtt/integration_tests.rs index 0701fab7a128a..fd0631deb2bec 100644 --- a/src/sources/mqtt/integration_tests.rs +++ b/src/sources/mqtt/integration_tests.rs @@ -3,8 +3,8 @@ use std::{collections::HashSet, time::Duration}; -use futures::StreamExt; -use rumqttc::{AsyncClient, MqttOptions, QoS}; +use futures::{Stream, StreamExt}; +use rumqttc::{AsyncClient, Event as MqttEvent, Incoming, MqttOptions, QoS}; use tokio::time::timeout; use crate::{ @@ -223,6 +223,58 @@ async fn mqtt_redelivers_unacknowledged_messages() { drop(source2.await); } +/// Forces a server-side disconnect of whatever client currently holds +/// `client_id`: a second connection with the same client ID makes the broker +/// disconnect the existing one ([MQTT-3.1.4-2]). With `clean_session = false` +/// the persistent session (subscriptions and unacknowledged in-flight +/// messages) survives the takeover; with `clean_session = true` the broker +/// additionally discards that session. The takeover connection is closed +/// immediately afterwards, freeing the client ID for the victim to +/// reconnect. +async fn kick_mqtt_client(client_id: &str, discard_session: bool) { + let mut options = MqttOptions::new(client_id, mqtt_broker_address(), mqtt_broker_port()); + options.set_keep_alive(Duration::from_secs(5)); + options.set_clean_session(discard_session); + let (client, mut eventloop) = AsyncClient::new(options, 10); + + // Wait until the broker accepts the takeover (kicking the victim). + loop { + match timeout(Duration::from_secs(5), eventloop.poll()) + .await + .expect("timed out connecting the takeover client") + { + Ok(MqttEvent::Incoming(Incoming::ConnAck(_))) => break, + Ok(_) => {} + Err(error) => panic!("takeover client failed to connect: {error:?}"), + } + } + + client.disconnect().await.unwrap(); + // Drive the DISCONNECT out; polling errors once the connection closes. + while timeout(Duration::from_secs(5), eventloop.poll()) + .await + .expect("timed out closing the takeover client") + .is_ok() + {} +} + +/// Receives from the source until every message in `expected` has been seen +/// at least once (duplicates are permitted: QoS 1 is at-least-once), +/// panicking if `each_timeout` passes without progress. +async fn collect_expected_messages( + rx: &mut (impl Stream + Unpin), + mut expected: HashSet, + each_timeout: Duration, +) { + while !expected.is_empty() { + let event = timeout(each_timeout, rx.next()) + .await + .unwrap_or_else(|_| panic!("timed out; still missing {} messages", expected.len())) + .expect("source stream ended unexpectedly"); + expected.remove(&message_body(&event)); + } +} + /// A withheld ack (the sink rejects the message, so its PUBACK is never sent) /// must force a reconnect on its own -- without an external trigger like a /// process restart or manual disconnect -- so the broker redelivers the @@ -294,6 +346,210 @@ async fn mqtt_forces_reconnect_after_withheld_ack() { drop(source.await); } +/// A server-initiated disconnect (broker kicks the client) must be survived +/// transparently when the session is preserved: the source reconnects, +/// resumes the session (`session_present = true`, no resubscribe needed), +/// and messages published while it was offline are delivered from the +/// session's queue. +#[tokio::test] +async fn mqtt_recovers_from_server_side_disconnect() { + trace_init(); + + let topic = "source-server-kick-test"; + let client_id = format!("sourceKickTest{}", random_string(6)); + let msg_before = random_string(32); + let msg_after = random_string(32); + + let 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() + }; + + let (tx, mut 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; + + let producer = get_mqtt_client().await; + producer + .publish(topic, QoS::AtLeastOnce, false, msg_before.as_bytes()) + .await + .unwrap(); + let received = timeout(Duration::from_secs(5), rx.next()) + .await + .expect("timed out waiting for pre-kick delivery") + .expect("source stream ended unexpectedly"); + assert_eq!(message_body(&received), msg_before); + + // Kick the source off the broker, keeping its session intact, and + // publish while it is (briefly) offline: the persistent session queues + // the message for delivery once the source has reconnected. + kick_mqtt_client(&client_id, false).await; + producer + .publish(topic, QoS::AtLeastOnce, false, msg_after.as_bytes()) + .await + .unwrap(); + + let redelivered = timeout(Duration::from_secs(10), rx.next()) + .await + .expect("timed out: the source did not recover from the server-side disconnect") + .expect("source stream ended unexpectedly"); + assert_eq!(message_body(&redelivered), msg_after); + + source.abort(); + drop(source.await); +} + +/// When the broker discards the persistent session entirely (here: a +/// takeover connection with `clean_session = true`, but a broker restart +/// without persistence behaves the same), the source's reconnect sees +/// `session_present = false` for a session with no subscriptions and must +/// resubscribe on its own before messages flow again. +#[tokio::test] +async fn mqtt_resubscribes_after_broker_discards_the_session() { + trace_init(); + + let topic = "source-session-loss-test"; + let client_id = format!("sourceSessionLossTest{}", random_string(6)); + let msg_before = random_string(32); + let msg_after = random_string(32); + + let 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() + }; + + let (tx, mut 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; + + let producer = get_mqtt_client().await; + producer + .publish(topic, QoS::AtLeastOnce, false, msg_before.as_bytes()) + .await + .unwrap(); + let received = timeout(Duration::from_secs(5), rx.next()) + .await + .expect("timed out waiting for pre-kick delivery") + .expect("source stream ended unexpectedly"); + assert_eq!(message_body(&received), msg_before); + + // Kick the source AND discard its session (subscriptions included). + kick_mqtt_client(&client_id, true).await; + + // A publish only reaches the source once it has reconnected and its + // resubscribe has been processed; QoS 1 publishes to a topic with no + // subscribers are simply dropped, so publish repeatedly until one lands + // rather than guessing at the resubscribe timing. + let mut received_after = None; + for _ in 0..20 { + producer + .publish(topic, QoS::AtLeastOnce, false, msg_after.as_bytes()) + .await + .unwrap(); + if let Ok(event) = timeout(Duration::from_millis(500), rx.next()).await { + received_after = event; + break; + } + } + let received_after = received_after + .expect("the source never resubscribed after the broker discarded its session"); + assert_eq!(message_body(&received_after), msg_after); + + source.abort(); + drop(source.await); +} + +/// The default mode (acknowledgements off, rumqttc auto-acks) must not lose +/// messages across a server-side disconnect either: in this mode a buffered +/// publish can already be acknowledged to the broker before the source +/// processes it, so anything the disconnect path drops locally would be gone +/// for good (the broker won't redeliver what it saw acked). +#[tokio::test] +async fn mqtt_does_not_lose_messages_across_server_side_disconnect_without_acks() { + trace_init(); + + let topic = "source-kick-no-acks-test"; + let client_id = format!("sourceKickNoAcksTest{}", random_string(6)); + let num_messages = 10; + + let 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()), + ..MqttSourceConfig::default() + }; + + let (tx, mut rx) = SourceSender::new_test(); + 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; + + let producer = get_mqtt_client().await; + + let (batch_before, ..) = random_lines_with_stream(100, num_messages, None); + send_test_events(&producer, topic, &batch_before).await; + collect_expected_messages( + &mut rx, + batch_before.into_iter().collect(), + Duration::from_secs(5), + ) + .await; + + // Kick the source (session preserved) and publish the second batch while + // it is offline; the session queues it for delivery after reconnect. + kick_mqtt_client(&client_id, false).await; + let (batch_after, ..) = random_lines_with_stream(100, num_messages, None); + send_test_events(&producer, topic, &batch_after).await; + + collect_expected_messages( + &mut rx, + batch_after.into_iter().collect(), + Duration::from_secs(10), + ) + .await; + + 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 a9b3f5090c74f..f392d544ab566 100644 --- a/src/sources/mqtt/source.rs +++ b/src/sources/mqtt/source.rs @@ -1414,4 +1414,376 @@ mod tests { let actions = state.on_connack(true, true); assert!(!actions.resubscribe); } + + // ------------------------------------------------------------------ + // Exhaustive protocol-state model check. + // + // The tests above sample hand-picked transition sequences; review kept + // finding holes in the sequences nobody picked (a forced reconnect + // skipped inside the cooldown with nothing left to re-trigger it, SUBACK + // failures treated as confirmation, ...). This test explores EVERY + // reachable abstract state of the protocol machinery instead: + // `ProtocolState` + `ForcedReconnectSchedule` + the resubscribe retry + // timer, under every legal interleaving of connection, finalization and + // timer events, checking on every transition the invariants that past + // findings violated. + // + // The event glue in `Model::apply` mirrors `run()`'s select branches; + // when changing `run()`, change it to match. Timer state is abstracted + // to armed/not-armed (which timestamps hold never affects WHETHER a + // transition is legal, only when it happens), so the abstract state + // space is finite and small; the BFS runs to fixpoint. + // ------------------------------------------------------------------ + + /// Cap for state-space purposes; behavior only distinguishes + /// `generation == 0` from `> 0`, so higher counts add nothing. + const MODEL_MAX_GENERATION: u64 = 2; + + /// Enough to have one suppressed publish plus one still-in-flight. + const MODEL_MAX_INFLIGHT: u8 = 2; + + #[derive(Clone, Copy, Debug)] + enum ModelEvent { + /// `poll()` yielded ConnAck; `subscribe_queued` is the outcome of the + /// `try_subscribe` the handler issues when `actions.resubscribe`. + ConnAck { + session_present: bool, + subscribe_queued: bool, + }, + /// `poll()` yielded a Publish which was processed and registered + /// with the finalizer. + ReceivePublish, + /// The ack stream yielded a finalization for a current-generation + /// entry (stale-generation entries cannot be yielded: the finalizer + /// is recreated on every ConnAck and disconnect). + FinalizePublish { status: BatchStatus }, + /// `poll()` yielded a SUBACK for the outstanding (re)subscribe. + SubAck { all_granted: bool }, + /// `poll()` returned an error (includes failed reconnect attempts, + /// so it is legal in any state). + PollError, + /// The forced-reconnect timer fired; `disconnect_queued` is the + /// outcome of `try_disconnect`. + ForcedReconnectDue { disconnect_queued: bool }, + /// The escalation deadline fired: the broker never closed after our + /// `Disconnect`, so the connection is dropped locally. + EscalationDue, + /// The paced resubscribe retry timer fired. + ResubscribeRetryDue { subscribe_queued: bool }, + /// The top-of-loop resubscribe retry ran (it runs before every + /// select; modeling it as an optional event covers both orders). + LoopTopResubscribe { subscribe_queued: bool }, + } + + const MODEL_EVENTS: [ModelEvent; 15] = [ + ModelEvent::ConnAck { + session_present: false, + subscribe_queued: false, + }, + ModelEvent::ConnAck { + session_present: false, + subscribe_queued: true, + }, + ModelEvent::ConnAck { + session_present: true, + subscribe_queued: false, + }, + ModelEvent::ConnAck { + session_present: true, + subscribe_queued: true, + }, + ModelEvent::ReceivePublish, + ModelEvent::FinalizePublish { + status: BatchStatus::Delivered, + }, + ModelEvent::FinalizePublish { + status: BatchStatus::Rejected, + }, + ModelEvent::SubAck { all_granted: true }, + ModelEvent::SubAck { all_granted: false }, + ModelEvent::PollError, + ModelEvent::ForcedReconnectDue { + disconnect_queued: true, + }, + ModelEvent::ForcedReconnectDue { + disconnect_queued: false, + }, + ModelEvent::EscalationDue, + ModelEvent::ResubscribeRetryDue { + subscribe_queued: true, + }, + ModelEvent::LoopTopResubscribe { + subscribe_queued: true, + }, + ]; + + #[derive(Clone)] + struct Model { + state: ProtocolState, + forced: ForcedReconnectSchedule, + resubscribe_retry_scheduled: bool, + /// Unfinalized publishes registered with the current finalizer. + inflight: u8, + now: tokio::time::Instant, + } + + /// The abstraction the BFS dedups on: timer instants collapse to + /// armed/not-armed, the generation saturates. + type ModelKey = (bool, bool, bool, bool, u64, bool, bool, bool, bool, u8); + + impl Model { + fn initial() -> Self { + Self { + state: ProtocolState { + pending_resubscribe: true, + ..ProtocolState::default() + }, + forced: ForcedReconnectSchedule::default(), + resubscribe_retry_scheduled: false, + inflight: 0, + now: tokio::time::Instant::now(), + } + } + + fn key(&self) -> ModelKey { + ( + self.state.connected, + self.state.pending_resubscribe, + self.state.resubscribe_awaiting_suback, + self.state.ack_suppressed, + self.state.connection_generation.min(MODEL_MAX_GENERATION), + self.forced.due_at().is_some(), + self.forced.cooldown_until.is_some(), + self.forced.escalate_at().is_some(), + self.resubscribe_retry_scheduled, + self.inflight, + ) + } + + fn legal(&self, event: ModelEvent) -> bool { + match event { + // ConnAck is strictly the first event of a connection. + ModelEvent::ConnAck { .. } => !self.state.connected, + ModelEvent::ReceivePublish => { + self.state.connected && self.inflight < MODEL_MAX_INFLIGHT + } + ModelEvent::FinalizePublish { .. } => self.inflight > 0, + // A SUBACK only arrives for an outstanding SUBSCRIBE. + ModelEvent::SubAck { .. } => { + self.state.connected && self.state.resubscribe_awaiting_suback + } + ModelEvent::PollError => true, + ModelEvent::ForcedReconnectDue { .. } => { + self.forced.due_at().is_some() && self.state.connected + } + ModelEvent::EscalationDue => self.forced.escalate_at().is_some(), + ModelEvent::ResubscribeRetryDue { .. } => self.resubscribe_retry_scheduled, + ModelEvent::LoopTopResubscribe { .. } => { + self.state.loop_actions().retry_resubscribe + && !self.forced.disconnect_sent() + } + } + } + + /// Mirrors the corresponding handler glue in `run()`. + fn apply(&mut self, event: ModelEvent, acknowledgements: bool) { + self.now += Duration::from_secs(1); + match event { + ModelEvent::ConnAck { + session_present, + subscribe_queued, + } => { + let actions = self.state.on_connack(acknowledgements, session_present); + // recreate_finalizer drops all registered entries. + assert!(actions.recreate_finalizer); + self.inflight = 0; + if actions.resubscribe { + self.state.on_resubscribe_result(subscribe_queued); + } + self.resubscribe_retry_scheduled = false; + } + ModelEvent::ReceivePublish => self.inflight += 1, + ModelEvent::FinalizePublish { status } => { + self.inflight -= 1; + let decision = self + .state + .finalize_publish(status, self.state.connection_generation); + if decision.should_ack { + // Emitting an ack while a `Disconnect` is unanswered + // would violate [MQTT-3.14.4-1]; `disconnect_sent` + // implies the generation is suppressed, which implies + // `should_ack` is false. + assert!( + !self.forced.disconnect_sent(), + "ack emitted while draining after a Disconnect" + ); + } + if decision.just_suppressed { + self.forced.schedule(self.now); + } + } + ModelEvent::SubAck { all_granted } => { + if all_granted { + self.state.on_suback(); + self.resubscribe_retry_scheduled = false; + } else { + self.resubscribe_retry_scheduled = true; + } + } + ModelEvent::PollError => self.connection_lost(), + ModelEvent::ForcedReconnectDue { disconnect_queued } => { + // Firing the timer while a previous `Disconnect` is + // still unanswered would send a second one. + assert!( + !self.forced.disconnect_sent(), + "forced-reconnect timer fired while a Disconnect is already unanswered" + ); + if disconnect_queued { + self.forced.sent(self.now); + } else { + self.forced.retry(self.now); + } + } + ModelEvent::EscalationDue => { + self.forced.escalated(); + // `connection.clean()` + shared disconnect handling. + self.connection_lost(); + } + ModelEvent::ResubscribeRetryDue { subscribe_queued } => { + self.resubscribe_retry_scheduled = false; + if self.state.connected + && self.state.pending_resubscribe + && !self.forced.disconnect_sent() + { + self.state.on_resubscribe_result(subscribe_queued); + } + } + ModelEvent::LoopTopResubscribe { subscribe_queued } => { + self.state.on_resubscribe_result(subscribe_queued); + } + } + } + + /// Mirrors `handle_connection_lost`. + fn connection_lost(&mut self) { + let actions = self.state.on_disconnect(); + assert!(actions.recreate_finalizer); + self.inflight = 0; + self.forced.cancel(); + self.resubscribe_retry_scheduled = false; + } + + fn check_invariants(&self, trace: &[ModelEvent]) { + let describe = || { + format!( + "state: {:?}, forced: {:?}, retry_scheduled: {}, inflight: {}\ntrace: {trace:?}", + self.state, self.forced, self.resubscribe_retry_scheduled, self.inflight, + ) + }; + + // The invariant behind the cooldown-skip finding: a suppressed, + // still-connected generation must always have a recovery action + // armed (the reconnect timer or the escalation deadline), or the + // source is stuck forever with nothing left to trigger recovery. + if self.state.connected && self.state.ack_suppressed { + assert!( + self.forced.due_at().is_some() || self.forced.escalate_at().is_some(), + "suppressed generation with no recovery action armed\n{}", + describe() + ); + } + + // Per-connection schedules must die with the connection: a timer + // surviving a disconnect would fire against the next connection. + if !self.state.connected { + assert!( + self.forced.due_at().is_none() + && self.forced.escalate_at().is_none() + && !self.resubscribe_retry_scheduled, + "per-connection schedule survived a disconnect\n{}", + describe() + ); + } + + // The invariant behind the SUBACK findings: while connected with + // an unconfirmed subscription, some path toward retrying it must + // remain live (the loop-top retry, an outstanding SUBACK that + // will either confirm or reject-and-pace, or the paced retry + // timer). While draining, the pending flag itself survives to + // the next connection, whose ConnAck re-issues the subscribe. + if self.state.connected + && self.state.pending_resubscribe + && !self.forced.disconnect_sent() + { + assert!( + self.state.loop_actions().retry_resubscribe + || self.state.resubscribe_awaiting_suback + || self.resubscribe_retry_scheduled, + "unconfirmed subscription with no retry path armed\n{}", + describe() + ); + } + } + } + + #[test] + fn exhaustive_protocol_state_model_check() { + for acknowledgements in [true, false] { + let mut visited = std::collections::HashSet::::new(); + // Keep a short trace per state for debuggable failures. + let mut queue = std::collections::VecDeque::<(Model, Vec)>::new(); + + let initial = Model::initial(); + initial.check_invariants(&[]); + visited.insert(initial.key()); + queue.push_back((initial, Vec::new())); + + let mut transitions = 0usize; + let mut reached_suppressed = false; + let mut reached_draining = false; + let mut reached_retry_pacing = false; + + while let Some((model, trace)) = queue.pop_front() { + for event in MODEL_EVENTS { + if !model.legal(event) { + continue; + } + let mut next = model.clone(); + next.apply(event, acknowledgements); + transitions += 1; + + let mut next_trace = trace.clone(); + next_trace.push(event); + next.check_invariants(&next_trace); + + // Transition-level invariant: `pending_resubscribe` may + // only be cleared by a fully granted SUBACK. + if model.state.pending_resubscribe && !next.state.pending_resubscribe { + assert!( + matches!(event, ModelEvent::SubAck { all_granted: true }), + "pending_resubscribe cleared by {event:?}, not a granted SUBACK\ntrace: {next_trace:?}" + ); + } + + reached_suppressed |= next.state.ack_suppressed; + reached_draining |= next.forced.disconnect_sent(); + reached_retry_pacing |= next.resubscribe_retry_scheduled; + + if visited.insert(next.key()) { + queue.push_back((next, next_trace)); + } + } + } + + // Guard against the model check passing vacuously. + assert!(visited.len() > 50, "suspiciously small state space"); + assert!(transitions > 200, "suspiciously few transitions"); + assert!(reached_suppressed, "never reached a suppressed state"); + assert!(reached_draining, "never reached a draining state"); + assert!( + reached_retry_pacing, + "never reached a paced-resubscribe-retry state" + ); + } + } }