ygndotgg commented on code in PR #4303: URL: https://github.com/apache/iggy/pull/4303#discussion_r4158455736
########## core/connectors/sources/mqtt_source/src/driver.rs: ########## @@ -0,0 +1,1157 @@ +// 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 iggy_common::{HeaderKey, HeaderValue}; +use rumqttc::tokio_rustls::rustls::{self, ClientConfig, RootCertStore}; +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 std::{ + collections::{BTreeMap, VecDeque}, + io::{BufReader, Cursor}, + sync::Arc, + time::Duration, +}; +use tokio::time::timeout; +use url::Url; + +const ACK_RETRY_ATTEMPTS: usize = 5; +const ACK_RETRY_DELAY: Duration = Duration::from_millis(10); + +// 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, +} + +/// 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()) + })?; + } + poll_mqtt311(&mut event_loop, poll_timeout).await?; + + Ok(Self { + connection: MqttConnection::Mqtt311 { + client: Box::new(client), + event_loop: Box::new(event_loop), + }, + buffered_messages: VecDeque::new(), + }) + } + 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()) + })?; + } + poll_mqtt5(&mut event_loop, poll_timeout).await?; + + Ok(Self { + connection: MqttConnection::Mqtt5 { + client: Box::new(client), + event_loop: Box::new(event_loop), + }, + buffered_messages: VecDeque::new(), + }) + } + } + } + + 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, Review Comment: Fixed by validating every SUBACK during initialization and failing startup when any subscription is rejected, preventing silent connectors. -- 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]
