diff --git a/.github/workflows/doc-deploy.yml b/.github/workflows/doc-deploy.yml index 63a2d6ed..b4550b8a 100644 --- a/.github/workflows/doc-deploy.yml +++ b/.github/workflows/doc-deploy.yml @@ -20,7 +20,7 @@ jobs: version: 8.10.5 - uses: actions/setup-node@v4 with: - node-version: '20' + node-version: '22' cache: 'pnpm' cache-dependency-path: ./docs/pnpm-lock.yaml diff --git a/.github/workflows/test-doc-deploy.yml b/.github/workflows/test-doc-deploy.yml index 9ac1639f..19ce93b6 100644 --- a/.github/workflows/test-doc-deploy.yml +++ b/.github/workflows/test-doc-deploy.yml @@ -18,7 +18,7 @@ jobs: version: 8.10.5 - uses: actions/setup-node@v4 with: - node-version: '20' + node-version: '22' cache: 'pnpm' cache-dependency-path: ./docs/pnpm-lock.yaml diff --git a/crates/arkflow-plugin/src/buffer/common_window.rs b/crates/arkflow-plugin/src/buffer/common_window.rs new file mode 100644 index 00000000..28928e0e --- /dev/null +++ b/crates/arkflow-plugin/src/buffer/common_window.rs @@ -0,0 +1,193 @@ +/* + * Licensed 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. + */ + +//! Common window infrastructure +//! +//! This module provides shared utilities for window implementations +//! to reduce code duplication across different window types. + +use crate::buffer::join::JoinConfig; +use crate::buffer::window::BaseWindow; +use arkflow_core::{Error, Resource}; +use std::sync::Arc; +use std::time::Instant; +use tokio::sync::Notify; +use tokio::time::sleep; +use tokio_util::sync::CancellationToken; + +/// Common window context +/// +/// Provides shared infrastructure for window implementations including +/// timer management, notification, and cancellation handling. +pub struct CommonWindowContext { + /// Notification mechanism for signaling between threads + pub notify: Arc, + /// Token for cancellation of background tasks + pub close_token: CancellationToken, + /// Last time the window was triggered + pub last_trigger: Arc>, +} + +impl CommonWindowContext { + /// Create a new common window context + /// + /// # Returns + /// + /// A new CommonWindowContext instance + pub fn new() -> Self { + Self { + notify: Arc::new(Notify::new()), + close_token: CancellationToken::new(), + last_trigger: Arc::new(std::sync::RwLock::new(Instant::now())), + } + } + + /// Start the background timer task + /// + /// This spawns a background task that periodically notifies waiters + /// based on the check interval. The task runs until cancelled. + /// + /// # Arguments + /// + /// * `check_interval` - How often to check if window should trigger + pub fn start_timer(&self, check_interval: std::time::Duration) { + let notify = Arc::clone(&self.notify); + let close = self.close_token.clone(); + + tokio::spawn(async move { + loop { + let timer = sleep(check_interval); + tokio::select! { + _ = timer => { + notify.notify_waiters(); + } + _ = close.cancelled() => { + notify.notify_waiters(); + break; + } + _ = notify.notified() => { + if close.is_cancelled(){ + break; + } + } + } + } + }); + } + + /// Update the last trigger time + /// + /// # Arguments + /// + /// * `time` - The new last trigger time + pub fn update_last_trigger(&self, time: Instant) { + if let Ok(mut last) = self.last_trigger.write() { + *last = time; + } + } + + /// Get the last trigger time + /// + /// # Returns + /// + /// The last trigger time + pub fn get_last_trigger(&self) -> Instant { + self.last_trigger + .read() + .map(|t| *t) + .unwrap_or_else(|_| Instant::now()) + } + + /// Check if the window is closed + /// + /// # Returns + /// + /// * `bool` - true if closed, false otherwise + pub fn is_closed(&self) -> bool { + self.close_token.is_cancelled() + } + + /// Close the window context + /// + /// Cancels the background timer task + pub fn close(&self) { + self.close_token.cancel(); + } + + /// Create a BaseWindow with join support + /// + /// # Arguments + /// + /// * `join_config` - Optional join configuration + /// * `gap` - Time interval for the timer + /// * `resource` - Resource reference + /// + /// # Returns + /// + /// A BaseWindow instance or an error + pub fn create_base_window( + &self, + join_config: Option, + gap: std::time::Duration, + resource: &Resource, + ) -> Result { + BaseWindow::new( + join_config, + Arc::clone(&self.notify), + self.close_token.clone(), + gap, + resource, + ) + } +} + +impl Default for CommonWindowContext { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_common_window_context_creation() { + let ctx = CommonWindowContext::new(); + assert!(!ctx.is_closed()); + assert!(ctx.get_last_trigger() <= Instant::now()); + } + + #[test] + fn test_common_window_context_close() { + let ctx = CommonWindowContext::new(); + assert!(!ctx.is_closed()); + ctx.close(); + assert!(ctx.is_closed()); + } + + #[test] + fn test_common_window_context_update_trigger() { + let ctx = CommonWindowContext::new(); + let now = Instant::now(); + ctx.update_last_trigger(now); + assert!(ctx.get_last_trigger() >= now); + } + + #[test] + fn test_common_window_context_default() { + let ctx = CommonWindowContext::default(); + assert!(!ctx.is_closed()); + } +} diff --git a/crates/arkflow-plugin/src/buffer/mod.rs b/crates/arkflow-plugin/src/buffer/mod.rs index 6ac3cb42..bb0ff399 100644 --- a/crates/arkflow-plugin/src/buffer/mod.rs +++ b/crates/arkflow-plugin/src/buffer/mod.rs @@ -11,12 +11,14 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +pub mod common_window; mod join; pub mod memory; pub mod session_window; pub mod sliding_window; pub mod tumbling_window; pub(crate) mod window; +pub mod window_strategy; use arkflow_core::Error; diff --git a/crates/arkflow-plugin/src/buffer/session_window.rs b/crates/arkflow-plugin/src/buffer/session_window.rs index 578d00be..78ea19e9 100644 --- a/crates/arkflow-plugin/src/buffer/session_window.rs +++ b/crates/arkflow-plugin/src/buffer/session_window.rs @@ -19,6 +19,7 @@ //! if they arrive within the gap duration of each other. When the gap duration elapses //! without new messages, the session is closed and all accumulated messages are emitted. +use crate::buffer::common_window::CommonWindowContext; use crate::buffer::join::JoinConfig; use crate::buffer::window::BaseWindow; use crate::time::deserialize_duration; @@ -30,9 +31,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::sync::Arc; use std::time; -use tokio::sync::{Notify, RwLock}; +use tokio::sync::RwLock; use tokio::time::Instant; -use tokio_util::sync::CancellationToken; /// Configuration for the session window buffer #[derive(Debug, Clone, Serialize, Deserialize)] @@ -52,10 +52,8 @@ struct SessionWindow { /// Configuration parameters for the session window config: SessionWindowConfig, base_window: BaseWindow, - /// Notification mechanism for signaling between threads - notify: Arc, - /// Token for cancellation of background tasks - close: CancellationToken, + /// Common window context for timer and notification management + context: CommonWindowContext, /// Timestamp of the last received message, used to determine session boundaries last_message_time: Arc>, } @@ -69,26 +67,16 @@ impl SessionWindow { /// # Returns /// * `Result` - A new session window instance or an error fn new(config: SessionWindowConfig, resource: &Resource) -> Result { - let notify = Arc::new(Notify::new()); - let notify_clone = Arc::clone(¬ify); - let gap = config.gap; - let close = CancellationToken::new(); - let close_clone = close.clone(); - let last_message_time = Arc::new(RwLock::new(Instant::now())); - let base_window = BaseWindow::new( - config.join.clone(), - notify_clone, - close_clone, - gap, - resource, - )?; + let context = CommonWindowContext::new(); + + // BaseWindow already starts a background timer, no need to start another one here + let base_window = context.create_base_window(config.join.clone(), config.gap, resource)?; Ok(Self { - close, - notify, config, base_window, - last_message_time, + context, + last_message_time: Arc::new(RwLock::new(Instant::now())), }) } } @@ -107,6 +95,8 @@ impl Buffer for SessionWindow { self.base_window.write(msg, ack).await?; // Update the last message timestamp to track session activity *self.last_message_time.write().await = Instant::now(); + // Notify waiting readers that a new message has arrived + self.context.notify.notify_waiters(); Ok(()) } @@ -117,7 +107,7 @@ impl Buffer for SessionWindow { /// * `Result)>, Error>` - The merged message batch and combined acknowledgment, /// or None if the buffer is closed and empty async fn read(&self) -> Result)>, Error> { - if self.close.is_cancelled() { + if self.context.is_closed() { return Ok(None); } @@ -134,8 +124,7 @@ impl Buffer for SessionWindow { } // Wait for notification from timer or write operation - let notify = Arc::clone(&self.notify); - notify.notified().await; + self.context.notify.notified().await; } // Process and return the current session self.base_window.process_window().await diff --git a/crates/arkflow-plugin/src/buffer/sliding_window.rs b/crates/arkflow-plugin/src/buffer/sliding_window.rs index 44aa53b4..3f2b0d0c 100644 --- a/crates/arkflow-plugin/src/buffer/sliding_window.rs +++ b/crates/arkflow-plugin/src/buffer/sliding_window.rs @@ -76,6 +76,7 @@ impl SlidingWindow { let close = CancellationToken::new(); let close_clone = close.clone(); + // SlidingWindow needs its own timer since it doesn't use BaseWindow tokio::spawn(async move { loop { let timer = sleep(interval); @@ -97,10 +98,10 @@ impl SlidingWindow { }); Ok(Self { - close, - notify, config, queue: Arc::new(Default::default()), + notify, + close, }) } @@ -171,6 +172,9 @@ impl Buffer for SlidingWindow { async fn write(&self, msg: MessageBatchRef, ack: Arc) -> Result<(), Error> { let mut queue_lock = self.queue.write().await; queue_lock.push_back((msg, ack)); + drop(queue_lock); + // Notify waiting readers that a new message has arrived + self.notify.notify_waiters(); Ok(()) } @@ -186,9 +190,11 @@ impl Buffer for SlidingWindow { } loop { + if self.close.is_cancelled() { + return Ok(None); + } { - let queue_arc = Arc::clone(&self.queue); - let queue_lock = queue_arc.read().await; + let queue_lock = self.queue.read().await; // If there are enough messages to form a window, break the loop and process them if queue_lock.len() >= self.config.window_size as usize { break; @@ -196,8 +202,7 @@ impl Buffer for SlidingWindow { // If the buffer is closed, return None } // Wait for notification from timer, write operation, or close - let notify = Arc::clone(&self.notify); - notify.notified().await; + self.notify.notified().await; } // Process the current window and slide forward self.process_slide().await @@ -209,12 +214,10 @@ impl Buffer for SlidingWindow { /// * `Result<(), Error>` - Success or an error async fn flush(&self) -> Result<(), Error> { self.close.cancel(); - let queue_arc = Arc::clone(&self.queue); - let queue_lock = queue_arc.read().await; + let queue_lock = self.queue.read().await; if !queue_lock.is_empty() { // Notify any waiting readers to process remaining messages - let notify = Arc::clone(&self.notify); - notify.notify_waiters(); + self.notify.notify_waiters(); } Ok(()) } diff --git a/crates/arkflow-plugin/src/buffer/tumbling_window.rs b/crates/arkflow-plugin/src/buffer/tumbling_window.rs index 88695801..e68c57c5 100644 --- a/crates/arkflow-plugin/src/buffer/tumbling_window.rs +++ b/crates/arkflow-plugin/src/buffer/tumbling_window.rs @@ -19,6 +19,7 @@ //! period elapses, all accumulated messages are emitted as a single batch and a new //! window begins immediately. +use crate::buffer::common_window::CommonWindowContext; use crate::buffer::join::JoinConfig; use crate::buffer::window::BaseWindow; use crate::time::deserialize_duration; @@ -30,8 +31,6 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::sync::Arc; use std::time; -use tokio::sync::Notify; -use tokio_util::sync::CancellationToken; /// Configuration for the tumbling window buffer #[derive(Debug, Clone, Serialize, Deserialize)] @@ -48,12 +47,10 @@ struct TumblingWindowConfig { /// Tumbling window buffer implementation /// Groups messages into fixed-size, non-overlapping time windows struct TumblingWindow { - /// Thread-safe queue to store message batches and their acknowledgments + /// Base window implementation for queue management base_window: BaseWindow, - /// Notification mechanism for signaling between threads - notify: Arc, - /// Token for cancellation of background tasks - close: CancellationToken, + /// Common window context for timer and notification management + context: CommonWindowContext, } impl TumblingWindow { @@ -65,23 +62,14 @@ impl TumblingWindow { /// # Returns /// * `Result` - A new tumbling window instance or an error fn new(config: TumblingWindowConfig, resource: &Resource) -> Result { - let notify = Arc::new(Notify::new()); - let notify_clone = Arc::clone(¬ify); - let interval = config.interval; - let close = CancellationToken::new(); - let close_clone = close.clone(); - let base_window = BaseWindow::new( - config.join.clone(), - notify_clone, - close_clone, - interval, - resource, - )?; + let context = CommonWindowContext::new(); + + // BaseWindow already starts a background timer, no need to start another one here + let base_window = context.create_base_window(config.join, config.interval, resource)?; Ok(Self { - close, - notify, base_window, + context, }) } } @@ -108,7 +96,7 @@ impl Buffer for TumblingWindow { /// or None if the buffer is closed and empty async fn read(&self) -> Result)>, Error> { // If the buffer is closed, return None - if self.close.is_cancelled() { + if self.context.is_closed() { return Ok(None); } @@ -120,8 +108,7 @@ impl Buffer for TumblingWindow { } } // Wait for notification from timer, write operation, or close - let notify = Arc::clone(&self.notify); - notify.notified().await; + self.context.notify.notified().await; } // Process and return the current window self.base_window.process_window().await diff --git a/crates/arkflow-plugin/src/buffer/window_strategy.rs b/crates/arkflow-plugin/src/buffer/window_strategy.rs new file mode 100644 index 00000000..ff940952 --- /dev/null +++ b/crates/arkflow-plugin/src/buffer/window_strategy.rs @@ -0,0 +1,217 @@ +/* + * Licensed 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. + */ + +//! Window strategy trait for buffer implementations +//! +//! This module provides a unified interface for different windowing strategies +//! to reduce code duplication across tumbling, sliding, and session windows. + +use std::time::Instant; + +/// Window strategy trait +/// +/// Defines how windows should be triggered and managed. +/// Different window types implement this trait with their specific logic. +pub trait WindowStrategy: Send + Sync { + /// Check if the window should trigger based on current time and last trigger + /// + /// # Arguments + /// + /// * `current_time` - Current time instant + /// * `last_trigger` - Time when the window was last triggered + /// + /// # Returns + /// + /// * `bool` - true if window should trigger, false otherwise + fn should_trigger(&self, current_time: Instant, last_trigger: Instant) -> bool; + + /// Get the interval for window checking + /// + /// # Returns + /// + /// * `Duration` - How often to check if window should trigger + fn check_interval(&self) -> std::time::Duration; + + /// Reset the strategy state (called after window triggers) + /// + /// Some strategies may need to reset state after triggering + fn reset(&mut self); +} + +/// Tumbling window strategy +/// +/// Windows are fixed-size, non-overlapping, and consecutive. +/// Example: 5-minute windows trigger at 0, 5, 10, 15... +#[derive(Clone, Debug)] +pub struct TumblingWindowStrategy { + /// Fixed duration of each window + window_size: std::time::Duration, +} + +impl TumblingWindowStrategy { + /// Create a new tumbling window strategy + /// + /// # Arguments + /// + /// * `window_size` - Fixed duration of each window + pub fn new(window_size: std::time::Duration) -> Self { + Self { window_size } + } +} + +impl WindowStrategy for TumblingWindowStrategy { + fn should_trigger(&self, current_time: Instant, last_trigger: Instant) -> bool { + current_time.duration_since(last_trigger) >= self.window_size + } + + fn check_interval(&self) -> std::time::Duration { + // Check more frequently than window size to ensure timely triggers + std::time::Duration::from_millis(self.window_size.as_millis().min(1000) as u64) + } + + fn reset(&mut self) { + // Tumbling windows don't need to reset state + // Each trigger is independent + } +} + +/// Sliding window strategy +/// +/// Windows overlap and slide forward by a fixed interval. +/// Example: 5-minute windows sliding every 1 minute +#[derive(Clone, Debug)] +pub struct SlidingWindowStrategy { + /// Size of each window + window_size: std::time::Duration, + /// How often to slide the window + slide_interval: std::time::Duration, +} + +impl SlidingWindowStrategy { + /// Create a new sliding window strategy + /// + /// # Arguments + /// + /// * `window_size` - Size of each window + /// * `slide_interval` - How often to slide the window + pub fn new(window_size: std::time::Duration, slide_interval: std::time::Duration) -> Self { + Self { + window_size, + slide_interval, + } + } +} + +impl WindowStrategy for SlidingWindowStrategy { + fn should_trigger(&self, current_time: Instant, last_trigger: Instant) -> bool { + current_time.duration_since(last_trigger) >= self.slide_interval + } + + fn check_interval(&self) -> std::time::Duration { + // Check at slide interval to ensure timely triggers + self.slide_interval + } + + fn reset(&mut self) { + // Sliding windows don't need to reset state + } +} + +/// Session window strategy +/// +/// Windows are dynamic and close after a period of inactivity (gap). +#[derive(Clone, Debug)] +pub struct SessionWindowStrategy { + /// Gap duration of inactivity before closing a session + gap_duration: std::time::Duration, +} + +impl SessionWindowStrategy { + /// Create a new session window strategy + /// + /// # Arguments + /// + /// * `gap_duration` - Period of inactivity before closing session + pub fn new(gap_duration: std::time::Duration) -> Self { + Self { gap_duration } + } +} + +impl WindowStrategy for SessionWindowStrategy { + fn should_trigger(&self, current_time: Instant, last_trigger: Instant) -> bool { + // Trigger when gap duration has passed since last activity + current_time.duration_since(last_trigger) >= self.gap_duration + } + + fn check_interval(&self) -> std::time::Duration { + // Check frequently to detect session gaps quickly + std::time::Duration::from_millis(self.gap_duration.as_millis().min(500) as u64) + } + + fn reset(&mut self) { + // Session windows reset on each new message + // (managed externally via last_trigger updates) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn test_tumbling_window_strategy() { + let strategy = TumblingWindowStrategy::new(Duration::from_secs(5)); + let base = Instant::now(); + + // Should not trigger immediately + assert!(!strategy.should_trigger(base, base)); + + // Should trigger after 5 seconds + assert!(strategy.should_trigger(base + Duration::from_secs(5), base)); + + // Check interval should be reasonable + assert!(strategy.check_interval() <= Duration::from_secs(1)); + } + + #[test] + fn test_sliding_window_strategy() { + let strategy = SlidingWindowStrategy::new(Duration::from_secs(5), Duration::from_secs(1)); + let base = Instant::now(); + + // Should not trigger immediately + assert!(!strategy.should_trigger(base, base)); + + // Should trigger after 1 second (slide interval) + assert!(strategy.should_trigger(base + Duration::from_secs(1), base)); + + // Check interval should match slide interval + assert_eq!(strategy.check_interval(), Duration::from_secs(1)); + } + + #[test] + fn test_session_window_strategy() { + let strategy = SessionWindowStrategy::new(Duration::from_secs(10)); + let base = Instant::now(); + + // Should not trigger immediately + assert!(!strategy.should_trigger(base, base)); + + // Should trigger after 10 seconds of inactivity + assert!(strategy.should_trigger(base + Duration::from_secs(10), base)); + + // Check interval should be frequent + assert!(strategy.check_interval() <= Duration::from_millis(500)); + } +} diff --git a/crates/arkflow-plugin/src/input/mqtt.rs b/crates/arkflow-plugin/src/input/mqtt.rs index ff784425..d4d9a2a7 100644 --- a/crates/arkflow-plugin/src/input/mqtt.rs +++ b/crates/arkflow-plugin/src/input/mqtt.rs @@ -16,6 +16,7 @@ //! //! Receive data from the MQTT broker +use crate::mqtt_client::{create_mqtt_options, parse_qos}; use arkflow_core::codec::Codec; use arkflow_core::error_helpers::parse_config; use arkflow_core::input::{register_input_builder, Ack, Input, InputBuilder}; @@ -23,7 +24,7 @@ use arkflow_core::{Error, MessageBatch, MessageBatchRef, Resource}; use async_trait::async_trait; use flume::{Receiver, Sender}; -use rumqttc::{AsyncClient, Event, MqttOptions, Packet, Publish, QoS}; +use rumqttc::{AsyncClient, Event, Packet, Publish}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use tokio::sync::Mutex; @@ -92,34 +93,23 @@ impl MqttInput { #[async_trait] impl Input for MqttInput { async fn connect(&self) -> Result<(), Error> { - // Create MQTT options - let mut mqtt_options = - MqttOptions::new(&self.config.client_id, &self.config.host, self.config.port); + // Create MQTT options using shared utility + let mut mqtt_options = create_mqtt_options( + &self.config.client_id, + &self.config.host, + self.config.port, + self.config.username.as_deref(), + self.config.password.as_deref(), + self.config.keep_alive, + self.config.clean_session, + ); mqtt_options.set_manual_acks(true); - // Set the authentication information - if let (Some(username), Some(password)) = (&self.config.username, &self.config.password) { - mqtt_options.set_credentials(username, password); - } - - // Set the keep-alive time - if let Some(keep_alive) = self.config.keep_alive { - mqtt_options.set_keep_alive(std::time::Duration::from_secs(keep_alive)); - } - - // Set up a clean session - if let Some(clean_session) = self.config.clean_session { - mqtt_options.set_clean_session(clean_session); - } // Create an MQTT client let (client, mut eventloop) = AsyncClient::new(mqtt_options, 10); + // Subscribe to topics - let qos_level = match self.config.qos { - Some(0) => QoS::AtMostOnce, - Some(1) => QoS::AtLeastOnce, - Some(2) => QoS::ExactlyOnce, - _ => QoS::AtLeastOnce, // Default is QoS 1 - }; + let qos_level = parse_qos(self.config.qos); for topic in &self.config.topics { client.subscribe(topic, qos_level).await.map_err(|e| { diff --git a/crates/arkflow-plugin/src/lib.rs b/crates/arkflow-plugin/src/lib.rs index bd7e682a..5cf760d6 100644 --- a/crates/arkflow-plugin/src/lib.rs +++ b/crates/arkflow-plugin/src/lib.rs @@ -18,6 +18,7 @@ pub mod component; pub mod context_pool; pub mod expr; pub mod input; +pub mod mqtt_client; pub mod output; pub mod processor; pub mod pulsar; diff --git a/crates/arkflow-plugin/src/mqtt_client.rs b/crates/arkflow-plugin/src/mqtt_client.rs new file mode 100644 index 00000000..50cbc682 --- /dev/null +++ b/crates/arkflow-plugin/src/mqtt_client.rs @@ -0,0 +1,168 @@ +/* + * Licensed 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. + */ + +//! MQTT client utilities for ArkFlow +//! +//! This module provides common utilities for MQTT client configuration +//! to reduce code duplication between input and output components. + +use rumqttc::{MqttOptions, QoS}; +use std::time::Duration; + +/// Create MqttOptions from common configuration parameters +/// +/// This function eliminates code duplication between MQTT input and output +/// components by providing a centralized way to configure MQTT options. +/// +/// # Arguments +/// +/// * `client_id` - Unique client identifier +/// * `host` - MQTT broker address +/// * `port` - MQTT broker port +/// * `username` - Optional username for authentication +/// * `password` - Optional password for authentication +/// * `keep_alive` - Optional keep-alive interval in seconds +/// * `clean_session` - Optional clean session flag +/// +/// # Returns +/// +/// Configured MqttOptions ready for client creation +/// +/// # Examples +/// +/// ```rust,no_run +/// use arkflow_plugin::mqtt_client::create_mqtt_options; +/// +/// let options = create_mqtt_options( +/// "my_client", +/// "localhost", +/// 1883, +/// Some("user"), +/// Some("pass"), +/// Some(60), +/// Some(true), +/// ); +/// ``` +pub fn create_mqtt_options( + client_id: &str, + host: &str, + port: u16, + username: Option<&str>, + password: Option<&str>, + keep_alive: Option, + clean_session: Option, +) -> MqttOptions { + let mut mqtt_options = MqttOptions::new(client_id, host, port); + + // Set authentication credentials if provided + if let (Some(username), Some(password)) = (username, password) { + mqtt_options.set_credentials(username, password); + } + + // Set keep-alive interval if provided + if let Some(keep_alive) = keep_alive { + mqtt_options.set_keep_alive(Duration::from_secs(keep_alive)); + } + + // Set clean session flag if provided + if let Some(clean_session) = clean_session { + mqtt_options.set_clean_session(clean_session); + } + + mqtt_options +} + +/// Convert QoS level from u8 to rumqttc QoS enum +/// +/// # Arguments +/// +/// * `qos` - QoS level as u8 (0, 1, or 2) +/// +/// # Returns +/// +/// Corresponding QoS enum value +/// +/// # Examples +/// +/// ```rust,no_run +/// use arkflow_plugin::mqtt_client::parse_qos; +/// +/// let qos = parse_qos(Some(1)); // Returns QoS::AtLeastOnce +/// let qos_default = parse_qos(None); // Returns QoS::AtLeastOnce (default) +/// ``` +pub fn parse_qos(qos: Option) -> QoS { + match qos { + Some(0) => QoS::AtMostOnce, + Some(1) => QoS::AtLeastOnce, + Some(2) => QoS::ExactlyOnce, + _ => QoS::AtLeastOnce, // Default is QoS 1 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_create_mqtt_options_basic() { + let options = create_mqtt_options("test_client", "localhost", 1883, None, None, None, None); + + assert_eq!(options.client_id(), "test_client"); + let broker_addr = options.broker_address(); + assert_eq!(broker_addr.0, "localhost"); + assert_eq!(broker_addr.1, 1883); + } + + #[test] + fn test_create_mqtt_options_with_auth() { + let options = create_mqtt_options( + "test_client", + "localhost", + 1883, + Some("user"), + Some("pass"), + None, + None, + ); + + // Verify credentials are set (not directly accessible, but no panic = success) + assert_eq!(options.client_id(), "test_client"); + } + + #[test] + fn test_create_mqtt_options_full() { + let options = create_mqtt_options( + "test_client", + "localhost", + 1883, + Some("user"), + Some("pass"), + Some(60), + Some(true), + ); + + assert_eq!(options.client_id(), "test_client"); + let broker_addr = options.broker_address(); + assert_eq!(broker_addr.1, 1883); + } + + #[test] + fn test_parse_qos() { + assert!(matches!(parse_qos(Some(0)), QoS::AtMostOnce)); + assert!(matches!(parse_qos(Some(1)), QoS::AtLeastOnce)); + assert!(matches!(parse_qos(Some(2)), QoS::ExactlyOnce)); + assert!(matches!(parse_qos(None), QoS::AtLeastOnce)); // Default + assert!(matches!(parse_qos(Some(99)), QoS::AtLeastOnce)); // Invalid -> default + } +} diff --git a/crates/arkflow-plugin/src/output/mqtt.rs b/crates/arkflow-plugin/src/output/mqtt.rs index 92fb6137..fe5e4206 100644 --- a/crates/arkflow-plugin/src/output/mqtt.rs +++ b/crates/arkflow-plugin/src/output/mqtt.rs @@ -17,6 +17,7 @@ //! Send the processed data to the MQTT broker use crate::expr::Expr; +use crate::mqtt_client::{create_mqtt_options, parse_qos}; use arkflow_core::error_helpers::parse_config; use arkflow_core::{ codec::Codec, @@ -83,24 +84,16 @@ impl MqttOutput { #[async_trait] impl Output for MqttOutput { async fn connect(&self) -> Result<(), Error> { - // Create MQTT options - let mut mqtt_options = - MqttOptions::new(&self.config.client_id, &self.config.host, self.config.port); - - // Set the authentication information - if let (Some(username), Some(password)) = (&self.config.username, &self.config.password) { - mqtt_options.set_credentials(username, password); - } - - // Set the keep-alive time - if let Some(keep_alive) = self.config.keep_alive { - mqtt_options.set_keep_alive(std::time::Duration::from_secs(keep_alive)); - } - - // Set up a purge session - if let Some(clean_session) = self.config.clean_session { - mqtt_options.set_clean_session(clean_session); - } + // Create MQTT options using shared utility + let mqtt_options = create_mqtt_options( + &self.config.client_id, + &self.config.host, + self.config.port, + self.config.username.as_deref(), + self.config.password.as_deref(), + self.config.keep_alive, + self.config.clean_session, + ); // Create an MQTT client let (client, mut eventloop) = T::create(mqtt_options, 10).await?; diff --git a/docs/blog/2025/04/01-v0.2.0-rc1.md b/docs/blog/2025/04/01-v0.2.0-rc1.md index 9c9ee08f..e29ed886 100644 --- a/docs/blog/2025/04/01-v0.2.0-rc1.md +++ b/docs/blog/2025/04/01-v0.2.0-rc1.md @@ -7,6 +7,8 @@ authors: chenquan We are excited to announce the release of ArkFlow v0.2.0-rc1! ArkFlow is a high-performance Rust stream processing engine that provides powerful data stream processing capabilities. This version brings many important features and improvements. + + ## Key Features ### Diverse Component Support diff --git a/docs/blog/2025/04/15-v0.2.0-release.md b/docs/blog/2025/04/15-v0.2.0-release.md index c7753b69..92ee3f53 100644 --- a/docs/blog/2025/04/15-v0.2.0-release.md +++ b/docs/blog/2025/04/15-v0.2.0-release.md @@ -6,6 +6,8 @@ authors: chenquan We are excited to announce the official release of ArkFlow v0.2.0! ArkFlow is a high-performance Rust stream processing engine that provides powerful data flow processing capabilities. This version builds upon RC1/RC2 with further optimizations and improvements, delivering more stable performance and richer functionality. + + ## Key Features ### Diverse Component Support