ygndotgg commented on code in PR #4303: URL: https://github.com/apache/iggy/pull/4303#discussion_r4158468620
########## core/connectors/sources/mqtt_source/src/lib.rs: ########## @@ -0,0 +1,1059 @@ +// 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 humantime::Duration as HumanDuration; +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}; + +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"; + +/// 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, + /// 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. + 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>, + /// 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>, + /// TLS server name, which must match the broker URL host. + pub server_name: 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>>, + // Committed only after the runtime confirms the corresponding Iggy batch. + state: Mutex<State>, + // At most one batch is staged because the SDK reports one batch result at a + // time. It pairs candidate state with QoS acknowledgement tokens. + 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, Clone, Default, Serialize, Deserialize)] +struct State { + // This is an operational counter rather than a broker offset. MQTT + // redelivery is controlled by the broker session and acknowledgement state. + acknowledged_messages: u64, +} + +#[derive(Debug)] +struct PendingBatch { + // Tokens are kept separately from ProducedMessage because the runtime + // acknowledges Iggy before the driver acknowledges MQTT. + ack_tokens: Vec<AckToken>, + // State is staged until the same batch has been persisted and acknowledged. + candidate_state: State, +} + +impl MqttSource { + pub fn new(id: u32, config: MqttSourceConfig, state: Option<ConnectorState>) -> Self { + // Invalid or incompatible state starts from a safe empty counter. The + // broker can still redeliver unacknowledged QoS messages after reconnect. + let restored_state = state.and_then(|state| { + state.deserialize::<State>(CONNECTOR_NAME, id).inspect(|state| { + info!( + "Restored MQTT source state for connector with ID {id}. Acknowledged messages: {}", + state.acknowledged_messages + ); + }) + }); + + Self { + id, + config, + poll_timeout: Duration::from_secs(1), + batch_size: DEFAULT_BATCH_SIZE, + batch_timeout: Duration::from_millis(10), + driver: Mutex::new(None), + state: Mutex::new(restored_state.unwrap_or_default()), + pending_batch: Mutex::new(None), + } + } + + fn serialize_state(&self, state: &State) -> Option<ConnectorState> { + ConnectorState::serialize(state, CONNECTOR_NAME, self.id) + } + + 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.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, + )) + } + + async fn current_state(&self) -> State { + self.state.lock().await.clone() + } + + 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::with_capacity(self.batch_size); + 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. + if self.pending_batch.lock().await.is_some() { + return Err(Error::InvalidState); + } + + 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 state = self.current_state().await; + // This is only a candidate. The runtime persists the returned batch and + // reports Ack or Nack through on_batch_result afterward. + let candidate_state = State { + acknowledged_messages: state + .acknowledged_messages + .saturating_add(received.len() as u64), + }; + let persisted_state = self.serialize_state(&candidate_state).ok_or_else(|| { + Error::Serialization("failed to serialize MQTT source state".to_string()) + })?; + let mut 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); + } + messages.push(ProducedMessage { + id: None, + checksum: None, + timestamp: None, + origin_timestamp: None, + headers: Some(received.message.headers), + payload: received.message.payload, + }); + } + *self.pending_batch.lock().await = Some(PendingBatch { + ack_tokens, + candidate_state, + }); + + Ok(ProducedMessages { + schema: Schema::Raw, + messages, + state: Some(persisted_state), + }) + } + + async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { + // A Nack means the runtime could not confirm the Iggy write or state save. + // Keep committed state unchanged so QoS 1/2 messages can be redelivered. + let Some(mut pending_batch) = self.pending_batch.lock().await.take() else { + return Ok(()); + }; + if result == SourceBatchResult::Nack { Review Comment: Fixed by Retaining and reply MQTT batches after runtime NACKs so messages are not lost while the broker connection remains active -- 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]
