diff --git a/Cargo.lock b/Cargo.lock index 80e9c146e34d6..ab725634370b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1408,6 +1408,7 @@ dependencies = [ "anstyle", "bitflags 1.3.2", "clap_lex 0.4.1", + "once_cell", "strsim", ] @@ -5765,6 +5766,8 @@ dependencies = [ "risingwave_frontend", "risingwave_meta", "risingwave_rt", + "strum", + "strum_macros", "task_stats_alloc", "tempfile", "tikv-jemallocator", diff --git a/Makefile.toml b/Makefile.toml index 0f455533b8382..65d7ab742f913 100644 --- a/Makefile.toml +++ b/Makefile.toml @@ -461,10 +461,11 @@ script = ''' set -ex -RUST_BACKTRACE=1 RW_NODE=playground \ +RUST_BACKTRACE=1 \ cargo run --bin risingwave \ --profile "${RISINGWAVE_BUILD_PROFILE}" \ - ${RISINGWAVE_FEATURE_FLAGS} + ${RISINGWAVE_FEATURE_FLAGS} \ + -- playground ''' [tasks.ctl] diff --git a/src/cmd_all/Cargo.toml b/src/cmd_all/Cargo.toml index 759ba7d77c4b2..10ccdeee23a9d 100644 --- a/src/cmd_all/Cargo.toml +++ b/src/cmd_all/Cargo.toml @@ -20,7 +20,7 @@ ignored = ["workspace-hack", "workspace-config", "task_stats_alloc"] [dependencies] anyhow = "1" -clap = { version = "4", features = ["derive"] } +clap = { version = "4", features = ["cargo", "derive"] } console = "0.15" risingwave_common = { path = "../common" } risingwave_compactor = { path = "../storage/compactor" } @@ -29,6 +29,8 @@ risingwave_ctl = { path = "../ctl" } risingwave_frontend = { path = "../frontend" } risingwave_meta = { path = "../meta" } risingwave_rt = { path = "../utils/runtime" } +strum = "0.24" +strum_macros = "0.24" tempfile = "3" tokio = { version = "0.2", package = "madsim-tokio", features = [ "rt", diff --git a/src/cmd_all/src/bin/risingwave.rs b/src/cmd_all/src/bin/risingwave.rs index b2389c0939c9f..8290e3c6ea88a 100644 --- a/src/cmd_all/src/bin/risingwave.rs +++ b/src/cmd_all/src/bin/risingwave.rs @@ -13,168 +13,152 @@ // limitations under the License. #![cfg_attr(coverage, feature(no_coverage))] -#![feature(let_chains)] -use std::collections::HashMap; -use std::env; +use std::str::FromStr; -use anyhow::{bail, Result}; -use clap::Parser; -use risingwave_cmd_all::playground; -#[cfg(enable_task_local_alloc)] -use risingwave_common::enable_task_local_jemalloc_on_unix; +use anyhow::Result; +use clap::{command, Arg, Command, Parser}; +use strum::IntoEnumIterator; +use strum_macros::{Display, EnumIter, EnumString, IntoStaticStr}; use tracing::Level; #[cfg(enable_task_local_alloc)] -enable_task_local_jemalloc_on_unix!(); - -#[cfg(not(enable_task_local_alloc))] -use risingwave_common::enable_jemalloc_on_unix; +risingwave_common::enable_task_local_jemalloc_on_unix!(); #[cfg(not(enable_task_local_alloc))] -enable_jemalloc_on_unix!(); - -type RwFns = HashMap<&'static str, Box) -> Result<()>>>; - -#[cfg_attr(coverage, no_coverage)] -fn main() -> Result<()> { - let mut fns: RwFns = HashMap::new(); - - // compute node configuration - for fn_name in ["compute", "compute-node", "compute_node"] { - fns.insert( - fn_name, - Box::new(|args: Vec| { - eprintln!("launching compute node"); - - let opts = risingwave_compute::ComputeNodeOpts::parse_from(args); - - risingwave_rt::init_risingwave_logger( - risingwave_rt::LoggerSettings::new().enable_tokio_console(false), - ); - - risingwave_rt::main_okk(risingwave_compute::start(opts)); - - Ok(()) - }), - ); - } - - // meta node configuration - for fn_name in ["meta", "meta-node", "meta_node"] { - fns.insert( - fn_name, - Box::new(move |args: Vec| { - eprintln!("launching meta node"); - - let opts = risingwave_meta::MetaNodeOpts::parse_from(args); - - risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); - - risingwave_rt::main_okk(risingwave_meta::start(opts)); +risingwave_common::enable_jemalloc_on_unix!(); + +const BINARY_NAME: &str = "risingwave"; +const ARGS_ID: &str = "args"; + +/// Component to lanuch. +#[derive(Clone, Copy, EnumIter, EnumString, Display, IntoStaticStr)] +#[strum(serialize_all = "snake_case")] +enum Component { + Compute, + Meta, + Frontend, + Compactor, + Ctl, + Playground, +} - Ok(()) - }), - ); +impl Component { + /// Start the component from the given `args` without `argv[0]`. + fn start(self, mut args: Vec) { + eprintln!("launching `{}` with args `{:?}`", self, args); + args.insert(0, format!("{} {}", BINARY_NAME, self)); // mock argv[0] + + match self { + Self::Compute => compute(args), + Self::Meta => meta(args), + Self::Frontend => frontend(args), + Self::Compactor => compactor(args), + Self::Ctl => ctl(args), + Self::Playground => playground(args), + } } - // frontend node configuration - for fn_name in ["frontend", "frontend-node", "frontend_node"] { - fns.insert( - fn_name, - Box::new(move |args: Vec| { - eprintln!("launching frontend node"); - - let opts = risingwave_frontend::FrontendOpts::parse_from(args); - - risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); - - risingwave_rt::main_okk(risingwave_frontend::start(opts)); - - Ok(()) - }), - ); + /// Aliases that can be used to launch the component. + fn aliases(self) -> Vec<&'static str> { + match self { + Component::Compute => vec!["compute-node", "compute_node"], + Component::Meta => vec!["meta-node", "meta_node"], + Component::Frontend => vec!["frontend-node", "frontend_node"], + Component::Compactor => vec!["compactor-node", "compactor_node"], + Component::Ctl => vec!["risectl"], + Component::Playground => vec!["play"], + } } - // compactor node configuration - for fn_name in ["compactor", "compactor-node", "compactor_node"] { - fns.insert( - fn_name, - Box::new(move |args: Vec| { - eprintln!("launching compactor node"); - - let opts = risingwave_compactor::CompactorOpts::parse_from(args); - - risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); - - risingwave_rt::main_okk(risingwave_compactor::start(opts)); - - Ok(()) - }), - ); + /// `clap` commands for all components. + fn commands() -> Vec { + Self::iter() + .map(|c| { + let name: &'static str = c.into(); + let args = Arg::new(ARGS_ID) + // make arguments transaprent to `clap` + .num_args(0..) + .allow_hyphen_values(true) + .trailing_var_arg(true); + Command::new(name).visible_aliases(c.aliases()).arg(args) + }) + .collect() } +} - // risectl - for fn_name in ["ctl", "risectl"] { - fns.insert( - fn_name, - Box::new(move |args: Vec| { - eprintln!("launching risectl"); - - let opts = risingwave_ctl::CliOpts::parse_from(args); - risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); +fn main() -> Result<()> { + let risingwave = || command!(BINARY_NAME); + let command = risingwave() + // `$ ./meta ` + .multicall(true) + .subcommands(Component::commands()) + // `$ ./risingwave meta ` + .subcommand( + risingwave() + .subcommand_value_name("COMPONENT") + .subcommand_help_heading("Components") + .subcommand_required(true) + .subcommands(Component::commands()), + ) + .disable_help_flag(true); // avoid top-level options + + let matches = command.get_matches(); + + let multicall = matches.subcommand().unwrap(); + let argv_1 = multicall.1.subcommand(); + let subcommand = argv_1.unwrap_or(multicall); + + let component = Component::from_str(subcommand.0)?; + let args = subcommand + .1 + .get_many::(ARGS_ID) + .into_iter() + .flatten() + .cloned() + .collect(); + + component.start(args); - risingwave_rt::main_okk(risingwave_ctl::start(opts)) - }), - ); - } + Ok(()) +} - // playground - for fn_name in ["play", "playground"] { - fns.insert( - fn_name, - Box::new(move |_: Vec| { - let settings = risingwave_rt::LoggerSettings::new() - .enable_tokio_console(false) - .with_target("risingwave_storage", Level::WARN); - risingwave_rt::init_risingwave_logger(settings); - - risingwave_rt::main_okk(playground()) - }), - ); - } +fn compute(args: Vec) { + let opts = risingwave_compute::ComputeNodeOpts::parse_from(args); + risingwave_rt::init_risingwave_logger( + risingwave_rt::LoggerSettings::new().enable_tokio_console(false), + ); + risingwave_rt::main_okk(risingwave_compute::start(opts)); +} - /// Get the launch target of this all-in-one binary - fn get_target(cmds: Vec<&str>) -> (String, Vec) { - if let Some(cmd) = env::args().nth(1) && cmds.contains(&cmd.as_str()) { - // ./risingwave meta - return (cmd, env::args().skip(1).collect()); - } +fn meta(args: Vec) { + let opts = risingwave_meta::MetaNodeOpts::parse_from(args); + risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); + risingwave_rt::main_okk(risingwave_meta::start(opts)); +} - if let Ok(target) = env::var("RW_NODE") { - // RW_NODE=meta ./risingwave - (target, env::args().collect()) - } else { - // ./meta-node - let x = env::args().next().expect("cannot find argv[0]"); - let x = x.rsplit('/').next().expect("cannot find binary name"); - let target = x.to_string(); - (target, env::args().collect()) - } - } +fn frontend(args: Vec) { + let opts = risingwave_frontend::FrontendOpts::parse_from(args); + risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); + risingwave_rt::main_okk(risingwave_frontend::start(opts)); +} - let (target, args) = get_target(fns.keys().copied().collect()); +fn compactor(args: Vec) { + let opts = risingwave_compactor::CompactorOpts::parse_from(args); + risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); + risingwave_rt::main_okk(risingwave_compactor::start(opts)); +} - match fns.remove(target.as_str()) { - Some(func) => { - func(args)?; - } - None => { - let mut components = fns.keys().collect::>(); - components.sort(); - bail!("unknown target: {}\nPlease either:\n* set `RW_NODE` env variable (`RW_NODE=`)\n* create a symbol link to `risingwave` binary (ln -s risingwave )\n* start with subcommand `risingwave ``\nwith one of the following: {:?}", target, components); - } - } +fn ctl(args: Vec) { + let opts = risingwave_ctl::CliOpts::parse_from(args); + risingwave_rt::init_risingwave_logger(risingwave_rt::LoggerSettings::new()); + risingwave_rt::main_okk(risingwave_ctl::start(opts)).unwrap(); +} - Ok(()) +fn playground(_args: Vec) { + let settings = risingwave_rt::LoggerSettings::new() + .enable_tokio_console(false) + .with_target("risingwave_storage", Level::WARN); + risingwave_rt::init_risingwave_logger(settings); + risingwave_rt::main_okk(risingwave_cmd_all::playground()).unwrap() }