From 6f390750a2e6a3bb3ae6de8c636267cc2c6185a3 Mon Sep 17 00:00:00 2001 From: colin Date: Fri, 19 Jun 2026 11:23:24 -0700 Subject: [PATCH 1/2] 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 5460a5975ed28bd16b124af1d1bd788545110bf0 Mon Sep 17 00:00:00 2001 From: colin Date: Wed, 1 Jul 2026 13:42:06 -0700 Subject: [PATCH 2/2] 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 + ); + } } }