From f4a434cad84084a5017ab826d690ebed7b432c07 Mon Sep 17 00:00:00 2001 From: QuakeWang <1677980708@qq.com> Date: Fri, 29 Aug 2025 14:53:03 +0800 Subject: [PATCH] feat: add BE and FE host selection tools and enhance HTTP client functionality --- src/tools/be/be_http_client.rs | 75 ++++++++++++++++++++++++------ src/tools/be/list.rs | 59 +++++++++++++++++++++++ src/tools/be/mod.rs | 2 + src/tools/common/host_selection.rs | 29 ++++++++++++ src/tools/common/mod.rs | 1 + src/tools/fe/list.rs | 53 +++++++++++++++++++++ src/tools/fe/mod.rs | 2 + src/tools/fe/routine_load/mod.rs | 6 +-- src/tools/mod.rs | 11 +++-- src/tools/mysql/cluster.rs | 19 +++++++- src/ui/menu.rs | 26 +++++++---- src/ui/service_handlers.rs | 17 +++++-- 12 files changed, 266 insertions(+), 34 deletions(-) create mode 100644 src/tools/be/list.rs create mode 100644 src/tools/common/host_selection.rs create mode 100644 src/tools/fe/list.rs diff --git a/src/tools/be/be_http_client.rs b/src/tools/be/be_http_client.rs index 1142308..36de62c 100644 --- a/src/tools/be/be_http_client.rs +++ b/src/tools/be/be_http_client.rs @@ -1,17 +1,41 @@ use crate::config_loader; use crate::error::{CliError, Result}; use crate::executor; +use crate::tools::{be, mysql}; use crate::ui; +use std::collections::BTreeSet; use std::process::Command; const BE_DEFAULT_IP: &str = "127.0.0.1"; /// Send an HTTP GET request to a BE API endpoint pub fn request_be_webserver_port(endpoint: &str, filter_pattern: Option<&str>) -> Result { - let be_http_ports = get_be_http_ports()?; + let mut be_targets: BTreeSet<(String, u16)> = BTreeSet::new(); - for &port in &be_http_ports { - let url = format!("http://{BE_DEFAULT_IP}:{port}{endpoint}"); + let ports = get_be_http_ports()?; + + let selected_host = be::list::get_selected_be_host(); + + let cluster_hosts = get_be_ip().unwrap_or_default(); + + let mut all_hosts = BTreeSet::new(); + if let Some(host) = &selected_host { + all_hosts.insert(host.clone()); + } + for host in cluster_hosts { + all_hosts.insert(host); + } + + if all_hosts.is_empty() { + all_hosts.insert(BE_DEFAULT_IP.to_string()); + } + + for host in all_hosts { + be_targets.extend(ports.iter().map(|p| (host.clone(), *p))); + } + + for (host, port) in &be_targets { + let url = format!("http://{host}:{port}{endpoint}"); let mut curl_cmd = Command::new("curl"); curl_cmd.args(["-sS", &url]); @@ -31,12 +55,15 @@ pub fn request_be_webserver_port(endpoint: &str, filter_pattern: Option<&str>) - } } - let ports_str = be_http_ports + let ports_str = be_targets .iter() - .map(|p| p.to_string()) + .map(|(h, p)| format!("{h}:{p}")) .collect::>() .join(", "); + ui::print_warning( + "Could not connect to any BE http endpoint. You can select a host via 'be-list'.", + ); Err(CliError::ToolExecutionFailed(format!( "Could not connect to any BE http port ({ports_str}). Check if BE is running." ))) @@ -44,14 +71,36 @@ pub fn request_be_webserver_port(endpoint: &str, filter_pattern: Option<&str>) - /// Get BE HTTP ports from configuration or use defaults pub fn get_be_http_ports() -> Result> { - match config_loader::load_config() { - Ok(doris_config) => Ok(doris_config.get_be_http_ports()), - Err(_) => { - // Fallback to default ports if configuration cannot be loaded - ui::print_warning( - "Could not load configuration, using default BE HTTP ports (8040, 8041)", - ); - Ok(vec![8040, 8041]) + if let Ok(doris_config) = config_loader::load_config() { + let config_ports = doris_config.get_be_http_ports(); + if !config_ports.is_empty() && config_ports != vec![8040, 8041] { + return Ok(config_ports); } } + + if let Ok(info) = mysql::ClusterInfo::load_from_file() { + let be_ports: Vec = info + .backends + .iter() + .filter(|b| b.alive) + .map(|b| b.http_port) + .collect(); + + if !be_ports.is_empty() { + return Ok(be_ports); + } + } + + Ok(vec![8040, 8041]) +} + +pub fn get_be_ip() -> Result> { + if let Ok(info) = mysql::ClusterInfo::load_from_file() { + let hosts = info.list_be_hosts(); + if !hosts.is_empty() { + return Ok(hosts); + } + } + + Ok(vec![BE_DEFAULT_IP.to_string()]) } diff --git a/src/tools/be/list.rs b/src/tools/be/list.rs new file mode 100644 index 0000000..81bdf3d --- /dev/null +++ b/src/tools/be/list.rs @@ -0,0 +1,59 @@ +use crate::config::Config; +use crate::error::{CliError, Result}; +use crate::tools::{ExecutionResult, Tool}; +use crate::ui; + +pub use crate::tools::common::host_selection::{ + get_selected_host as get_selected_be_host_generic, + set_selected_host as set_selected_be_host_generic, +}; +pub fn set_selected_be_host(host: String) { + set_selected_be_host_generic(true, host); +} +pub fn get_selected_be_host() -> Option { + get_selected_be_host_generic(true) +} + +pub struct BeListTool; + +impl Tool for BeListTool { + fn name(&self) -> &str { + "be-list" + } + + fn description(&self) -> &str { + "List and select a BE host (IP) for this session" + } + + fn requires_pid(&self) -> bool { + false + } + + fn execute(&self, _config: &Config, _pid: u32) -> Result { + let info = crate::tools::mysql::ClusterInfo::load_from_file()?; + let hosts = info.list_be_hosts(); + if hosts.is_empty() { + return Err(CliError::ConfigError( + "No BE hosts found in clusters.toml".to_string(), + )); + } + + let items: Vec = hosts; + + let selection = dialoguer::Select::with_theme(&dialoguer::theme::ColorfulTheme::default()) + .with_prompt("Select Backend (BE) host") + .items(&items) + .default(0) + .interact() + .map_err(|e| CliError::InvalidInput(format!("BE selection failed: {e}")))?; + + let host = items[selection].clone(); + set_selected_be_host(host.clone()); + ui::print_success(&format!("Selected BE host: {host}")); + + Ok(ExecutionResult { + output_path: std::path::PathBuf::from("console_output"), + message: "BE host updated for this session".to_string(), + }) + } +} diff --git a/src/tools/be/mod.rs b/src/tools/be/mod.rs index 82a0f5e..9209239 100644 --- a/src/tools/be/mod.rs +++ b/src/tools/be/mod.rs @@ -1,6 +1,7 @@ mod be_http_client; mod be_vars; mod jmap; +mod list; mod memz; mod pipeline_tasks; mod pstack; @@ -8,6 +9,7 @@ mod response_handler; pub use be_vars::BeVarsTool; pub use jmap::{JmapDumpTool, JmapHistoTool}; +pub use list::BeListTool; pub use memz::{MemzGlobalTool, MemzTool}; pub use pipeline_tasks::PipelineTasksTool; pub use pstack::PstackTool; diff --git a/src/tools/common/host_selection.rs b/src/tools/common/host_selection.rs new file mode 100644 index 0000000..baa7700 --- /dev/null +++ b/src/tools/common/host_selection.rs @@ -0,0 +1,29 @@ +use once_cell::sync::OnceCell; +use std::sync::Mutex; + +static SELECTED_FE_HOST: OnceCell>> = OnceCell::new(); +static SELECTED_BE_HOST: OnceCell>> = OnceCell::new(); + +fn storage(cell: &OnceCell>>) -> &Mutex> { + cell.get_or_init(|| Mutex::new(None)) +} + +pub fn set_selected_host(is_be: bool, host: String) { + let cell = if is_be { + &SELECTED_BE_HOST + } else { + &SELECTED_FE_HOST + }; + if let Ok(mut guard) = storage(cell).lock() { + *guard = Some(host); + } +} + +pub fn get_selected_host(is_be: bool) -> Option { + let cell = if is_be { + &SELECTED_BE_HOST + } else { + &SELECTED_FE_HOST + }; + storage(cell).lock().ok().and_then(|g| g.clone()) +} diff --git a/src/tools/common/mod.rs b/src/tools/common/mod.rs index e5846ab..3f2f401 100644 --- a/src/tools/common/mod.rs +++ b/src/tools/common/mod.rs @@ -1,3 +1,4 @@ pub mod format_utils; pub mod fs_utils; +pub mod host_selection; pub mod jmap; diff --git a/src/tools/fe/list.rs b/src/tools/fe/list.rs new file mode 100644 index 0000000..103b9ee --- /dev/null +++ b/src/tools/fe/list.rs @@ -0,0 +1,53 @@ +use crate::config::Config; +use crate::error::{CliError, Result}; +use crate::tools::Tool; +use crate::ui; +use std::collections::BTreeSet; + +pub struct FeListTool; + +impl Tool for FeListTool { + fn name(&self) -> &str { + "fe-list" + } + + fn description(&self) -> &str { + "List and select a FE host (IP) for this session" + } + + fn requires_pid(&self) -> bool { + false + } + + fn execute(&self, _config: &Config, _pid: u32) -> Result { + let info = crate::tools::mysql::ClusterInfo::load_from_file()?; + let mut hosts: BTreeSet = BTreeSet::new(); + for fe in info.frontends.iter().filter(|f| f.alive) { + if !fe.host.is_empty() { + hosts.insert(fe.host.clone()); + } + } + if hosts.is_empty() { + return Err(CliError::ConfigError( + "No FE hosts found in clusters.toml".to_string(), + )); + } + let items: Vec = hosts.iter().cloned().collect(); + + let selection = dialoguer::Select::with_theme(&dialoguer::theme::ColorfulTheme::default()) + .with_prompt("Select Frontend (FE) host") + .items(&items) + .default(0) + .interact() + .map_err(|e| CliError::InvalidInput(format!("FE selection failed: {e}")))?; + + let host = items[selection].clone(); + crate::tools::common::host_selection::set_selected_host(false, host.clone()); + ui::print_success(&format!("Selected FE host: {host}")); + + Ok(crate::tools::ExecutionResult { + output_path: std::path::PathBuf::from("console_output"), + message: "FE target updated for this session".to_string(), + }) + } +} diff --git a/src/tools/fe/mod.rs b/src/tools/fe/mod.rs index 022ee3a..c388d8c 100644 --- a/src/tools/fe/mod.rs +++ b/src/tools/fe/mod.rs @@ -1,11 +1,13 @@ mod jmap; mod jstack; +mod list; mod profiler; pub mod routine_load; pub mod table_info; pub use jmap::{JmapDumpTool, JmapHistoTool}; pub use jstack::JstackTool; +pub use list::FeListTool; pub use profiler::FeProfilerTool; pub use routine_load::{RoutineLoadJobLister, get_routine_load_tools}; pub use table_info::{FeTableInfoTool, TableIdentity, TableInfoReport}; diff --git a/src/tools/fe/routine_load/mod.rs b/src/tools/fe/routine_load/mod.rs index b64862a..11b1e8d 100644 --- a/src/tools/fe/routine_load/mod.rs +++ b/src/tools/fe/routine_load/mod.rs @@ -18,9 +18,9 @@ pub use traffic_monitor::RoutineLoadTrafficMonitor; /// Routine Load tool index enum to avoid hardcoded indices #[derive(Debug, Clone, Copy)] pub enum RoutineLoadToolIndex { - JobLister = 4, - PerformanceAnalyzer = 5, - TrafficMonitor = 6, + JobLister = 5, + PerformanceAnalyzer = 6, + TrafficMonitor = 7, } impl RoutineLoadToolIndex { diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 80f0280..5578a26 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -10,9 +10,7 @@ use std::path::PathBuf; /// Result of executing a tool #[derive(Debug)] pub struct ExecutionResult { - /// Path to the generated output file pub output_path: PathBuf, - /// Success message describing the operation pub message: String, } @@ -25,7 +23,6 @@ pub trait Tool { fn execute(&self, config: &Config, pid: u32) -> Result; /// Indicates whether the tool requires a process PID to execute. - /// Most tools do, so the default is true. fn requires_pid(&self) -> bool { true } @@ -47,11 +44,13 @@ impl ToolRegistry { /// Creates a new tool registry with all available tools pub fn new() -> Self { use crate::tools::be::{ - BeVarsTool, MemzGlobalTool, MemzTool, PipelineTasksTool, PstackTool, + BeListTool, BeVarsTool, MemzGlobalTool, MemzTool, PipelineTasksTool, PstackTool, }; use crate::tools::be::{JmapDumpTool as BeJmapDumpTool, JmapHistoTool as BeJmapHistoTool}; use crate::tools::fe::routine_load::get_routine_load_tools; - use crate::tools::fe::{FeProfilerTool, JmapDumpTool, JmapHistoTool, JstackTool}; + use crate::tools::fe::{ + FeListTool, FeProfilerTool, JmapDumpTool, JmapHistoTool, JstackTool, + }; let mut registry = Self { fe_tools: Vec::new(), @@ -59,6 +58,7 @@ impl ToolRegistry { }; // Register FE tools + registry.fe_tools.push(Box::new(FeListTool)); registry.fe_tools.push(Box::new(JmapDumpTool)); registry.fe_tools.push(Box::new(JmapHistoTool)); registry.fe_tools.push(Box::new(JstackTool)); @@ -68,6 +68,7 @@ impl ToolRegistry { registry.fe_tools.extend(get_routine_load_tools()); // Register BE tools + registry.be_tools.push(Box::new(BeListTool)); registry.be_tools.push(Box::new(PstackTool)); registry.be_tools.push(Box::new(BeVarsTool)); registry.be_tools.push(Box::new(BeJmapDumpTool)); diff --git a/src/tools/mysql/cluster.rs b/src/tools/mysql/cluster.rs index 6772113..83f604b 100644 --- a/src/tools/mysql/cluster.rs +++ b/src/tools/mysql/cluster.rs @@ -138,7 +138,6 @@ impl Backend { } /// Parse Tag information and extract cloud cluster information - /// This is a private helper method specifically for Backend Tag field parsing fn parse_tag_info(tag_str: &str) -> Option { if tag_str.is_empty() || tag_str == "{}" { return None; @@ -192,6 +191,24 @@ pub struct ClusterInfo { } impl ClusterInfo { + pub fn load_from_file() -> Result { + let config_dir = fs_utils::get_user_config_dir()?; + let file_path = config_dir.join("clusters.toml"); + let content = fs_utils::read_file_content(&file_path)?; + let info: ClusterInfo = toml::from_str(&content).map_err(|e| { + crate::error::CliError::ConfigError(format!("Failed to parse clusters.toml: {e}")) + })?; + Ok(info) + } + + pub fn list_be_hosts(&self) -> Vec { + self.backends + .iter() + .filter(|b| b.alive) + .map(|b| b.host.clone()) + .collect() + } + pub fn save_to_file(&self) -> Result { self.validate()?; let config_dir = fs_utils::get_user_config_dir()?; diff --git a/src/ui/menu.rs b/src/ui/menu.rs index 391448d..5e9c152 100644 --- a/src/ui/menu.rs +++ b/src/ui/menu.rs @@ -84,7 +84,10 @@ fn show_interactive_menu(step: u8, title: &str, items: &[String]) -> Result Result { title: "Select FE tool".to_string(), options: vec![ MenuOption { - action: FeToolAction::JmapDump, + action: FeToolAction::FeList, key: "[1]".to_string(), + name: "fe-list".to_string(), + description: "List and select FE host (IP)".to_string(), + }, + MenuOption { + action: FeToolAction::JmapDump, + key: "[2]".to_string(), name: "jmap-dump".to_string(), description: "Generate heap dump (.hprof)".to_string(), }, MenuOption { action: FeToolAction::JmapHisto, - key: "[2]".to_string(), + key: "[3]".to_string(), name: "jmap-histo".to_string(), description: "Generate histogram (.log)".to_string(), }, MenuOption { action: FeToolAction::Jstack, - key: "[3]".to_string(), + key: "[4]".to_string(), name: "jstack".to_string(), description: "Generate thread stack trace (.log)".to_string(), }, MenuOption { action: FeToolAction::FeProfiler, - key: "[4]".to_string(), + key: "[5]".to_string(), name: "fe-profiler".to_string(), description: "Generate flame graph for FE performance analysis using async-profiler" @@ -194,19 +204,19 @@ pub fn show_fe_tools_menu() -> Result { }, MenuOption { action: FeToolAction::TableInfo, - key: "[5]".to_string(), + key: "[6]".to_string(), name: "table-info".to_string(), description: "Collect table info for a selected table".to_string(), }, MenuOption { action: FeToolAction::RoutineLoad, - key: "[6]".to_string(), + key: "[7]".to_string(), name: "routine-load".to_string(), description: "Routine Load management tools".to_string(), }, MenuOption { action: FeToolAction::Back, - key: "[7]".to_string(), + key: "[8]".to_string(), name: "← Back".to_string(), description: "Return to main menu".to_string(), }, diff --git a/src/ui/service_handlers.rs b/src/ui/service_handlers.rs index 0f70f46..08995e1 100644 --- a/src/ui/service_handlers.rs +++ b/src/ui/service_handlers.rs @@ -20,8 +20,17 @@ pub fn handle_service_loop( pub fn handle_fe_service_loop(config: &Config, tools: &[Box]) -> Result<()> { loop { match crate::ui::show_fe_tools_menu()? { + crate::ui::FeToolAction::FeList => { + let tool = &*tools[0]; + if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { + match e { + error::CliError::GracefulExit => {} + _ => print_error(&format!("Tool execution failed: {e}")), + } + } + } crate::ui::FeToolAction::JmapDump => { - let tool = &*tools[0]; // jmap-dump + let tool = &*tools[1]; if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { match e { error::CliError::GracefulExit => { /* Do nothing, just loop again */ } @@ -38,7 +47,7 @@ pub fn handle_fe_service_loop(config: &Config, tools: &[Box]) -> Resul } } crate::ui::FeToolAction::JmapHisto => { - let tool = &*tools[1]; // jmap-histo + let tool = &*tools[2]; if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { match e { error::CliError::GracefulExit => { /* Do nothing, just loop again */ } @@ -55,7 +64,7 @@ pub fn handle_fe_service_loop(config: &Config, tools: &[Box]) -> Resul } } crate::ui::FeToolAction::Jstack => { - let tool = &*tools[2]; // jstack + let tool = &*tools[3]; if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { match e { error::CliError::GracefulExit => { /* Do nothing, just loop again */ } @@ -72,7 +81,7 @@ pub fn handle_fe_service_loop(config: &Config, tools: &[Box]) -> Resul } } crate::ui::FeToolAction::FeProfiler => { - let tool = &*tools[3]; // fe-profiler + let tool = &*tools[4]; if let Err(e) = crate::execute_tool_enhanced(config, tool, "FE") { match e { error::CliError::GracefulExit => { /* Do nothing, just loop again */ }