ygndotgg commented on code in PR #4303: URL: https://github.com/apache/iggy/pull/4303#discussion_r4269389590
########## core/connectors/sources/mqtt_source/src/lib.rs: ########## @@ -0,0 +1,1227 @@ +// 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. + +mod driver; + +use async_trait::async_trait; +use base64::Engine; +use humantime::Duration as HumanDuration; +use iggy_common::{HeaderKey, HeaderValue}; +use iggy_connector_sdk::{ + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + source::SourceBatchResult, source_connector, +}; +use secrecy::SecretString; +use serde::{Deserialize, Serialize}; +use std::{ + collections::{BTreeMap, HashSet}, + fmt, + str::FromStr, + time::Duration, +}; +use tokio::{sync::Mutex, time::sleep}; +use tracing::{debug, error, info, warn}; +use url::Url; + +use driver::{AckToken, MqttDriver}; +use driver::{Mqtt5EnvelopeProperties, MqttMessage}; + +source_connector!(MqttSource); + +const CONNECTOR_NAME: &str = "MQTT source"; +const DEFAULT_KEEP_ALIVE: &str = "30s"; +const DEFAULT_POLL_TIMEOUT: &str = "1s"; +const DEFAULT_REQUEST_CAPACITY: usize = 32; +const DEFAULT_BATCH_SIZE: usize = 100; +const DEFAULT_BATCH_TIMEOUT: &str = "10ms"; +const DEFAULT_MAX_RETRIES: u32 = 5; +const DEFAULT_SESSION_EXPIRY_INTERVAL: u32 = 3600; + +/// MQTT wire protocol selected for the broker connection. +#[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum MqttProtocol { + #[serde(rename = "mqtt311")] + Mqtt311, + #[serde(rename = "mqtt5")] + #[default] + Mqtt5, +} + +/// Subscription delivery level requested from the broker. +/// +/// QoS 0 has no acknowledgement token. QoS 1 and QoS 2 produce a token that +/// remains pending until the corresponding Iggy batch is acknowledged. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Qos { + Zero, + One, + Two, +} + +impl TryFrom<u8> for Qos { + type Error = Error; + + fn try_from(value: u8) -> Result<Self, Self::Error> { + match value { + 0 => Ok(Self::Zero), + 1 => Ok(Self::One), + 2 => Ok(Self::Two), + value => Err(Error::InvalidConfigValue(format!( + "qos {value} is unsupported; expected 0, 1, or 2" + ))), + } + } +} + +#[derive(Debug, Deserialize)] +pub struct MqttSourceConfig { + /// MQTT broker URL, using `mqtt://`, `mqtts://`, or `ssl://`. + pub broker_url: String, + /// Topic filters subscribed to by this connector instance. + pub subscriptions: Vec<String>, + /// Exact topic-filter QoS overrides applied on top of `qos`. + #[serde(default)] + pub subscription_qos: BTreeMap<String, u8>, + /// MQTT protocol version used for the connection. + #[serde(default)] + pub protocol: MqttProtocol, + /// Preserves MQTT 5 extended properties in a JSON payload envelope. + #[serde(default)] + pub include_metadata: bool, + /// Default subscription QoS when no per-filter override exists. + #[serde(default = "default_qos")] + pub qos: u8, + /// Optional CA and client-authentication material for TLS connections. + #[serde(default)] + pub tls: Option<MqttTlsConfig>, + /// Explicit broker client ID. A connector-specific ID is generated when absent. + pub client_id: Option<String>, + /// Broker username, which must be paired with `password`. + pub username: Option<String>, + /// Broker password, kept secret in memory and never included in debug output. + pub password: Option<SecretString>, + /// Whether the broker should discard the previous session on connect. + #[serde(default)] + pub clean_start: bool, + /// MQTT 5 session expiry interval in seconds. + #[serde(default = "default_session_expiry_interval")] + pub session_expiry_interval: Option<u32>, + /// MQTT keep-alive interval. + pub keep_alive: Option<String>, + /// Maximum time to wait for a message during batch collection. + pub poll_timeout: Option<String>, + /// Capacity of rumqttc's request channel. + pub request_capacity: Option<usize>, + /// Maximum number of messages in one Iggy source batch. + pub batch_size: Option<usize>, + /// Maximum time to wait after the first message before flushing a batch. + pub batch_timeout: Option<String>, + /// Number of retries for a broker acknowledgement after the initial attempt. + /// After retries are exhausted, the source abandons the broker acknowledgement + /// and continues. The broker may redeliver the message later, producing duplicates. + #[serde(default)] + pub max_retries: Option<u32>, + /// Enables per-message MQTT logging for troubleshooting. + pub verbose_logging: Option<bool>, +} + +#[derive(Debug, Deserialize)] +pub struct MqttTlsConfig { + /// Optional custom CA bundle. System roots are used when this is absent. + pub ca_file: Option<String>, + /// Client certificate for mutual TLS, paired with `client_key_file`. + pub client_cert_file: Option<String>, + /// Client private key for mutual TLS, paired with `client_cert_file`. + pub client_key_file: Option<String>, +} + +pub struct MqttSource { + id: u32, + config: MqttSourceConfig, + poll_timeout: Duration, + batch_size: usize, + batch_timeout: Duration, + // The driver is moved out briefly during event-loop I/O to avoid holding a + // mutex guard across await points. + driver: Mutex<Option<MqttDriver>>, + // At most one batch is staged because the SDK reports one batch result at a + // time. It holds the MQTT acknowledgement tokens for that batch. + pending_batch: Mutex<Option<PendingBatch>>, +} + +impl fmt::Debug for MqttSource { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("MqttSource") + .field("id", &self.id) + .field("config", &self.config) + .field("poll_timeout", &self.poll_timeout) + .finish_non_exhaustive() + } +} + +#[derive(Debug)] +struct PendingBatch { + replay_messages: Vec<ReplayMessage>, + schema: Schema, + // Tokens are kept separately because the runtime acknowledges Iggy before + // the driver acknowledges MQTT. + ack_tokens: Vec<AckToken>, + // The SDK permits one batch in flight. A NACK leaves this false so poll() + // can replay the batch before reading newer MQTT messages. + in_flight: bool, +} + +#[derive(Debug, Clone)] +struct ReplayMessage { + headers: Option<BTreeMap<HeaderKey, HeaderValue>>, + payload: Vec<u8>, +} + +impl PendingBatch { + fn replay(&self) -> ProducedMessages { + ProducedMessages { + schema: self.schema, + messages: self + .replay_messages + .iter() + .cloned() + .map(|message| ProducedMessage { + id: None, + checksum: None, + timestamp: None, + origin_timestamp: None, + headers: message.headers, + payload: message.payload, + }) + .collect(), + state: None, + } + } +} + +impl MqttSource { + pub fn new(id: u32, config: MqttSourceConfig, _state: Option<ConnectorState>) -> Self { + Self { + id, + config, + poll_timeout: Duration::from_secs(1), + batch_size: DEFAULT_BATCH_SIZE, + batch_timeout: Duration::from_millis(10), + driver: Mutex::new(None), + pending_batch: Mutex::new(None), + } + } + + fn validate_config(&self) -> Result<(Qos, Duration, Duration, usize, usize, Duration), Error> { + // Validate static values before opening the network connection so an + // operator sees configuration errors during initialization. + if self.config.broker_url.trim().is_empty() { + return Err(Error::InvalidConfigValue( + "broker_url must not be empty".to_string(), + )); + } + if self.config.include_metadata && self.config.protocol == MqttProtocol::Mqtt311 { + return Err(Error::InvalidConfigValue( + "include_metadata is only supported with protocol = \"mqtt5\"".to_string(), + )); + } + if self.config.subscriptions.is_empty() + || self + .config + .subscriptions + .iter() + .any(|topic| topic.trim().is_empty()) + { + return Err(Error::InvalidConfigValue( + "subscriptions must contain at least one non-empty topic".to_string(), + )); + } + let mut unique_subscriptions = HashSet::with_capacity(self.config.subscriptions.len()); + for subscription in &self.config.subscriptions { + if !unique_subscriptions.insert(subscription) { + return Err(Error::InvalidConfigValue(format!( + "subscriptions contains duplicate topic filter: {subscription}" + ))); + } + } + // An override is keyed by the exact configured filter. An unknown key + // would otherwise be accepted but never used. + for subscription in self.config.subscription_qos.keys() { + if !self + .config + .subscriptions + .iter() + .any(|topic| topic == subscription) + { + return Err(Error::InvalidConfigValue(format!( + "subscription_qos contains unknown topic filter: {subscription}" + ))); + } + } + if self.config.username.is_some() != self.config.password.is_some() { + return Err(Error::InvalidConfigValue( + "username and password must be configured together".to_string(), + )); + } + // Certificate and hostname checks happen before rumqttc is constructed, + // so malformed TLS configuration cannot become a reconnect loop. + validate_tls_config(&self.config)?; + + let qos = Qos::try_from(self.config.qos)?; + for subscription in &self.config.subscriptions { + qos_for_subscription(&self.config, subscription, qos)?; + } + let keep_alive = parse_duration( + self.config.keep_alive.as_deref(), + DEFAULT_KEEP_ALIVE, + "keep_alive", + )?; + let poll_timeout = parse_duration( + self.config.poll_timeout.as_deref(), + DEFAULT_POLL_TIMEOUT, + "poll_timeout", + )?; + if poll_timeout.is_zero() { + return Err(Error::InvalidConfigValue( + "poll_timeout must be greater than zero".to_string(), + )); + } + let request_capacity = self + .config + .request_capacity + .unwrap_or(DEFAULT_REQUEST_CAPACITY); + if request_capacity == 0 { + return Err(Error::InvalidConfigValue( + "request_capacity must be greater than zero".to_string(), + )); + } + if request_capacity < self.config.subscriptions.len() { + return Err(Error::InvalidConfigValue(format!( + "request_capacity must be at least the number of subscriptions ({})", + self.config.subscriptions.len() + ))); + } + + // These bounds protect both the plugin-owned batch and the MQTT request + // channel from configurations that would otherwise never make progress. + let batch_size = self.config.batch_size.unwrap_or(DEFAULT_BATCH_SIZE); + if batch_size == 0 { + return Err(Error::InvalidConfigValue( + "batch_size must be greater than zero".to_string(), + )); + } + let batch_timeout = parse_duration( + self.config.batch_timeout.as_deref(), + DEFAULT_BATCH_TIMEOUT, + "batch_timeout", + )?; + + let minimum_keep_alive = match self.config.protocol { + MqttProtocol::Mqtt311 => Duration::from_secs(1), + MqttProtocol::Mqtt5 => Duration::from_secs(5), + }; + if !keep_alive.is_zero() && keep_alive < minimum_keep_alive { + return Err(Error::InvalidConfigValue(format!( + "keep_alive must be at least {:?} for {:?}", + minimum_keep_alive, self.config.protocol + ))); + } + + Ok(( + qos, + keep_alive, + poll_timeout, + request_capacity, + batch_size, + batch_timeout, + )) + } + + fn max_retries(&self) -> u32 { + self.config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES) + } + + async fn collect_batch( + &self, + driver: &mut MqttDriver, + ) -> Result<Vec<driver::ReceivedMessage>, Error> { + // Start the batch timeout only after the first message arrives. A quiet + // connector therefore remains cheap while partial batches still flush. + let Some(first) = driver.next_message(self.poll_timeout).await? else { + return Ok(Vec::new()); + }; + + let mut messages = Vec::new(); + messages.push(first); + let deadline = tokio::time::Instant::now() + self.batch_timeout; + while messages.len() < self.batch_size { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + break; + } + let Some(message) = driver.next_message(remaining).await? else { + break; + }; + messages.push(message); + } + Ok(messages) + } +} + +pub(crate) fn qos_for_subscription( + config: &MqttSourceConfig, + subscription: &str, + default_qos: Qos, +) -> Result<Qos, Error> { + // The global QoS is the fallback; an exact filter override wins when present. + match config.subscription_qos.get(subscription).copied() { + Some(value) => Qos::try_from(value), + None => Ok(default_qos), + } +} + +#[async_trait] +impl Source for MqttSource { + async fn open(&mut self) -> Result<(), Error> { + // Opening validates configuration and creates the broker session. Polling + // assumes this driver is ready. + let (qos, keep_alive, poll_timeout, request_capacity, batch_size, batch_timeout) = + match self.validate_config() { + Ok(values) => values, + Err(error) => { + error!( + "Failed to validate MQTT source connector with ID {}: {error}", + self.id + ); + return Err(error); + } + }; + let driver = match MqttDriver::connect( + self.id, + &self.config, + qos, + keep_alive, + poll_timeout, + request_capacity, + ) + .await + { + Ok(driver) => driver, + Err(error) => { + error!( + "Failed to connect MQTT source connector with ID {} to broker: {error}", + self.id + ); + return Err(error); + } + }; + self.poll_timeout = poll_timeout; + self.batch_size = batch_size; + self.batch_timeout = batch_timeout; + *self.driver.lock().await = Some(driver); + info!( + "Opened {CONNECTOR_NAME} connector with ID {} using {:?}, QoS {}, batch_size {}, batch_timeout {:?}, for {} subscription(s)", + self.id, + self.config.protocol, + self.config.qos, + self.batch_size, + self.batch_timeout, + self.config.subscriptions.len() + ); + Ok(()) + } + + async fn poll(&self) -> Result<ProducedMessages, Error> { + sleep(Duration::from_millis(10)).await; + // The SDK permits one source batch in flight. A second poll would lose + // the association between messages, candidate state, and ACK tokens. + { + let mut pending_batch = self.pending_batch.lock().await; + if let Some(pending_batch) = pending_batch.as_mut() { + if pending_batch.in_flight { + return Err(Error::InvalidState); + } + pending_batch.in_flight = true; + return Ok(pending_batch.replay()); + } + } + + let mut driver = self + .driver + .lock() + .await + .take() + .ok_or_else(|| Error::InitError("MQTT driver is not initialized".to_string()))?; + // Move the driver out while polling so network I/O does not occur while + // holding the source mutex. Restore it before returning. + let received = self.collect_batch(&mut driver).await; + *self.driver.lock().await = Some(driver); + let received = received?; + if received.is_empty() { + return Ok(empty_messages()); + } + + let mut messages = Vec::with_capacity(received.len()); + let mut replay_messages = Vec::with_capacity(received.len()); + let mut ack_tokens = Vec::with_capacity(received.len()); + for received in received { + if self.config.verbose_logging.unwrap_or(false) { + debug!( + "Received MQTT message for {CONNECTOR_NAME} connector with ID {}: topic={}, qos={:?}, packet_id={:?}, retain={}", + self.id, + received.message.topic, + received.message.metadata.qos, + received.message.metadata.packet_id, + received.message.metadata.retain + ); + } + if let Some(ack_token) = received.ack_token { + ack_tokens.push(ack_token); + } + let payload = if self.config.include_metadata { + serialize_mqtt5_envelope(&received.message)? + } else { + received.message.payload + }; + replay_messages.push(ReplayMessage { + headers: Some(received.message.headers.clone()), + payload: payload.clone(), + }); + messages.push(ProducedMessage { + id: None, + checksum: None, + timestamp: None, + origin_timestamp: None, + headers: Some(received.message.headers), + payload, + }); + } + let schema = if self.config.include_metadata { + Schema::Json + } else { + Schema::Raw + }; + *self.pending_batch.lock().await = Some(PendingBatch { + replay_messages, + schema, + ack_tokens, + in_flight: true, + }); + + Ok(ProducedMessages { + schema, + messages, + state: None, + }) + } + + async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { + let mut pending_batches = self.pending_batch.lock().await; + let Some(pending_batch) = pending_batches.as_mut() else { + return Ok(()); + }; + if result == SourceBatchResult::Nack { + pending_batch.in_flight = false; + warn!( + "NACK received for MQTT source connector with ID {}; replaying the batch before reading new messages", + self.id + ); + return Ok(()); + } + + let Some(mut pending_batch) = pending_batches.take() else { + return Ok(()); + }; + drop(pending_batches); + + if pending_batch.ack_tokens.is_empty() { + // QoS 0 has no broker token, so there is no broker acknowledgement + // to perform after Iggy accepts the batch. + return Ok(()); + } + + let Some(mut driver) = self.driver.lock().await.take() else { + *self.pending_batch.lock().await = Some(pending_batch); + return Err(Error::InitError( + "MQTT driver is not initialized".to_string(), + )); + }; + let mut acknowledgement = driver + // MQTT acknowledgements are sent only after Iggy acknowledges the + // batch, preserving at-least-once delivery. + .acknowledge_batch( + &mut pending_batch.ack_tokens, + self.poll_timeout, + self.batch_size, + self.max_retries(), + ) + .await; + for _ in 0..driver::ACK_BATCH_RETRY_ATTEMPTS { + if acknowledgement.is_ok() || pending_batch.ack_tokens.is_empty() { + break; + } + tokio::time::sleep(driver::ACK_RETRY_DELAY.min(self.poll_timeout)).await; + acknowledgement = driver + .acknowledge_batch( + &mut pending_batch.ack_tokens, + self.poll_timeout, + self.batch_size, + self.max_retries(), + ) + .await; + } + *self.driver.lock().await = Some(driver); + if let Err(error) = acknowledgement { Review Comment: Fixed by retaining unsent MQTT acknowledgement tokens and retrying them during later polls, including recovery through a broker disconnect, so max_retries = 0 no longer drops broker acknowledgements or blocks QoS delivery slots. -- 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]
