From 531d5f07532b4c349b569c9742da06466b303a1c Mon Sep 17 00:00:00 2001 From: QuakeWang <1677980708@qq.com> Date: Tue, 26 Aug 2025 14:26:27 +0800 Subject: [PATCH] feat: introduce AppState management and background task handling --- src/core/app_state.rs | 53 ++ src/core/background_tasks.rs | 97 ++++ src/core/mod.rs | 5 + src/lib.rs | 535 +----------------- src/tools/be/be_vars.rs | 1 - src/tools/fe/routine_load/job_lister.rs | 6 +- .../fe/routine_load/performance_analyzer.rs | 5 +- src/tools/fe/routine_load/traffic_monitor.rs | 3 +- src/ui/dialogs.rs | 26 +- src/ui/error_handlers.rs | 181 ++++++ src/ui/mod.rs | 6 + src/ui/selector.rs | 3 +- src/ui/service_handlers.rs | 182 ++++++ src/ui/tool_executor.rs | 65 +++ 14 files changed, 644 insertions(+), 524 deletions(-) create mode 100644 src/core/app_state.rs create mode 100644 src/core/background_tasks.rs create mode 100644 src/core/mod.rs create mode 100644 src/ui/error_handlers.rs create mode 100644 src/ui/service_handlers.rs create mode 100644 src/ui/tool_executor.rs diff --git a/src/core/app_state.rs b/src/core/app_state.rs new file mode 100644 index 0000000..8e5a2f2 --- /dev/null +++ b/src/core/app_state.rs @@ -0,0 +1,53 @@ +use crate::config::Config; +use crate::config_loader; +use crate::tools::ToolRegistry; + +pub struct AppState { + pub config: Config, + pub doris_config: crate::config_loader::DorisConfig, + pub registry: ToolRegistry, + pub background_handle: Option>, +} + +impl AppState { + pub fn new() -> crate::error::Result { + let doris_config = config_loader::load_config()?; + let config = config_loader::to_app_config(doris_config.clone()); + let registry = ToolRegistry::new(); + + Ok(Self { + config, + doris_config, + registry, + background_handle: None, + }) + } + + pub fn spawn_background_tasks_if_needed(&mut self) { + let fe_process_exists = + config_loader::process_detector::get_pid_by_env(config_loader::Environment::FE).is_ok(); + let has_mysql = self.doris_config.mysql.is_some(); + if fe_process_exists && has_mysql { + self.background_handle = + Some(crate::core::background_tasks::spawn_cluster_info_collector( + self.doris_config.clone(), + )); + } + } + + pub fn update_config(&mut self, new_config: Config) { + self.config = new_config.clone(); + self.doris_config = self.doris_config.clone().with_app_config(&new_config); + config_loader::persist_configuration(&self.doris_config); + } + + pub fn reset_runtime_config(&mut self) { + self.config = Config::new(); + } + + pub fn cleanup(&mut self) { + if let Some(handle) = self.background_handle.take() { + let _ = handle.join(); + } + } +} diff --git a/src/core/background_tasks.rs b/src/core/background_tasks.rs new file mode 100644 index 0000000..14cf28b --- /dev/null +++ b/src/core/background_tasks.rs @@ -0,0 +1,97 @@ +use crate::error::Result; + +/// Collect cluster info asynchronously in the background +pub fn spawn_cluster_info_collector( + doris_config: crate::config_loader::DorisConfig, +) -> std::thread::JoinHandle<()> { + std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(100)); + if should_update_cluster_info() { + collect_cluster_info_with_retry(&doris_config); + } + }) +} + +/// Collect cluster info with retry mechanism and timeout +pub fn collect_cluster_info_with_retry(doris_config: &crate::config_loader::DorisConfig) { + const MAX_RETRIES: u32 = 3; + const RETRY_DELAY_SECS: u64 = 2; + const TIMEOUT_SECS: u64 = 30; + + let start = std::time::Instant::now(); + let timeout = std::time::Duration::from_secs(TIMEOUT_SECS); + let mut retry_count = 0; + + while retry_count < MAX_RETRIES && start.elapsed() < timeout { + match collect_cluster_info_background(doris_config) { + Ok(_) => break, + Err(e) => { + retry_count += 1; + + if let crate::error::CliError::MySQLAccessDenied(_) = e { + break; + } + if let crate::error::CliError::ConfigError(_) = e { + break; + } + + if retry_count >= MAX_RETRIES || start.elapsed() >= timeout { + if std::env::var("CLOUD_CLI_DEBUG").is_ok() { + eprintln!( + "Background cluster info collection failed after {retry_count} attempts: {e}" + ); + } + break; + } else { + std::thread::sleep(std::time::Duration::from_secs(RETRY_DELAY_SECS)); + } + } + } + } +} + +/// Check if cluster info needs to be updated +pub fn should_update_cluster_info() -> bool { + let clusters_file = match dirs::home_dir() { + Some(home) => home.join(".config").join("cloud-cli").join("clusters.toml"), + None => return true, + }; + + if !clusters_file.exists() { + return true; + } + + let metadata = match std::fs::metadata(&clusters_file) { + Ok(m) => m, + Err(_) => return true, + }; + + if metadata.len() < 100 { + return true; + } + + let modified = match metadata.modified() { + Ok(m) => m, + Err(_) => return true, + }; + + let duration = match std::time::SystemTime::now().duration_since(modified) { + Ok(d) => d, + Err(_) => return true, + }; + + duration.as_secs() > 300 +} + +/// Implementation for collecting cluster info in the background +pub fn collect_cluster_info_background( + doris_config: &crate::config_loader::DorisConfig, +) -> Result<()> { + if doris_config.mysql.is_none() { + return Ok(()); + } + let mysql_tool = crate::tools::mysql::MySQLTool; + let cluster_info = mysql_tool.query_cluster_info(doris_config)?; + cluster_info.save_to_file()?; + Ok(()) +} diff --git a/src/core/mod.rs b/src/core/mod.rs new file mode 100644 index 0000000..8bfd573 --- /dev/null +++ b/src/core/mod.rs @@ -0,0 +1,5 @@ +pub mod app_state; +pub mod background_tasks; + +pub use app_state::*; +pub use background_tasks::*; diff --git a/src/lib.rs b/src/lib.rs index 426e0d1..114042a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod config; pub mod config_loader; +pub mod core; pub mod error; pub mod executor; pub mod process; @@ -7,26 +8,24 @@ pub mod tools; pub mod ui; use config::Config; -use config_loader::{load_config, persist_configuration}; +use config_loader::persist_configuration; use dialoguer::Confirm; use error::Result; -use std::thread; +use tools::Tool; use tools::mysql::CredentialManager; -use tools::{Tool, ToolRegistry}; use ui::*; /// Main CLI application runner pub fn run_cli() -> Result<()> { - let mut doris_config = load_config()?; + let mut app_state = crate::core::AppState::new()?; - let config = config_loader::to_app_config(doris_config.clone()); - if let Err(e) = config.validate() { + if let Err(e) = app_state.config.validate() { ui::print_error(&format!("Config warning: {e}")); } let fe_process_exists = config_loader::process_detector::get_pid_by_env(config_loader::Environment::FE).is_ok(); - let has_mysql = doris_config.mysql.is_some(); + let has_mysql = app_state.doris_config.mysql.is_some(); let cred_mgr = CredentialManager::new()?; if fe_process_exists @@ -39,10 +38,10 @@ pub fn run_cli() -> Result<()> { match cred_mgr.prompt_credentials_with_connection_test() { Ok((user, password)) => { let mysql_config = cred_mgr.encrypt_credentials(&user, &password)?; - doris_config.mysql = Some(mysql_config); - persist_configuration(&doris_config); + app_state.doris_config.mysql = Some(mysql_config); + persist_configuration(&app_state.doris_config); - match tools::mysql::MySQLTool.query_cluster_info(&doris_config) { + match tools::mysql::MySQLTool.query_cluster_info(&app_state.doris_config) { Ok(cluster_info) => { if let Err(e) = cluster_info.save_to_file() { ui::print_warning(&format!("Failed to save cluster info: {e}")); @@ -61,19 +60,16 @@ pub fn run_cli() -> Result<()> { } // Collect cluster info asynchronously in the background - let background_handle = if fe_process_exists && has_mysql { - Some(spawn_cluster_info_collector(doris_config.clone())) - } else { - None - }; + app_state.spawn_background_tasks_if_needed(); - let registry = ToolRegistry::new(); - let mut current_config = config; + let mut current_config = app_state.config.clone(); loop { match show_main_menu()? { MainMenuAction::Fe => { - if let Err(e) = handle_service_loop(¤t_config, "FE", registry.fe_tools()) { + if let Err(e) = + ui::handle_service_loop(¤t_config, "FE", app_state.registry.fe_tools()) + { print_error(&format!("FE service error: {e}")); if !ask_continue("Would you like to return to the main menu?")? { break; @@ -81,7 +77,9 @@ pub fn run_cli() -> Result<()> { } } MainMenuAction::Be => { - if let Err(e) = handle_service_loop(¤t_config, "BE", registry.be_tools()) { + if let Err(e) = + ui::handle_service_loop(¤t_config, "BE", app_state.registry.be_tools()) + { print_error(&format!("BE service error: {e}")); if !ask_continue("Would you like to return to the main menu?")? { break; @@ -91,507 +89,16 @@ pub fn run_cli() -> Result<()> { MainMenuAction::Exit => break, } - current_config = Config::new(); + app_state.reset_runtime_config(); + current_config = app_state.config.clone(); } - // Wait for background task to complete - if let Some(handle) = background_handle { - let _ = handle.join(); - } + app_state.cleanup(); ui::print_goodbye(); Ok(()) } -/// Collect cluster info asynchronously in the background -fn spawn_cluster_info_collector( - doris_config: crate::config_loader::DorisConfig, -) -> std::thread::JoinHandle<()> { - thread::spawn(move || { - // Delay a short time to avoid blocking main program startup - std::thread::sleep(std::time::Duration::from_millis(100)); - - // Check if cluster info needs to be updated - if should_update_cluster_info() { - collect_cluster_info_with_retry(&doris_config); - } - }) -} - -/// Collect cluster info with retry mechanism and timeout -fn collect_cluster_info_with_retry(doris_config: &crate::config_loader::DorisConfig) { - const MAX_RETRIES: u32 = 3; - const RETRY_DELAY_SECS: u64 = 2; - const TIMEOUT_SECS: u64 = 30; - - let start = std::time::Instant::now(); - let timeout = std::time::Duration::from_secs(TIMEOUT_SECS); - let mut retry_count = 0; - - while retry_count < MAX_RETRIES && start.elapsed() < timeout { - match collect_cluster_info_background(doris_config) { - Ok(_) => { - // Successfully collected, exit retry loop - break; - } - Err(e) => { - retry_count += 1; - - // Don't retry on authentication errors - if let crate::error::CliError::MySQLAccessDenied(_) = e { - break; - } - - // Don't retry on configuration errors - if let crate::error::CliError::ConfigError(_) = e { - break; - } - - if retry_count >= MAX_RETRIES || start.elapsed() >= timeout { - // Only log in debug mode and avoid excessive output - if std::env::var("CLOUD_CLI_DEBUG").is_ok() { - eprintln!( - "Background cluster info collection failed after {retry_count} attempts: {e}" - ); - } - break; - } else { - // Wait before retrying - std::thread::sleep(std::time::Duration::from_secs(RETRY_DELAY_SECS)); - } - } - } - } -} - -/// Check if cluster info needs to be updated -fn should_update_cluster_info() -> bool { - let clusters_file = match dirs::home_dir() { - Some(home) => home.join(".config").join("cloud-cli").join("clusters.toml"), - None => return true, // Unable to determine path, default to update - }; - - if !clusters_file.exists() { - return true; - } - - let metadata = match std::fs::metadata(&clusters_file) { - Ok(m) => m, - Err(_) => return true, // Unable to get metadata, default to update - }; - - if metadata.len() < 100 { - return true; - } - - let modified = match metadata.modified() { - Ok(m) => m, - Err(_) => return true, // Unable to get modification time, default to update - }; - - let duration = match std::time::SystemTime::now().duration_since(modified) { - Ok(d) => d, - Err(_) => return true, // Time error, default to update - }; - - duration.as_secs() > 300 // 5 minutes -} - -/// Implementation for collecting cluster info in the background -fn collect_cluster_info_background(doris_config: &crate::config_loader::DorisConfig) -> Result<()> { - if doris_config.mysql.is_none() { - return Ok(()); - } - let mysql_tool = tools::mysql::MySQLTool; - let cluster_info = mysql_tool.query_cluster_info(doris_config)?; - cluster_info.save_to_file()?; - Ok(()) -} - -/// Generic loop for handling a service type (FE or BE). -fn handle_service_loop(config: &Config, service_name: &str, tools: &[Box]) -> Result<()> { - if service_name == "FE" { - handle_fe_service_loop(config, tools) - } else { - handle_be_service_loop(config, tools) - } -} - -/// Handle FE service loop with nested menu structure -fn handle_fe_service_loop(config: &Config, tools: &[Box]) -> Result<()> { - loop { - match ui::show_fe_tools_menu()? { - ui::FeToolAction::JmapDump => { - let tool = &*tools[0]; // jmap-dump - if let Err(e) = execute_tool_enhanced(config, tool, "FE") { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Tool execution failed: {e}")), - } - } - match ui::show_post_execution_menu(tool.name())? { - ui::PostExecutionAction::Continue => continue, - ui::PostExecutionAction::BackToMain => return Ok(()), - ui::PostExecutionAction::Exit => { - ui::print_goodbye(); - std::process::exit(0); - } - } - } - ui::FeToolAction::JmapHisto => { - let tool = &*tools[1]; // jmap-histo - if let Err(e) = execute_tool_enhanced(config, tool, "FE") { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Tool execution failed: {e}")), - } - } - match ui::show_post_execution_menu(tool.name())? { - ui::PostExecutionAction::Continue => continue, - ui::PostExecutionAction::BackToMain => return Ok(()), - ui::PostExecutionAction::Exit => { - ui::print_goodbye(); - std::process::exit(0); - } - } - } - ui::FeToolAction::Jstack => { - let tool = &*tools[2]; // jstack - if let Err(e) = execute_tool_enhanced(config, tool, "FE") { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Tool execution failed: {e}")), - } - } - match ui::show_post_execution_menu(tool.name())? { - ui::PostExecutionAction::Continue => continue, - ui::PostExecutionAction::BackToMain => return Ok(()), - ui::PostExecutionAction::Exit => { - ui::print_goodbye(); - std::process::exit(0); - } - } - } - ui::FeToolAction::FeProfiler => { - let tool = &*tools[3]; // fe-profiler - if let Err(e) = execute_tool_enhanced(config, tool, "FE") { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Tool execution failed: {e}")), - } - } - match ui::show_post_execution_menu(tool.name())? { - ui::PostExecutionAction::Continue => continue, - ui::PostExecutionAction::BackToMain => return Ok(()), - ui::PostExecutionAction::Exit => { - ui::print_goodbye(); - std::process::exit(0); - } - } - } - ui::FeToolAction::RoutineLoad => { - if let Err(e) = handle_routine_load_loop(config, tools) { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Routine Load error: {e}")), - } - } - } - ui::FeToolAction::Back => return Ok(()), - } - } -} - -/// Handle Routine Load sub-menu loop -fn handle_routine_load_loop(config: &Config, tools: &[Box]) -> Result<()> { - loop { - match ui::show_routine_load_menu()? { - ui::RoutineLoadAction::GetJobId => execute_routine_load_tool( - config, - tools, - crate::tools::fe::routine_load::RoutineLoadToolIndex::JobLister, - )?, - - ui::RoutineLoadAction::Performance => execute_routine_load_tool( - config, - tools, - crate::tools::fe::routine_load::RoutineLoadToolIndex::PerformanceAnalyzer, - )?, - ui::RoutineLoadAction::Traffic => execute_routine_load_tool( - config, - tools, - crate::tools::fe::routine_load::RoutineLoadToolIndex::TrafficMonitor, - )?, - ui::RoutineLoadAction::Back => return Ok(()), - } - } -} - -fn execute_routine_load_tool( - config: &Config, - tools: &[Box], - tool_index: crate::tools::fe::routine_load::RoutineLoadToolIndex, -) -> Result<()> { - let tool = tool_index.get_tool(tools).ok_or_else(|| { - error::CliError::ToolExecutionFailed(format!( - "Tool not found at index {}", - tool_index as usize - )) - })?; - - if let Err(e) = execute_tool_enhanced(config, tool, "FE") { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Tool execution failed: {e}")), - } - // For failed Routine Load tools, continue loop to show menu again - return Ok(()); - } - - // Only show post execution menu for successful execution - match ui::show_post_execution_menu(tool.name())? { - ui::PostExecutionAction::Continue => Ok(()), - ui::PostExecutionAction::BackToMain => Err(error::CliError::GracefulExit), - ui::PostExecutionAction::Exit => { - ui::print_goodbye(); - std::process::exit(0); - } - } -} - -/// Handle BE service loop (original logic) -fn handle_be_service_loop(config: &Config, tools: &[Box]) -> Result<()> { - loop { - match show_tool_selection_menu(2, "Select BE tool", tools)? { - Some(tool) => { - if let Err(e) = execute_tool_enhanced(config, tool, "BE") { - match e { - error::CliError::GracefulExit => { /* Do nothing, just loop again */ } - _ => print_error(&format!("Tool execution failed: {e}")), - } - } - - match show_post_execution_menu(tool.name())? { - PostExecutionAction::Continue => continue, - PostExecutionAction::BackToMain => return Ok(()), - PostExecutionAction::Exit => { - ui::print_goodbye(); - std::process::exit(0); - } - } - } - None => return Ok(()), // "Back" was selected - } - } -} - -/// Enhanced tool execution function that uses the new configuration system fn execute_tool_enhanced(config: &Config, tool: &dyn Tool, service_name: &str) -> Result<()> { - let pid = if tool.requires_pid() { - // Try to get PID from configuration first - match config_loader::get_current_pid() { - Some(pid) => pid, - None => { - // Fallback: try to detect and select process interactively - match process::select_process_interactively() { - Ok(pid) => pid, - Err(_) => { - let tool_name = tool.name(); - print_error(&format!("No {tool_name} processes found.")); - return Ok(()); - } - } - } - } - } else { - 0 // PID is not required, provide a dummy value - }; - - print_info(&format!("Executing {}...", tool.name())); - - match tool.execute(config, pid) { - Ok(result) => { - print_success(&result.message); - if result - .output_path - .to_str() - .filter(|p| !p.is_empty() && *p != "console_output") - .is_some() - { - print_info(&format!( - "Output saved to: {}", - result.output_path.display() - )); - } - Ok(()) - } - Err(error::CliError::GracefulExit) => Ok(()), // Simply return to the menu - Err(e) => { - // Handle the error and get the potentially updated config - match handle_tool_execution_error(config, &e, service_name, tool.name())? { - Some(updated_config) => { - // Try executing the tool again with the updated config - execute_tool_enhanced(&updated_config, tool, service_name) - } - None => Ok(()), - } - } - } -} - -fn handle_tool_execution_error( - config: &Config, - error: &error::CliError, - service_name: &str, - tool_name: &str, -) -> Result> { - print_info(""); - - // Special handling for Routine Load tools when Job ID is missing - if service_name == "FE" - && tool_name.contains("routine_load") - && error.to_string().contains("No Job ID in memory") - { - print_warning("Routine Load tool execution failed: No Job ID selected."); - print_error(&format!("Error: {error}")); - - print_info(""); - print_info("Would you like to:"); - - let options = vec![ - "Go to Get Job ID".to_string(), - "Return to Routine Load menu".to_string(), - "Cancel and return to menu".to_string(), - ]; - - let selection = dialoguer::Select::with_theme(&dialoguer::theme::ColorfulTheme::default()) - .with_prompt("Choose an option") - .items(&options) - .default(0) - .interact() - .map_err(|e| { - error::CliError::InvalidInput(format!("Error fix selection failed: {e}")) - })?; - - match selection { - 0 => { - // Signal to go to Get Job ID - this will be handled by the calling loop - Err(error::CliError::GracefulExit) - } - 1 => { - // Signal to return to Routine Load menu - Err(error::CliError::GracefulExit) - } - 2 => Ok(None), - _ => Err(error::CliError::InvalidInput( - "Invalid selection".to_string(), - )), - } - } else { - // Original generic error handling for other tools - print_warning("Tool execution failed due to configuration issues."); - print_error(&format!("Error: {error}")); - - print_info(""); - print_info("Would you like to:"); - - let options = vec![ - "Fix JDK path and retry".to_string(), - "Fix output directory and retry".to_string(), - "Cancel and return to menu".to_string(), - ]; - - let selection = dialoguer::Select::with_theme(&dialoguer::theme::ColorfulTheme::default()) - .with_prompt("Choose an option") - .items(&options) - .default(0) - .interact() - .map_err(|e| { - error::CliError::InvalidInput(format!("Error fix selection failed: {e}")) - })?; - - match selection { - 0 => { - // Fix JDK path - let new_path: String = - dialoguer::Input::with_theme(&dialoguer::theme::ColorfulTheme::default()) - .with_prompt("Enter the correct JDK path") - .with_initial_text(config.jdk_path.to_string_lossy().to_string()) - .interact_text() - .map_err(|e| { - error::CliError::InvalidInput(format!("JDK path input failed: {e}")) - })?; - - let new_path = std::path::PathBuf::from(new_path); - - // Validate the new path - if !new_path.exists() { - let path_display = new_path.display(); - print_error(&format!("Path does not exist: {path_display}")); - return Ok(None); - } - - let jmap_path = new_path.join("bin/jmap"); - let jstack_path = new_path.join("bin/jstack"); - - if !jmap_path.exists() || !jstack_path.exists() { - print_error("Required JDK tools (jmap/jstack) not found in the specified path"); - return Ok(None); - } - - let fixed_config = config.clone().with_jdk_path(new_path); - - // Persist the updated configuration - if let Err(e) = persist_updated_config(&fixed_config) { - print_warning(&format!("Failed to persist configuration: {e}")); - } - - print_success("JDK path updated successfully!"); - Ok(Some(fixed_config)) - } - 1 => { - // Fix output directory - let new_path: String = - dialoguer::Input::with_theme(&dialoguer::theme::ColorfulTheme::default()) - .with_prompt("Enter the output directory path") - .with_initial_text(config.output_dir.to_string_lossy().to_string()) - .interact_text() - .map_err(|e| { - error::CliError::InvalidInput(format!("Output dir input failed: {e}")) - })?; - - let new_path = std::path::PathBuf::from(new_path); - - // Test creating the directory - if let Err(e) = std::fs::create_dir_all(&new_path) { - print_error(&format!("Cannot create directory: {e}")); - return Ok(None); - } - - let fixed_config = config.clone().with_output_dir(new_path); - - // Persist the updated configuration - if let Err(e) = persist_updated_config(&fixed_config) { - print_warning(&format!("Failed to persist configuration: {e}")); - } - - print_success("Output directory updated successfully!"); - Ok(Some(fixed_config)) - } - 2 => Ok(None), - _ => Err(error::CliError::InvalidInput( - "Invalid selection".to_string(), - )), - } - } -} - -/// Persist updated configuration to disk -fn persist_updated_config(config: &Config) -> Result<()> { - let mut doris_config = config_loader::load_config()?; - doris_config = doris_config.with_app_config(config); - match config_loader::config_persister::persist_config(&doris_config) { - Ok(_) => Ok(()), - Err(e) => Err(e), - } + ui::tool_executor::execute_tool_enhanced(config, tool, service_name) } diff --git a/src/tools/be/be_vars.rs b/src/tools/be/be_vars.rs index 892a105..8f1c340 100644 --- a/src/tools/be/be_vars.rs +++ b/src/tools/be/be_vars.rs @@ -53,7 +53,6 @@ fn prompt_for_variable_name() -> Result { if input.trim().is_empty() { ui::print_warning("Variable name cannot be empty!"); - ui::print_info("Hint: e.g., tablet_map_shard_size, or just 'shard' to search."); Ok("".to_string()) } else { Ok(input) diff --git a/src/tools/fe/routine_load/job_lister.rs b/src/tools/fe/routine_load/job_lister.rs index 4288ac4..ff23572 100644 --- a/src/tools/fe/routine_load/job_lister.rs +++ b/src/tools/fe/routine_load/job_lister.rs @@ -36,7 +36,8 @@ impl Tool for RoutineLoadJobLister { let selected_job = self.prompt_job_selection(&jobs)?; self.save_selected_job(selected_job, &database)?; let report = self.generate_selection_report(selected_job)?; - ui::print_info(&format!("\n{report}")); + ui::print_info(""); + ui::print_info(&report); return Ok(ExecutionResult { output_path: std::path::PathBuf::from("console_output"), message: format!( @@ -114,7 +115,8 @@ impl RoutineLoadJobLister { } fn display_jobs(&self, jobs: &[RoutineLoadJob]) -> Result<()> { - ui::print_info("\nRoutine Load Jobs in Database:"); + ui::print_info(""); + ui::print_info("Routine Load Jobs in Database:"); ui::print_info(&"=".repeat(100)); for job in jobs.iter() { diff --git a/src/tools/fe/routine_load/performance_analyzer.rs b/src/tools/fe/routine_load/performance_analyzer.rs index 615b3f0..f69c81d 100644 --- a/src/tools/fe/routine_load/performance_analyzer.rs +++ b/src/tools/fe/routine_load/performance_analyzer.rs @@ -174,11 +174,10 @@ impl RoutineLoadPerformanceAnalyzer { } } - // Render table - ui::print_info("\nPer-commit stats"); + ui::print_info(""); + ui::print_info("Per-commit stats"); self.print_table(&headers, &rows, &widths)?; - // Summary stats.display_summary(); Ok(()) } diff --git a/src/tools/fe/routine_load/traffic_monitor.rs b/src/tools/fe/routine_load/traffic_monitor.rs index f57efaa..49aeec0 100644 --- a/src/tools/fe/routine_load/traffic_monitor.rs +++ b/src/tools/fe/routine_load/traffic_monitor.rs @@ -122,7 +122,8 @@ impl RoutineLoadTrafficMonitor { } fn display_traffic_results(&self, per_minute_data: &BTreeMap) -> Result<()> { - ui::print_info("\nPer-minute loadedRows (ascending time)"); + ui::print_info(""); + ui::print_info("Per-minute loadedRows (ascending time)"); ui::print_info(&"-".repeat(40)); for (minute, rows) in per_minute_data.iter() { diff --git a/src/ui/dialogs.rs b/src/ui/dialogs.rs index c25ac71..a3190b9 100644 --- a/src/ui/dialogs.rs +++ b/src/ui/dialogs.rs @@ -10,8 +10,9 @@ pub enum NoJobsNextAction { } pub fn show_no_jobs_recovery_menu(database: &str) -> Result { + ui::print_info(""); ui::print_warning(&format!( - "\n[!] No Routine Load jobs found in database '{database}'" + "No Routine Load jobs found in database '{database}'" )); ui::print_info("This could mean:"); ui::print_info(" - The database name is incorrect"); @@ -35,7 +36,8 @@ pub fn show_no_jobs_recovery_menu(database: &str) -> Result { } pub fn show_unknown_db_recovery_menu(database: &str) -> Result { - ui::print_warning(&format!("\n[!] Unknown database '{database}'")); + ui::print_info(""); + ui::print_warning(&format!("Unknown database '{database}'")); ui::print_info("Please verify the database name or choose another one."); let options = vec!["Choose another database", "Back to Routine Load menu"]; @@ -54,3 +56,23 @@ pub fn show_unknown_db_recovery_menu(database: &str) -> Result Ok(action) } + +// Generic prompt helpers for reuse across UI modules +pub fn select_index(prompt: &str, options: &[&str]) -> Result { + let selection = Select::new() + .with_prompt(prompt) + .items(options) + .default(0) + .interact() + .map_err(|e| CliError::InvalidInput(e.to_string()))?; + Ok(selection) +} + +pub fn input_text(prompt: &str, initial: &str) -> Result { + let text = dialoguer::Input::with_theme(&dialoguer::theme::ColorfulTheme::default()) + .with_prompt(prompt) + .with_initial_text(initial.to_string()) + .interact_text() + .map_err(|e| CliError::InvalidInput(e.to_string()))?; + Ok(text) +} diff --git a/src/ui/error_handlers.rs b/src/ui/error_handlers.rs new file mode 100644 index 0000000..0763604 --- /dev/null +++ b/src/ui/error_handlers.rs @@ -0,0 +1,181 @@ +use crate::config::Config; +use crate::config_loader; +use crate::error::{self, Result}; +use crate::ui::{print_error, print_info, print_success, print_warning}; + +pub fn handle_tool_execution_error( + config: &Config, + error: &error::CliError, + service_name: &str, + tool_name: &str, +) -> Result> { + print_info(""); + + if service_name == "FE" + && tool_name.contains("routine_load") + && error.to_string().contains("No Job ID in memory") + { + print_warning("Routine Load tool execution failed: No Job ID selected."); + print_error(&format!("Error: {error}")); + + print_info(""); + print_info("Would you like to:"); + + let options = [ + "Go to Get Job ID".to_string(), + "Return to Routine Load menu".to_string(), + "Cancel and return to menu".to_string(), + ]; + + let options_ref: Vec<&str> = options.iter().map(|s| s.as_str()).collect(); + let selection = crate::ui::dialogs::select_index("Choose an option", &options_ref)?; + + return match selection { + 0 | 1 => Err(error::CliError::GracefulExit), + 2 => Ok(None), + _ => Err(error::CliError::InvalidInput( + "Invalid selection".to_string(), + )), + }; + } + + // BE connectivity: provide network-centric guidance instead of config fixes + if service_name == "BE" && is_be_connectivity_error(error) { + print_warning("BE connectivity issue detected."); + print_error(&format!("Error: {error}")); + + let options = ["Retry", "Return to menu"]; + let selection = crate::ui::dialogs::select_index("Choose an option", &options)?; + return match selection { + 0 => Ok(Some(config.clone())), // retry with same config + _ => Ok(None), + }; + } + + // FE profiler script missing: show simple guidance + if service_name == "FE" && is_fe_profiler_script_missing(tool_name, error) { + print_warning("FE profiler script missing."); + print_error(&format!("Error: {error}")); + + let options = ["Return to menu"]; + let _ = crate::ui::dialogs::select_index("Choose an option", &options)?; + return Ok(None); + } + + print_warning("Tool execution failed due to configuration issues."); + print_error(&format!("Error: {error}")); + + print_info(""); + print_info("Would you like to:"); + // Build options conditionally + let mut labels: Vec<&str> = Vec::new(); + type ActionFn = fn(&Config) -> Result>; + let mut actions: Vec = Vec::new(); + + if is_jdk_missing(config, error) { + labels.push("Fix JDK path and retry"); + actions.push(fix_jdk_path as ActionFn); + } + if is_output_dir_invalid(config, error) { + labels.push("Fix output directory and retry"); + actions.push(fix_output_directory as ActionFn); + } + labels.push("Cancel and return to menu"); + + let selection = crate::ui::dialogs::select_index("Choose an option", &labels)?; + if selection < actions.len() { + actions[selection](config) + } else { + Ok(None) + } +} + +fn fix_jdk_path(config: &Config) -> Result> { + let new_path: String = crate::ui::dialogs::input_text( + "Enter the correct JDK path", + &config.jdk_path.to_string_lossy(), + )?; + + let new_path = std::path::PathBuf::from(new_path); + + if !new_path.exists() { + print_error(&format!("Path does not exist: {}", new_path.display())); + return Ok(None); + } + + let jmap_path = new_path.join("bin/jmap"); + let jstack_path = new_path.join("bin/jstack"); + + if !jmap_path.exists() || !jstack_path.exists() { + print_error("Required JDK tools (jmap/jstack) not found in the specified path"); + return Ok(None); + } + + let fixed_config = config.clone().with_jdk_path(new_path); + + if let Err(e) = persist_updated_config(&fixed_config) { + print_warning(&format!("Failed to persist configuration: {e}")); + } + + print_success("JDK path updated successfully!"); + Ok(Some(fixed_config)) +} + +fn fix_output_directory(config: &Config) -> Result> { + let new_path: String = crate::ui::dialogs::input_text( + "Enter the output directory path", + &config.output_dir.to_string_lossy(), + )?; + + let new_path = std::path::PathBuf::from(new_path); + + if let Err(e) = std::fs::create_dir_all(&new_path) { + print_error(&format!("Cannot create directory: {e}")); + return Ok(None); + } + + let fixed_config = config.clone().with_output_dir(new_path); + + if let Err(e) = persist_updated_config(&fixed_config) { + print_warning(&format!("Failed to persist configuration: {e}")); + } + + print_success("Output directory updated successfully!"); + Ok(Some(fixed_config)) +} + +fn persist_updated_config(config: &Config) -> Result<()> { + let mut doris_config = config_loader::load_config()?; + doris_config = doris_config.with_app_config(config); + match config_loader::config_persister::persist_config(&doris_config) { + Ok(_) => Ok(()), + Err(e) => Err(e), + } +} + +fn is_be_connectivity_error(error: &error::CliError) -> bool { + let s = error.to_string(); + s.contains("Could not connect to any BE http port") +} + +fn is_fe_profiler_script_missing(tool_name: &str, error: &error::CliError) -> bool { + tool_name.contains("fe-profiler") && error.to_string().contains("profile_fe.sh not found") +} + +fn is_jdk_missing(config: &Config, error: &error::CliError) -> bool { + let s = error.to_string(); + if s.contains("JDK path does not exist") || s.contains("jmap") || s.contains("jstack") { + return true; + } + let jmap = config.jdk_path.join("bin/jmap"); + let jstack = config.jdk_path.join("bin/jstack"); + !jmap.exists() || !jstack.exists() +} + +fn is_output_dir_invalid(config: &Config, error: &error::CliError) -> bool { + let s = error.to_string(); + if s.contains("Cannot create directory") || s.contains("Output dir input failed") { + return true; + } + !config.output_dir.exists() +} diff --git a/src/ui/mod.rs b/src/ui/mod.rs index ac69268..fd39266 100644 --- a/src/ui/mod.rs +++ b/src/ui/mod.rs @@ -1,13 +1,19 @@ use console::{Term, style}; pub mod dialogs; +pub mod error_handlers; pub mod menu; pub mod selector; +pub mod service_handlers; +pub mod tool_executor; pub mod utils; pub use dialogs::*; +pub use error_handlers::*; pub use menu::*; pub use selector::*; +pub use service_handlers::*; +pub use tool_executor::*; pub use utils::*; pub static SUCCESS: &str = "[+] "; diff --git a/src/ui/selector.rs b/src/ui/selector.rs index 6096860..92a57fc 100644 --- a/src/ui/selector.rs +++ b/src/ui/selector.rs @@ -38,7 +38,8 @@ impl InteractiveSelector { let mut selection: usize = 0; let mut last_drawn_lines: usize; - crate::ui::print_info(&format!("\n{title}", title = self.title)); + crate::ui::print_info(""); + crate::ui::print_info(&self.title.to_string()); crate::ui::print_info("Use ↑/↓ move, ←/→ page, 1-9 jump, Enter to select:"); term.hide_cursor() diff --git a/src/ui/service_handlers.rs b/src/ui/service_handlers.rs new file mode 100644 index 0000000..57b5854 --- /dev/null +++ b/src/ui/service_handlers.rs @@ -0,0 +1,182 @@ +use crate::config::Config; +use crate::error::{self, Result}; +use crate::tools::Tool; +use crate::ui::*; + +/// Generic loop for handling a service type (FE or BE). +pub fn handle_service_loop( + config: &Config, + service_name: &str, + tools: &[Box], +) -> Result<()> { + if service_name == "FE" { + handle_fe_service_loop(config, tools) + } else { + handle_be_service_loop(config, tools) + } +} + +/// Handle FE service loop with nested menu structure +pub fn handle_fe_service_loop(config: &Config, tools: &[Box]) -> Result<()> { + loop { + match crate::ui::show_fe_tools_menu()? { + crate::ui::FeToolAction::JmapDump => { + let tool = &*tools[0]; // jmap-dump + if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Tool execution failed: {e}")), + } + } + match crate::ui::show_post_execution_menu(tool.name())? { + crate::ui::PostExecutionAction::Continue => continue, + crate::ui::PostExecutionAction::BackToMain => return Ok(()), + crate::ui::PostExecutionAction::Exit => { + crate::ui::print_goodbye(); + std::process::exit(0); + } + } + } + crate::ui::FeToolAction::JmapHisto => { + let tool = &*tools[1]; // jmap-histo + if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Tool execution failed: {e}")), + } + } + match crate::ui::show_post_execution_menu(tool.name())? { + crate::ui::PostExecutionAction::Continue => continue, + crate::ui::PostExecutionAction::BackToMain => return Ok(()), + crate::ui::PostExecutionAction::Exit => { + crate::ui::print_goodbye(); + std::process::exit(0); + } + } + } + crate::ui::FeToolAction::Jstack => { + let tool = &*tools[2]; // jstack + if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Tool execution failed: {e}")), + } + } + match crate::ui::show_post_execution_menu(tool.name())? { + crate::ui::PostExecutionAction::Continue => continue, + crate::ui::PostExecutionAction::BackToMain => return Ok(()), + crate::ui::PostExecutionAction::Exit => { + crate::ui::print_goodbye(); + std::process::exit(0); + } + } + } + crate::ui::FeToolAction::FeProfiler => { + let tool = &*tools[3]; // fe-profiler + if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Tool execution failed: {e}")), + } + } + match crate::ui::show_post_execution_menu(tool.name())? { + crate::ui::PostExecutionAction::Continue => continue, + crate::ui::PostExecutionAction::BackToMain => return Ok(()), + crate::ui::PostExecutionAction::Exit => { + crate::ui::print_goodbye(); + std::process::exit(0); + } + } + } + crate::ui::FeToolAction::RoutineLoad => { + if let Err(e) = handle_routine_load_loop(config, tools) { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Routine Load error: {e}")), + } + } + } + crate::ui::FeToolAction::Back => return Ok(()), + } + } +} + +/// Handle Routine Load sub-menu loop +pub fn handle_routine_load_loop(config: &Config, tools: &[Box]) -> Result<()> { + loop { + match crate::ui::show_routine_load_menu()? { + crate::ui::RoutineLoadAction::GetJobId => execute_routine_load_tool( + config, + tools, + crate::tools::fe::routine_load::RoutineLoadToolIndex::JobLister, + )?, + + crate::ui::RoutineLoadAction::Performance => execute_routine_load_tool( + config, + tools, + crate::tools::fe::routine_load::RoutineLoadToolIndex::PerformanceAnalyzer, + )?, + crate::ui::RoutineLoadAction::Traffic => execute_routine_load_tool( + config, + tools, + crate::tools::fe::routine_load::RoutineLoadToolIndex::TrafficMonitor, + )?, + crate::ui::RoutineLoadAction::Back => return Ok(()), + } + } +} + +fn execute_routine_load_tool( + config: &Config, + tools: &[Box], + tool_index: crate::tools::fe::routine_load::RoutineLoadToolIndex, +) -> Result<()> { + let tool = tool_index.get_tool(tools).ok_or_else(|| { + error::CliError::ToolExecutionFailed(format!( + "Tool not found at index {}", + tool_index as usize + )) + })?; + + if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Tool execution failed: {e}")), + } + return Ok(()); + } + match crate::ui::show_post_execution_menu(tool.name())? { + crate::ui::PostExecutionAction::Continue => Ok(()), + crate::ui::PostExecutionAction::BackToMain => Err(error::CliError::GracefulExit), + crate::ui::PostExecutionAction::Exit => { + crate::ui::print_goodbye(); + std::process::exit(0); + } + } +} + +/// Handle BE service loop (original logic) +pub fn handle_be_service_loop(config: &Config, tools: &[Box]) -> Result<()> { + loop { + match show_tool_selection_menu(2, "Select BE tool", tools)? { + Some(tool) => { + if let Err(e) = crate::execute_tool_enhanced(config, tool, "BE") { + match e { + error::CliError::GracefulExit => { /* Do nothing, just loop again */ } + _ => print_error(&format!("Tool execution failed: {e}")), + } + } + + match show_post_execution_menu(tool.name())? { + PostExecutionAction::Continue => continue, + PostExecutionAction::BackToMain => return Ok(()), + PostExecutionAction::Exit => { + crate::ui::print_goodbye(); + std::process::exit(0); + } + } + } + None => return Ok(()), // "Back" was selected + } + } +} diff --git a/src/ui/tool_executor.rs b/src/ui/tool_executor.rs new file mode 100644 index 0000000..90a5c8a --- /dev/null +++ b/src/ui/tool_executor.rs @@ -0,0 +1,65 @@ +use crate::config::Config; +use crate::config_loader; +use crate::error::{self, Result}; +use crate::process; +use crate::tools::Tool; +use crate::ui::{print_error, print_info, print_success}; +use std::path::Path; + +pub fn execute_tool_enhanced(config: &Config, tool: &dyn Tool, service_name: &str) -> Result<()> { + let pid = match resolve_pid_if_required(tool) { + Some(pid) => pid, + None => return Ok(()), + }; + + print_info(&format!("Executing {}...", tool.name())); + + match tool.execute(config, pid) { + Ok(result) => { + print_success(&result.message); + maybe_print_output_path(&result.output_path); + Ok(()) + } + Err(error::CliError::GracefulExit) => Ok(()), + Err(e) => { + match crate::ui::error_handlers::handle_tool_execution_error( + config, + &e, + service_name, + tool.name(), + )? { + Some(updated_config) => execute_tool_enhanced(&updated_config, tool, service_name), + None => Ok(()), + } + } + } +} + +fn resolve_pid_if_required(tool: &dyn Tool) -> Option { + if !tool.requires_pid() { + return Some(0); + } + + if let Some(pid) = config_loader::get_current_pid() { + return Some(pid); + } + + match process::select_process_interactively() { + Ok(pid) => Some(pid), + Err(_) => { + let tool_name = tool.name(); + print_error(&format!("No {tool_name} processes found.")); + None + } + } +} + +fn maybe_print_output_path(output_path: &Path) { + if output_path + .to_str() + .filter(|p| !p.is_empty() && *p != "console_output") + .is_some() + { + print_info(&format!("Output saved to: {}", output_path.display())); + } +}