enhancement(mqtt source): support end-to-end acknowledgements

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.

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)]

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<String> {
@@ -77,14 +81,22 @@ impl SourceConfig for MqttSourceConfig {
     async fn build(&self, cx: SourceContext) -> crate::Result<crate::sources::Source> {
         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<SourceOutput> {
@@ -107,12 +119,22 @@ impl SourceConfig for MqttSourceConfig {
     }
 
     fn can_acknowledge(&self) -> bool {
-        false
+        true
     }
 }
 
 impl MqttSourceConfig {
-    fn build_connector(&self) -> Result<MqttConnector, MqttError> {
+    fn build_connector(&self, acknowledgements: bool) -> Result<MqttConnector, MqttError> {
+        // 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::<MqttSourceConfig>();
     }
+
+    #[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());
+    }
 }

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,13 +40,26 @@ async fn send_test_events(client: &AsyncClient, topic: &str, messages: &Vec<Stri
     }
 }
 
-async fn get_mqtt_client() -> AsyncClient {
+fn message_body(event: &Event) -> String {
+    event
+        .as_log()
+        .get(log_schema().message_key_target_path().unwrap())
+        .unwrap()
+        .to_string_lossy()
+        .into_owned()
+}
+
+async fn get_mqtt_client_with_options(configure: impl FnOnce(&mut MqttOptions)) -> AsyncClient {
+    // Unique client ID per producer: brokers that strictly enforce client-ID
+    // uniqueness (e.g. RabbitMQ) otherwise kick a previous connection when tests
+    // run concurrently, which manifests as spurious publish timeouts.
     let mut mqtt_options = MqttOptions::new(
-        "integration-test-producer",
+        format!("integration-test-producer-{}", random_string(6)),
         mqtt_broker_address(),
         mqtt_broker_port(),
     );
     mqtt_options.set_keep_alive(Duration::from_secs(5));
+    configure(&mut mqtt_options);
 
     let (client, mut eventloop) = AsyncClient::new(mqtt_options, 10);
 
@@ -59,6 +72,10 @@ async fn get_mqtt_client() -> AsyncClient {
     client
 }
 
+async fn get_mqtt_client() -> AsyncClient {
+    get_mqtt_client_with_options(|_| {}).await
+}
+
 #[tokio::test]
 async fn mqtt_one_topic_happy() {
     trace_init();
@@ -118,6 +135,197 @@ 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);
+}
+
+/// A downstream that can't keep up must not deadlock the source or grow memory
+/// without bound. `poll_mqtt_connection` forwards polled events to the main task
+/// over a channel bounded by `EVENTS_CHANNEL_CAPACITY`; once that backlog
+/// saturates, the event is dropped rather than buffered without bound or blocking
+/// (blocking would risk deadlocking against rumqttc's own outgoing request
+/// channel, since the poller is also the only thing that drains it). This test
+/// proves that path actually recovers instead of hanging, and that every message
+/// is still eventually delivered afterward: a dropped QoS 1 publish was never
+/// acknowledged, so the broker's own retry timer redelivers it (see
+/// `compose.yaml`'s shortened `retry_interval`) independently of whatever else
+/// is still flowing through.
+#[tokio::test]
+async fn mqtt_recovers_from_downstream_saturation() {
+    trace_init();
+
+    let topic = format!("source-saturation-test-{}", random_string(6));
+    let client_id = format!("sourceSaturationTest{}", random_string(6));
+    // Comfortably more than the poller's internal event backlog capacity, so the
+    // backlog saturates and starts dropping messages before the sink below starts
+    // draining.
+    let num_events = super::source::EVENTS_CHANNEL_CAPACITY * 3;
+    let (input, ..) = random_lines_with_stream(16, num_events, None);
+
+    let config = MqttSourceConfig {
+        common: MqttCommonConfig {
+            host: mqtt_broker_address(),
+            port: mqtt_broker_port(),
+            client_id: Some(client_id),
+            ..Default::default()
+        },
+        topic: OneOrMany::One(topic.clone()),
+        acknowledgements: true.into(),
+        ..MqttSourceConfig::default()
+    };
+
+    // A small, undrained buffer stands in for a stalled/slow sink: `out.send_batch`
+    // inside the source blocks almost immediately once it fills.
+    let (tx, rx) = SourceSender::new_test_finalize(EventStatus::Delivered);
+    let source = tokio::spawn(async move {
+        config
+            .build(SourceContext::new_test(tx, None))
+            .await
+            .unwrap()
+            .await
+            .unwrap()
+    });
+
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // rumqttc's default outgoing QoS-1 inflight window (100) would otherwise pace
+    // the producer to the broker's PUBACK round-trip time, making it too slow to
+    // ever build up a backlog at the subscriber. Raise it so publishing is only
+    // limited by the network, not by our own producer's flow control.
+    let producer = get_mqtt_client_with_options(|opts| {
+        opts.set_inflight(u16::MAX);
+    })
+    .await;
+    send_test_events(&producer, &topic, &input).await;
+
+    // Give the poller time to saturate its backlog and start dropping messages
+    // while nothing is draining the sink.
+    tokio::time::sleep(Duration::from_secs(2)).await;
+
+    // Now let the sink drain. Duplicates are expected (a dropped message is
+    // redelivered by the broker's retry timer, independently of whatever else is
+    // still flowing through), so drain until every distinct message has been seen
+    // rather than taking a fixed count. If the poller had deadlocked instead of
+    // recovering, this times out.
+    let mut expected_messages: HashSet<_> = input.into_iter().collect();
+    let mut total_received = 0usize;
+    let drain = async {
+        tokio::pin!(rx);
+        while !expected_messages.is_empty() {
+            let event = rx.next().await.expect("source stream ended unexpectedly");
+            expected_messages.remove(&message_body(&event));
+            total_received += 1;
+        }
+    };
+    timeout(Duration::from_secs(60), drain).await.expect(
+        "timed out waiting for all messages after downstream recovered: the source likely deadlocked",
+    );
+
+    // If nothing was ever actually dropped (e.g. a regression reverted the
+    // backlog channel to unbounded), every message would still eventually
+    // arrive via simple buffering and the check above would pass for the wrong
+    // reason. Requiring more received events than distinct inputs pins the
+    // pass on the drop-and-redeliver path having actually run.
+    assert!(
+        total_received > num_events,
+        "expected at least one duplicate from broker redelivery of a dropped message, \
+         got {total_received} events for {num_events} distinct inputs -- the backlog \
+         may never have actually saturated"
+    );
+
+    source.abort();
+    drop(source.await);
+}
+
 #[tokio::test]
 async fn mqtt_many_topics_happy() {
     trace_init();