ygndotgg commented on code in PR #4303: URL: https://github.com/apache/iggy/pull/4303#discussion_r4162930401
########## core/connectors/sources/mqtt_source/src/driver.rs: ########## @@ -0,0 +1,1226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::{MqttProtocol, MqttSourceConfig, Qos, qos_for_subscription}; +use base64::Engine; +use iggy_common::{HeaderKey, HeaderValue}; +use rumqttc::tokio_rustls::rustls::{ + self, ClientConfig, RootCertStore, + pki_types::{CertificateDer, PrivateKeyDer}, +}; +use rumqttc::v5::{ + AsyncClient as Mqtt5Client, Event as Mqtt5Event, EventLoop as Mqtt5EventLoop, + Incoming as Mqtt5Incoming, MqttOptions as Mqtt5Options, + mqttbytes::{QoS as Mqtt5Qos, v5::Publish as Mqtt5Publish}, +}; +use rumqttc::{ + AsyncClient as Mqtt311Client, Event as Mqtt311Event, EventLoop as Mqtt311EventLoop, + Incoming as Mqtt311Incoming, MqttOptions as Mqtt311Options, TlsConfiguration, Transport, + mqttbytes::{QoS as Mqtt311Qos, v4::Publish as Mqtt311Publish}, +}; +use rustls_native_certs::load_native_certs; +use rustls_pemfile::{certs, private_key}; +use secrecy::ExposeSecret; +use serde::Serialize; +use std::{ + collections::{BTreeMap, VecDeque}, + io::{BufReader, Cursor}, + sync::{Arc, LazyLock}, + time::Duration, +}; +use tokio::time::timeout; +use tracing::warn; +use url::Url; + +const ACK_RETRY_DELAY: Duration = Duration::from_millis(10); +const MAX_HEADER_VALUE_LENGTH: usize = 255; +const MQTT_PROTOCOL_HEADER: &str = "mqtt.protocol"; +const MQTT_TOPIC_HEADER: &str = "mqtt.topic"; + +static MQTT_PROTOCOL_HEADER_KEY: LazyLock<HeaderKey> = LazyLock::new(|| { + HeaderKey::try_from(MQTT_PROTOCOL_HEADER).expect("MQTT protocol header key is valid") +}); +static MQTT_TOPIC_HEADER_KEY: LazyLock<HeaderKey> = LazyLock::new(|| { + HeaderKey::try_from(MQTT_TOPIC_HEADER).expect("MQTT topic header key is valid") +}); +static MQTT_QOS_HEADER_KEY: LazyLock<HeaderKey> = + LazyLock::new(|| HeaderKey::try_from("mqtt.qos").expect("MQTT QoS header key is valid")); +static MQTT_DUP_HEADER_KEY: LazyLock<HeaderKey> = + LazyLock::new(|| HeaderKey::try_from("mqtt.dup").expect("MQTT duplicate header key is valid")); +static MQTT_RETAIN_HEADER_KEY: LazyLock<HeaderKey> = + LazyLock::new(|| HeaderKey::try_from("mqtt.retain").expect("MQTT retain header key is valid")); +#[cfg(test)] +static MQTT_PACKET_ID_HEADER_KEY: LazyLock<HeaderKey> = LazyLock::new(|| { + HeaderKey::try_from("mqtt.packet_id").expect("MQTT packet ID header key is valid") +}); +#[cfg(test)] +static MQTT_RESPONSE_TOPIC_HEADER_KEY: LazyLock<HeaderKey> = LazyLock::new(|| { + HeaderKey::try_from("mqtt.response_topic").expect("MQTT response topic header key is valid") +}); + +// Metadata is kept separate from the payload so protocol details can be +// preserved in Iggy headers without changing the application bytes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct MqttMessageMetadata { + pub(crate) qos: Qos, + pub(crate) packet_id: Option<u16>, + pub(crate) dup: bool, + pub(crate) retain: bool, +} + +/// MQTT payload plus normalized topic metadata ready for Iggy headers. +#[derive(Debug)] +pub(crate) struct MqttMessage { + pub(crate) payload: Vec<u8>, + pub(crate) topic: String, + pub(crate) headers: BTreeMap<HeaderKey, HeaderValue>, + pub(crate) metadata: MqttMessageMetadata, + pub(crate) mqtt5_properties: Option<Mqtt5EnvelopeProperties>, +} + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct Mqtt5EnvelopeProperties { + pub(crate) payload_format_indicator: Option<u8>, + pub(crate) message_expiry_interval: Option<u32>, + pub(crate) topic_alias: Option<u16>, + pub(crate) response_topic: Option<String>, + pub(crate) correlation_data_base64: Option<String>, + pub(crate) user_properties: Vec<Mqtt5UserProperty>, + pub(crate) subscription_identifiers: Vec<usize>, + pub(crate) content_type: Option<String>, +} + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct Mqtt5UserProperty { + pub(crate) key: String, + pub(crate) value: String, +} + +/// A normalized MQTT message and its optional deferred acknowledgement token. +#[derive(Debug)] +pub(crate) struct ReceivedMessage { + pub(crate) message: MqttMessage, + pub(crate) ack_token: Option<AckToken>, +} + +/// The original publish packet needed by rumqttc to acknowledge QoS 1 or QoS 2. +#[derive(Debug)] +pub(crate) struct AckToken(AckTokenKind); + +#[derive(Debug)] +enum AckTokenKind { + Mqtt311(Mqtt311Publish), + Mqtt5(Mqtt5Publish), +} + +enum MqttConnection { + Mqtt311 { + client: Box<Mqtt311Client>, + event_loop: Box<Mqtt311EventLoop>, + }, + Mqtt5 { + client: Box<Mqtt5Client>, + event_loop: Box<Mqtt5EventLoop>, + }, +} + +/// Protocol-specific MQTT client and event loop behind one common driver API. +pub(crate) struct MqttDriver { + // Messages received while flushing acknowledgements are retained here rather + // than dropped. The source batch size bounds this queue. + connection: MqttConnection, + buffered_messages: VecDeque<ReceivedMessage>, +} + +impl MqttDriver { + pub(crate) async fn connect( + id: u32, + config: &MqttSourceConfig, + qos: Qos, + keep_alive: Duration, + poll_timeout: Duration, + request_capacity: usize, + ) -> Result<Self, iggy_connector_sdk::Error> { + install_rustls_provider(); + let broker_url = broker_url_with_client_id(config, id)?; + match config.protocol { + MqttProtocol::Mqtt311 => { + let mut options = Mqtt311Options::parse_url(&broker_url).map_err(|error| { + iggy_connector_sdk::Error::InvalidConfigValue(format!("broker_url: {error}")) + })?; + options + .set_client_id(client_id(config, id)) + .set_clean_session(config.clean_start) + .set_keep_alive(keep_alive) + .set_request_channel_capacity(request_capacity) + // Manual acknowledgements let Iggy persistence happen before + // PUBACK/PUBREC is sent to the broker. + .set_manual_acks(true); + set_mqtt311_credentials(&mut options, config); + if let Some(transport) = tls_transport(config)? { + options.set_transport(transport); + } + + let (client, mut event_loop) = Mqtt311Client::new(options, request_capacity); + for topic in &config.subscriptions { + // Each filter may select its own QoS override; the global + // value is used only when no exact override is configured. + let subscription_qos = qos_for_subscription(config, topic, qos)?; + client + .subscribe(topic, subscription_qos.into()) + .await + .map_err(|error| { + iggy_connector_sdk::Error::Connection(error.to_string()) + })?; + } + let buffered_messages = + poll_mqtt311(&mut event_loop, config.subscriptions.len(), poll_timeout).await?; + + Ok(Self { + connection: MqttConnection::Mqtt311 { + client: Box::new(client), + event_loop: Box::new(event_loop), + }, + buffered_messages, + }) + } + MqttProtocol::Mqtt5 => { + let mut options = Mqtt5Options::parse_url(&broker_url).map_err(|error| { + iggy_connector_sdk::Error::InvalidConfigValue(format!("broker_url: {error}")) + })?; + options + .set_client_id(client_id(config, id)) + .set_clean_start(config.clean_start) + .set_keep_alive(keep_alive) + .set_request_channel_capacity(request_capacity) + // MQTT 5 uses the same deferred-acknowledgement lifecycle; + // rumqttc maps the token to PUBREC/PUBREL/PUBCOMP internally. + .set_manual_acks(true) + .set_session_expiry_interval(config.session_expiry_interval); + set_mqtt5_credentials(&mut options, config); + if let Some(transport) = tls_transport(config)? { + options.set_transport(transport); + } + + let (client, mut event_loop) = Mqtt5Client::new(options, request_capacity); + for topic in &config.subscriptions { + // Keep subscription QoS resolution identical across protocol + // versions so the configuration has one predictable meaning. + let subscription_qos = qos_for_subscription(config, topic, qos)?; + client + .subscribe(topic, subscription_qos.into()) + .await + .map_err(|error| { + iggy_connector_sdk::Error::Connection(error.to_string()) + })?; + } + let buffered_messages = + poll_mqtt5(&mut event_loop, config.subscriptions.len(), poll_timeout).await?; + + Ok(Self { + connection: MqttConnection::Mqtt5 { + client: Box::new(client), + event_loop: Box::new(event_loop), + }, + buffered_messages, + }) + } + } + } + + pub(crate) async fn next_message( + &mut self, + poll_timeout: Duration, + ) -> Result<Option<ReceivedMessage>, iggy_connector_sdk::Error> { + // Consume buffered messages first. They arrived while another message's + // acknowledgement was being flushed and are already valid source data. + if let Some(message) = self.buffered_messages.pop_front() { + return Ok(Some(message)); + } + + // Bound each driver poll so the source can flush partial batches and + // respond to runtime shutdown instead of waiting indefinitely. + let deadline = tokio::time::Instant::now() + poll_timeout; + loop { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Ok(None); + } + let event = match &mut self.connection { + MqttConnection::Mqtt311 { event_loop, .. } => { + match timeout(remaining, event_loop.poll()).await { + Ok(Ok(event)) => match event { + Mqtt311Event::Incoming(Mqtt311Incoming::Publish(publish)) => { + Some(normalize_mqtt311(publish)?) + } + // Connection and subscription events advance the + // event loop but are not source messages. + Mqtt311Event::Outgoing(_) | Mqtt311Event::Incoming(_) => None, + }, + Ok(Err(error)) => { + return Err(iggy_connector_sdk::Error::Connection(error.to_string())); + } + Err(_) => return Ok(None), + } + } + MqttConnection::Mqtt5 { event_loop, .. } => { + match timeout(remaining, event_loop.poll()).await { + Ok(Ok(event)) => match event { + Mqtt5Event::Incoming(Mqtt5Incoming::Publish(publish)) => { + Some(normalize_mqtt5(publish)?) + } + Mqtt5Event::Outgoing(_) | Mqtt5Event::Incoming(_) => None, + }, + Ok(Err(error)) => { + return Err(iggy_connector_sdk::Error::Connection(error.to_string())); + } + Err(_) => return Ok(None), + } + } + }; + if event.is_some() { + return Ok(event); + } + } + } + + pub(crate) async fn acknowledge_batch( + &mut self, + ack_tokens: &mut Vec<AckToken>, + poll_timeout: Duration, + max_buffered_messages: usize, + max_retries: u32, + ) -> Result<(), iggy_connector_sdk::Error> { + // Process tokens in order. On failure, remove only tokens already + // acknowledged so the remaining suffix can be retried. + let mut acknowledged = 0; + let retry_delay = ACK_RETRY_DELAY.min(poll_timeout); + while acknowledged < ack_tokens.len() { + let mut last_error = None; + let mut acknowledged_token = false; + for attempt in 0..=max_retries { + match self.try_acknowledge(&ack_tokens[acknowledged]) { + Ok(()) => { + acknowledged_token = true; + break; + } + Err(error) => { + last_error = Some(error); + if attempt == max_retries { + break; + } + // rumqttc may need event-loop progress before try_ack can + // enqueue the acknowledgement. + if let Err(error) = self + .poll_for_ack_progress(poll_timeout, max_buffered_messages) + .await + { + retain_unacknowledged_tokens(ack_tokens, acknowledged); + return Err(error); + } + if !retry_delay.is_zero() { + tokio::time::sleep(retry_delay).await; + } + } + } + } + + if !acknowledged_token { + retain_unacknowledged_tokens(ack_tokens, acknowledged); + return Err(last_error.unwrap_or(iggy_connector_sdk::Error::InvalidState)); + } + acknowledged += 1; + } + ack_tokens.clear(); + Ok(()) + } + + fn try_acknowledge(&self, ack_token: &AckToken) -> Result<(), iggy_connector_sdk::Error> { + // A token must be acknowledged by the same protocol client that created + // it. Mixing MQTT 3.1.1 and MQTT 5 tokens is an invalid internal state. + match (&self.connection, &ack_token.0) { + (MqttConnection::Mqtt311 { client, .. }, AckTokenKind::Mqtt311(publish)) => client + .try_ack(publish) + .map_err(|error| iggy_connector_sdk::Error::Connection(error.to_string())), + (MqttConnection::Mqtt5 { client, .. }, AckTokenKind::Mqtt5(publish)) => client + .try_ack(publish) + .map_err(|error| iggy_connector_sdk::Error::Connection(error.to_string())), + _ => Err(iggy_connector_sdk::Error::InvalidState), + } + } + + async fn poll_for_ack_progress( + &mut self, + poll_timeout: Duration, + max_buffered_messages: usize, + ) -> Result<(), iggy_connector_sdk::Error> { + // A publish can arrive while the event loop is being driven for an ACK. + // Preserve it for the next source poll, subject to the batch bound. + let received = match &mut self.connection { + MqttConnection::Mqtt311 { event_loop, .. } => { + let event = timeout(poll_timeout, event_loop.poll()) + .await + .map_err(|_| { + iggy_connector_sdk::Error::Connection( + "timed out while flushing MQTT acknowledgements".to_string(), + ) + })? + .map_err(|error| iggy_connector_sdk::Error::Connection(error.to_string()))?; + match event { + Mqtt311Event::Incoming(Mqtt311Incoming::Publish(publish)) => { + Some(normalize_mqtt311(publish)?) + } + Mqtt311Event::Outgoing(_) | Mqtt311Event::Incoming(_) => None, + } + } + MqttConnection::Mqtt5 { event_loop, .. } => { + let event = timeout(poll_timeout, event_loop.poll()) + .await + .map_err(|_| { + iggy_connector_sdk::Error::Connection( + "timed out while flushing MQTT acknowledgements".to_string(), + ) + })? + .map_err(|error| iggy_connector_sdk::Error::Connection(error.to_string()))?; + match event { + Mqtt5Event::Incoming(Mqtt5Incoming::Publish(publish)) => { + Some(normalize_mqtt5(publish)?) + } + Mqtt5Event::Outgoing(_) | Mqtt5Event::Incoming(_) => None, + } + } + }; + + if let Some(received) = received { Review Comment: Fixed by Preserving MQTT publishes in the buffer before reporting capacity errors to prevent message loss during acknowledgement flushing. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
