diff --git a/Cargo.lock b/Cargo.lock index 93cd7fcd..1f303a98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -254,6 +254,7 @@ dependencies = [ "futures", "lazy_static", "num_cpus", + "reqwest", "serde", "serde_json", "serde_yaml", @@ -4639,7 +4640,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.2", ] [[package]] @@ -5944,7 +5945,7 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "itertools 0.14.0", "log", "multimap", @@ -7365,7 +7366,7 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03c3c6b7927ffe7ecaa769ee0e3994da3b8cafc8f444578982c83ecb161af917" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.101", diff --git a/crates/arkflow-core/Cargo.toml b/crates/arkflow-core/Cargo.toml index d1986276..a32df716 100644 --- a/crates/arkflow-core/Cargo.toml +++ b/crates/arkflow-core/Cargo.toml @@ -27,4 +27,5 @@ clap = { workspace = true } colored = { workspace = true } flume = { workspace = true } axum = { workspace = true } +reqwest = { workspace = true, features = ["json"] } num_cpus = "1.17.0" \ No newline at end of file diff --git a/crates/arkflow-core/src/cli/mod.rs b/crates/arkflow-core/src/cli/mod.rs index f2d1686c..b51dc071 100644 --- a/crates/arkflow-core/src/cli/mod.rs +++ b/crates/arkflow-core/src/cli/mod.rs @@ -14,17 +14,24 @@ use crate::config::{EngineConfig, LogFormat}; use crate::engine::Engine; +use crate::remote_config::RemoteConfigManager; use clap::{Arg, Command}; use std::process; +use tokio::signal::unix::{signal, SignalKind}; +use tokio_util::sync::CancellationToken; use tracing::{info, Level}; use tracing_subscriber::fmt; pub struct Cli { pub config: Option, + pub remote_config_manager: Option, } impl Default for Cli { fn default() -> Self { - Self { config: None } + Self { + config: None, + remote_config_manager: None, + } } } @@ -33,54 +40,120 @@ impl Cli { let matches = Command::new("arkflow") .version("0.4.0-rc1") .author("chenquan") - .about("High-performance Rust stream processing engine, providing powerful data stream processing capabilities, supporting multiple input/output sources and processors.") + .about("High-performance Rust stream processing engine, providing powerful data stream processing capabilities, supporting multiple input/output sources and processors") .arg( Arg::new("config") .short('c') .long("config") .value_name("FILE") - .help("Specify the profile path.") - .required(true), + .help("Specify the profile path") ) .arg( Arg::new("validate") .short('v') .long("validate") - .help("Only the profile is verified, not the engine is started.") + .help("Only the profile is verified, not the engine is started") .action(clap::ArgAction::SetTrue), ) + .subcommand( + Command::new("remote").about("Use remote configuration for automatic stream management") + .arg( + Arg::new("url") + .long("url") + .value_name("URL") + .help("Remote configuration API endpoint URL for automatic stream management") + .required( true) + ) + .arg( + Arg::new("interval") + .long("interval") + .value_name("SECONDS") + .help("Interval in seconds for polling remote configuration") + .default_value("30"), + ) + .arg( + Arg::new("token") + .long("token") + .value_name("TOKEN") + .help("Authentication token for remote configuration API") + .required( true), + )) .get_matches(); - // Get the profile path - let config_path = matches.get_one::("config").unwrap(); + // Check if using remote configuration + if let Some(remote) = matches.subcommand_matches("remote") { + // Initialize remote configuration manager + let interval = remote + .get_one::("interval") + .and_then(|s| s.parse::().ok()) + .unwrap_or(30); + let token = remote.get_one::("token").cloned(); + let remote_url = remote + .get_one::("url") + .expect("Remote configuration URL not found"); - // Get the profile path - let config = match EngineConfig::from_file(config_path) { - Ok(config) => config, - Err(e) => { - println!("Failed to load configuration file: {}", e); - process::exit(1); + let remote_manager = RemoteConfigManager::new(remote_url.clone(), interval, token); + + self.remote_config_manager.replace(remote_manager); + info!("Using remote configuration from: {}", remote_url); + } else { + // Use local configuration file + let config_path = matches + .get_one::("config") + .ok_or("Configuration not found")?; + + let config = match EngineConfig::from_file(config_path) { + Ok(config) => config, + Err(e) => { + println!("Failed to load configuration file: {}", e); + process::exit(1); + } + }; + + // If you just verify the configuration, exit it + if matches.get_flag("validate") { + info!("The config is validated."); + return Ok(()); } - }; - // If you just verify the configuration, exit it - if matches.get_flag("validate") { - info!("The config is validated."); - return Ok(()); + self.config = Some(config); } - self.config = Some(config); Ok(()) } pub async fn run(&self) -> Result<(), Box> { - // Initialize the logging system - let config = self.config.clone().unwrap(); - init_logging(&config); - let engine = Engine::new(config); - engine.run().await?; + let token = CancellationToken::new(); + + if let Some(remote_manager) = &self.remote_config_manager { + // Run with remote configuration management + remote_manager.run(token.clone()).await?; + } else { + // Run with local configuration + let config = self.config.clone().unwrap(); + init_logging(&config); + let engine = Engine::new(config); + engine.run(token.clone()).await?; + } + + // Set up signal handlers + let mut sigint = signal(SignalKind::interrupt()).expect("Failed to set signal handler"); + let mut sigterm = signal(SignalKind::terminate()).expect("Failed to set signal handler"); + + tokio::spawn(async move { + tokio::select! { + _ = sigint.recv() => { + info!("Received SIGINT, exiting..."); + + }, + _ = sigterm.recv() => { + info!("Received SIGTERM, exiting..."); + } + } + token.cancel(); + }); Ok(()) } } -fn init_logging(config: &EngineConfig) -> () { +pub(crate) fn init_logging(config: &EngineConfig) -> () { let log_level = match config.logging.level.as_str() { "trace" => Level::TRACE, "debug" => Level::DEBUG, diff --git a/crates/arkflow-core/src/engine/mod.rs b/crates/arkflow-core/src/engine/mod.rs index e5f8d3a0..ef5ab9d2 100644 --- a/crates/arkflow-core/src/engine/mod.rs +++ b/crates/arkflow-core/src/engine/mod.rs @@ -16,7 +16,6 @@ use crate::config::EngineConfig; use std::process; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tokio::signal::unix::{signal, SignalKind}; use tokio_util::sync::CancellationToken; use tracing::{error, info}; @@ -217,9 +216,7 @@ impl Engine { /// 5. Waits for all streams to complete /// /// Returns an error if any part of the initialization or execution fails - pub async fn run(&self) -> Result<(), Box> { - let token = CancellationToken::new(); - + pub async fn run(&self, token: CancellationToken) -> Result<(), Box> { // Start the health check server self.start_health_check_server(token.clone()).await?; @@ -243,23 +240,6 @@ impl Engine { // Set the readiness status self.health_state.is_ready.store(true, Ordering::SeqCst); - // Set up signal handlers - let mut sigint = signal(SignalKind::interrupt()).expect("Failed to set signal handler"); - let mut sigterm = signal(SignalKind::terminate()).expect("Failed to set signal handler"); - let token_clone = token.clone(); - tokio::spawn(async move { - tokio::select! { - _ = sigint.recv() => { - info!("Received SIGINT, exiting..."); - - }, - _ = sigterm.recv() => { - info!("Received SIGTERM, exiting..."); - } - } - - token_clone.cancel(); - }); for (i, mut stream) in streams.into_iter().enumerate() { info!("Starting flow #{}", i + 1); diff --git a/crates/arkflow-core/src/lib.rs b/crates/arkflow-core/src/lib.rs index c7f73d0a..570b2485 100644 --- a/crates/arkflow-core/src/lib.rs +++ b/crates/arkflow-core/src/lib.rs @@ -35,6 +35,7 @@ pub mod input; pub mod output; pub mod pipeline; pub mod processor; +pub mod remote_config; pub mod stream; pub mod temporary; diff --git a/crates/arkflow-core/src/remote_config.rs b/crates/arkflow-core/src/remote_config.rs new file mode 100644 index 00000000..fe44a2da --- /dev/null +++ b/crates/arkflow-core/src/remote_config.rs @@ -0,0 +1,315 @@ +/* + * 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. + */ + +//! Remote Configuration Management Module +//! +//! This module provides functionality to automatically fetch configuration from remote APIs +//! and manage stream processing pipelines dynamically. + +use crate::config::{EngineConfig, LoggingConfig}; +use crate::stream::StreamConfig; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::RwLock; +use tokio::time::interval; +use tokio_util::sync::CancellationToken; +use tracing::{error, info, warn}; + +/// Remote configuration response structure +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RemoteConfigResponse { + /// Configuration version for change detection + pub version: String, + /// List of pipeline configurations + pub streams: Vec, +} + +/// Stream information from remote API +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamInfo { + /// Unique pipeline identifier + pub id: String, + /// Pipeline name + pub name: String, + /// Pipeline status (active, inactive, deleted) + pub status: StreamStatus, + /// Stream configuration + pub config: StreamConfig, + /// Configuration version + pub version: String, +} + +/// Steam status enumeration +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum StreamStatus { + Active, + Inactive, + Deleted, +} + +/// Stream runtime information +#[derive(Debug)] +struct StreamRuntime { + /// Pipeline information + info: StreamInfo, + /// Cancellation token for stopping the pipeline + cancellation_token: CancellationToken, + /// Task handle + handle: Option>, +} + +/// Remote configuration manager +pub struct RemoteConfigManager { + /// Remote API endpoint URL + api_url: String, + /// Polling interval in seconds + poll_interval: u64, + /// Authentication token + auth_token: Option, + /// HTTP client + client: reqwest::Client, + /// Currently running pipelines + streams: Arc>>, + /// Last known configuration version + last_version: Arc>>, +} + +impl RemoteConfigManager { + /// Create a new remote configuration manager + pub fn new(api_url: String, poll_interval: u64, auth_token: Option) -> Self { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .expect("Failed to create HTTP client"); + + Self { + api_url, + poll_interval, + auth_token, + client, + streams: Arc::new(RwLock::new(HashMap::new())), + last_version: Arc::new(RwLock::new(None)), + } + } + + /// Start the remote configuration management loop + pub async fn run(&self, token: CancellationToken) -> Result<(), Box> { + info!("Starting remote configuration manager"); + info!("Polling interval: {} seconds", self.poll_interval); + info!("API endpoint: {}", self.api_url); + + // Initialize default logging + self.init_default_logging(); + + let mut interval_timer = interval(Duration::from_secs(self.poll_interval)); + + loop { + tokio::select! { + _ = token.cancelled() => { + info!("Shutting down remote configuration manager"); + break; + } + _ = interval_timer.tick() => { + if let Err(e) = self.fetch_and_update_config().await { + error!("Failed to fetch remote configuration: {}", e); + } + } + } + } + Ok(()) + } + + /// Initialize default logging configuration + fn init_default_logging(&self) { + let default_config = EngineConfig { + streams: vec![], + logging: LoggingConfig::default(), + health_check: crate::config::HealthCheckConfig::default(), + }; + crate::cli::init_logging(&default_config); + } + + /// Fetch configuration from remote API and update pipelines + async fn fetch_and_update_config(&self) -> Result<(), Box> { + let config = self.fetch_remote_config().await?; + + // Check if configuration has changed + let last_version = self.last_version.read().await; + if let Some(ref last_ver) = *last_version { + if last_ver == &config.version { + // No changes, skip update + return Ok(()); + } + } + drop(last_version); + + info!( + "Configuration changed, updating pipelines (version: {})", + config.version + ); + + // Update pipelines + self.update_streams(config.streams).await?; + + // Update version + let mut last_version = self.last_version.write().await; + *last_version = Some(config.version); + + Ok(()) + } + + /// Fetch configuration from remote API + async fn fetch_remote_config( + &self, + ) -> Result> { + let mut request = self.client.get(&self.api_url); + + // Add authentication header if token is provided + if let Some(ref token) = self.auth_token { + request = request.header("Authorization", format!("Bearer {}", token)); + } + + let response = request.send().await?; + + if !response.status().is_success() { + return Err(format!("HTTP error: {}", response.status()).into()); + } + + let config: RemoteConfigResponse = response.json().await?; + Ok(config) + } + + /// Update pipelines based on remote configuration + async fn update_streams( + &self, + new_streams: Vec, + ) -> Result<(), Box> { + let mut streams = self.streams.write().await; + let mut new_streams_ids = std::collections::HashSet::new(); + + // Process new/updated pipelines + for stream_info in new_streams { + new_streams_ids.insert(stream_info.id.clone()); + + match stream_info.status { + StreamStatus::Active => { + if let Some(existing) = streams.get(&stream_info.id) { + // Check if pipeline needs to be restarted + if existing.info.version != stream_info.version { + info!( + "Restarting pipeline '{}' (version: {} -> {})", + stream_info.name, existing.info.version, stream_info.version + ); + + // Stop existing stream + existing.cancellation_token.cancel(); + if let Some(handle) = &existing.handle { + let _ = handle.abort(); + } + + // Start new stream + self.start_stream(&mut streams, stream_info).await?; + } + } else { + // Start new stream + info!("Starting new stream '{}'", stream_info.name); + self.start_stream(&mut streams, stream_info).await?; + } + } + StreamStatus::Inactive => { + if let Some(existing) = streams.get(&stream_info.id) { + info!("Stopping stream '{}'", stream_info.name); + existing.cancellation_token.cancel(); + if let Some(handle) = &existing.handle { + let _ = handle.abort(); + } + streams.remove(&stream_info.id); + } + } + StreamStatus::Deleted => { + if let Some(existing) = streams.get(&stream_info.id) { + info!("Deleting stream '{}'", stream_info.name); + existing.cancellation_token.cancel(); + if let Some(handle) = &existing.handle { + let _ = handle.abort(); + } + streams.remove(&stream_info.id); + } + } + } + } + + // Remove stream that are no longer in the configuration + let current_ids: Vec = streams.keys().cloned().collect(); + for id in current_ids { + if !new_streams_ids.contains(&id) { + if let Some(existing) = streams.get(&id) { + warn!( + "Removing stream '{}' (no longer in remote config)", + existing.info.name + ); + existing.cancellation_token.cancel(); + if let Some(handle) = &existing.handle { + let _ = handle.abort(); + } + } + streams.remove(&id); + } + } + + Ok(()) + } + + /// Start a new stream + async fn start_stream( + &self, + streams: &mut HashMap, + stream_info: StreamInfo, + ) -> Result<(), Box> { + let cancellation_token = CancellationToken::new(); + let token_clone = cancellation_token.clone(); + let config_clone = stream_info.config.clone(); + let stream_name = stream_info.name.clone(); + let stream_id = stream_info.id.clone(); + + // Build and start the stream + let handle = tokio::spawn(async move { + match config_clone.build() { + Ok(mut stream) => { + info!("Stream '{}' started successfully", stream_name); + if let Err(e) = stream.run(token_clone).await { + error!("Stream '{}' error: {}", stream_name, e); + } else { + info!("Stream '{}' completed", stream_name); + } + } + Err(e) => { + error!("Failed to build stream '{}': {}", stream_name, e); + } + } + }); + + let runtime = StreamRuntime { + info: stream_info, + cancellation_token, + handle: Some(handle), + }; + + streams.insert(stream_id, runtime); + Ok(()) + } +}