Convert apply arguments to type-safe clap derive

This commit is contained in:
i1i1 2023-08-10 01:23:57 +03:00 committed by Zhaofeng Li
parent b80b57cb48
commit 87f4e3a676
6 changed files with 210 additions and 187 deletions

View file

@ -1,135 +1,126 @@
use std::env; use std::env;
use std::path::PathBuf; use std::path::PathBuf;
use std::str::FromStr;
use clap::{ use clap::{builder::ArgPredicate, ArgMatches, Args, Command as ClapCommand, FromArgMatches};
builder::{ArgPredicate, PossibleValuesParser, ValueParser},
value_parser, Arg, ArgMatches, Command as ClapCommand, FromArgMatches,
};
use crate::nix::deployment::{ use crate::nix::{
Deployment, EvaluationNodeLimit, EvaluatorType, Goal, Options, ParallelismLimit, deployment::{Deployment, EvaluationNodeLimit, EvaluatorType, Goal, Options, ParallelismLimit},
node_filter::NodeFilterOpts,
}; };
use crate::nix::NodeFilter;
use crate::progress::SimpleProgressOutput; use crate::progress::SimpleProgressOutput;
use crate::util;
use crate::{error::ColmenaError, nix::hive::HiveArgs}; use crate::{error::ColmenaError, nix::hive::HiveArgs};
pub fn register_deploy_args(command: ClapCommand) -> ClapCommand { #[derive(Debug, Args)]
command pub struct DeployOpts {
.arg(Arg::new("eval-node-limit") #[arg(
.long("eval-node-limit") value_name = "LIMIT",
.value_name("LIMIT") default_value_t,
.help("Evaluation node limit") long,
.long_help(r#"Limits the maximum number of hosts to be evaluated at once. help = "Evaluation node limit",
long_help = r#"Limits the maximum number of hosts to be evaluated at once.
The evaluation process is RAM-intensive. The default behavior is to limit the maximum number of host evaluated at the same time based on naive heuristics. The evaluation process is RAM-intensive. The default behavior is to limit the maximum number of host evaluated at the same time based on naive heuristics.
Set to 0 to disable the limit. Set to 0 to disable the limit.
"#) "#
.default_value("auto") )]
.num_args(1) eval_node_limit: EvaluationNodeLimit,
.value_parser(ValueParser::new(|s: &str| -> Result<EvaluationNodeLimit, String> { #[arg(
if s == "auto" { value_name = "LIMIT",
return Ok(EvaluationNodeLimit::Heuristic); default_value_t = 10,
} long,
short,
help = "Deploy parallelism limit",
long_help = r#"Limits the maximum number of hosts to be deployed in parallel.
match s.parse::<usize>() { Set to 0 to disable parallelism limit.
Ok(0) => Ok(EvaluationNodeLimit::None), "#
Ok(n) => Ok(EvaluationNodeLimit::Manual(n)), )]
Err(_) => Err(String::from("The value must be a valid number")), parallel: usize,
} #[arg(
}))) long,
.arg(Arg::new("parallel") help = "Create GC roots for built profiles",
.short('p') long_help = r#"Create GC roots for built profiles.
.long("parallel")
.value_name("LIMIT")
.help("Deploy parallelism limit")
.long_help(r#"Limits the maximum number of hosts to be deployed in parallel.
Set to 0 to disable parallemism limit.
"#)
.default_value("10")
.num_args(1)
.value_parser(value_parser!(usize)))
.arg(Arg::new("keep-result")
.long("keep-result")
.help("Create GC roots for built profiles")
.long_help(r#"Create GC roots for built profiles.
The built system profiles will be added as GC roots so that they will not be removed by the garbage collector. The built system profiles will be added as GC roots so that they will not be removed by the garbage collector.
The links will be created under .gcroots in the directory the Hive configuration is located. The links will be created under .gcroots in the directory the Hive configuration is located.
"#) "#
.num_args(0)) )]
.arg(Arg::new("verbose") keep_result: bool,
.short('v') #[arg(
.long("verbose") short,
.help("Be verbose") long,
.long_help("Deactivates the progress spinner and prints every line of output.") help = "Be verbose",
.num_args(0)) long_help = "Deactivates the progress spinner and prints every line of output."
.arg(Arg::new("no-keys") )]
.long("no-keys") verbose: bool,
.help("Do not upload keys") #[arg(
.long_help(r#"Do not upload secret keys set in `deployment.keys`. long,
help = "Do not upload keys",
long_help = r#"Do not upload secret keys set in `deployment.keys`.
By default, Colmena will upload keys set in `deployment.keys` before deploying the new profile on a node. By default, Colmena will upload keys set in `deployment.keys` before deploying the new profile on a node.
To upload keys without building or deploying the rest of the configuration, use `colmena upload-keys`. To upload keys without building or deploying the rest of the configuration, use `colmena upload-keys`.
"#) "#
.num_args(0)) )]
.arg(Arg::new("reboot") no_keys: bool,
.long("reboot") #[arg(
.help("Reboot nodes after activation") long,
.long_help("Reboots nodes after activation and waits for them to come back up.") help = "Reboot nodes after activation",
.num_args(0)) long_help = "Reboots nodes after activation and waits for them to come back up."
.arg(Arg::new("no-substitute") )]
.long("no-substitute") reboot: bool,
.alias("no-substitutes") #[arg(
.help("Do not use substitutes") long,
.long_help("Disables the use of substituters when copying closures to the remote host.") alias = "no-substitutes",
.num_args(0)) help = "Do not use substitutes",
.arg(Arg::new("no-gzip") long_help = "Disables the use of substituters when copying closures to the remote host."
.long("no-gzip") )]
.help("Do not use gzip") no_substitute: bool,
.long_help("Disables the use of gzip when copying closures to the remote host.") #[arg(
.num_args(0)) long,
.arg(Arg::new("build-on-target") help = "Do not use gzip",
.long("build-on-target") long_help = "Disables the use of gzip when copying closures to the remote host."
.help("Build the system profiles on the target nodes") )]
.long_help(r#"Build the system profiles on the target nodes themselves. no_gzip: bool,
#[arg(
long,
help = "Build the system profiles on the target nodes",
long_help = r#"Build the system profiles on the target nodes themselves.
If enabled, the system profiles will be built on the target nodes themselves, not on the host running Colmena itself. If enabled, the system profiles will be built on the target nodes themselves, not on the host running Colmena itself.
This overrides per-node perferences set in `deployment.buildOnTarget`. This overrides per-node perferences set in `deployment.buildOnTarget`.
To temporarily disable remote build on all nodes, use `--no-build-on-target`. To temporarily disable remote build on all nodes, use `--no-build-on-target`.
"#) "#
.num_args(0)) )]
.arg(Arg::new("no-build-on-target") build_on_target: bool,
.long("no-build-on-target") #[arg(long, hide = true)]
.hide(true) no_build_on_target: bool,
.num_args(0)) #[arg(
.arg(Arg::new("force-replace-unknown-profiles") long,
.long("force-replace-unknown-profiles") help = "Ignore all targeted nodes deployment.replaceUnknownProfiles setting",
.help("Ignore all targeted nodes deployment.replaceUnknownProfiles setting") long_help = r#"If `deployment.replaceUnknownProfiles` is set for a target, using this switch
.long_help(r#"If `deployment.replaceUnknownProfiles` is set for a target, using this switch will treat deployment.replaceUnknownProfiles as though it was set true and perform unknown profile replacement."#
will treat deployment.replaceUnknownProfiles as though it was set true and perform unknown profile replacement."#) )]
.num_args(0)) force_replace_unknown_profiles: bool,
.arg(Arg::new("evaluator") #[arg(
.long("evaluator") long,
.help("The evaluator to use (experimental)") default_value_t,
.long_help(r#"If set to `chunked` (default), evaluation of nodes will happen in batches. If set to `streaming`, the experimental streaming evaluator (nix-eval-jobs) will be used and nodes will be evaluated in parallel. help = "The evaluator to use (experimental)",
long_help = r#"If set to `chunked` (default), evaluation of nodes will happen in batches. If set to `streaming`, the experimental streaming evaluator (nix-eval-jobs) will be used and nodes will be evaluated in parallel.
This is an experimental feature."#) This is an experimental feature."#
.default_value("chunked") )]
.value_parser(value_parser!(EvaluatorType))) evaluator: EvaluatorType,
} }
pub fn subcommand() -> ClapCommand { #[derive(Debug, Args)]
let command = ClapCommand::new("apply") #[command(name = "apply", about = "Apply configurations on remote machines")]
.about("Apply configurations on remote machines") struct Opts {
.arg( #[arg(
Arg::new("goal") help = "Deployment goal",
.help("Deployment goal") long_help = r#"The goal of the deployment.
.long_help(
r#"The goal of the deployment.
Same as the targets for switch-to-configuration, with the following extra pseudo-goals: Same as the targets for switch-to-configuration, with the following extra pseudo-goals:
@ -139,24 +130,17 @@ Same as the targets for switch-to-configuration, with the following extra pseudo
`switch` is the default goal unless `--reboot` is passed, in which case `boot` is the default. `switch` is the default goal unless `--reboot` is passed, in which case `boot` is the default.
"#, "#,
) default_value_if("reboot", ArgPredicate::IsPresent, Some("boot"))
.default_value("switch") )]
.default_value_if("reboot", ArgPredicate::IsPresent, Some("boot")) goal: Goal,
.default_value("switch") #[command(flatten)]
.index(1) deploy: DeployOpts,
.value_parser(PossibleValuesParser::new([ #[command(flatten)]
"build", node_filter: NodeFilterOpts,
"push", }
"switch",
"boot",
"test",
"dry-activate",
"keys",
])),
);
let command = register_deploy_args(command);
util::register_selector_args(command) pub fn subcommand() -> ClapCommand {
Opts::augment_args(ClapCommand::new("apply"))
} }
pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(), ColmenaError> { pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(), ColmenaError> {
@ -168,13 +152,27 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
let ssh_config = env::var("SSH_CONFIG_FILE").ok().map(PathBuf::from); let ssh_config = env::var("SSH_CONFIG_FILE").ok().map(PathBuf::from);
// FIXME: Just get_one::<Goal> let Opts {
let goal_arg = local_args.get_one::<String>("goal").unwrap(); goal,
let goal = Goal::from_str(goal_arg).unwrap(); deploy:
DeployOpts {
eval_node_limit,
parallel,
keep_result,
verbose,
no_keys,
reboot,
no_substitute,
no_gzip,
build_on_target,
no_build_on_target,
force_replace_unknown_profiles,
evaluator,
},
node_filter,
} = Opts::from_arg_matches(local_args).expect("Failed to parse `apply` args");
let filter = local_args.get_one::<NodeFilter>("on"); if node_filter.on.is_none() && goal != Goal::Build {
if filter.is_none() && goal != Goal::Build {
// User did not specify node, we should check meta and see rules // User did not specify node, we should check meta and see rules
let meta = hive.get_meta_config().await?; let meta = hive.get_meta_config().await?;
if !meta.allow_apply_all { if !meta.allow_apply_all {
@ -185,11 +183,15 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
} }
let targets = hive let targets = hive
.select_nodes(filter.cloned(), ssh_config, goal.requires_target_host()) .select_nodes(
node_filter.on.clone(),
ssh_config,
goal.requires_target_host(),
)
.await?; .await?;
let n_targets = targets.len(); let n_targets = targets.len();
let verbose = local_args.get_flag("verbose") || goal == Goal::DryActivate; let verbose = verbose || goal == Goal::DryActivate;
let mut output = SimpleProgressOutput::new(verbose); let mut output = SimpleProgressOutput::new(verbose);
let progress = output.get_sender(); let progress = output.get_sender();
@ -198,27 +200,20 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
// FIXME: Configure limits // FIXME: Configure limits
let options = { let options = {
let mut options = Options::default(); let mut options = Options::default();
options.set_substituters_push(!local_args.get_flag("no-substitute")); options.set_substituters_push(!no_substitute);
options.set_gzip(!local_args.get_flag("no-gzip")); options.set_gzip(!no_gzip);
options.set_upload_keys(!local_args.get_flag("no-keys")); options.set_upload_keys(!no_keys);
options.set_reboot(local_args.get_flag("reboot")); options.set_reboot(reboot);
options.set_force_replace_unknown_profiles( options.set_force_replace_unknown_profiles(force_replace_unknown_profiles);
local_args.get_flag("force-replace-unknown-profiles"), options.set_evaluator(evaluator);
);
options.set_evaluator(
local_args
.get_one::<EvaluatorType>("evaluator")
.unwrap()
.to_owned(),
);
if local_args.get_flag("keep-result") { if keep_result {
options.set_create_gc_roots(true); options.set_create_gc_roots(true);
} }
if local_args.get_flag("no-build-on-target") { if no_build_on_target {
options.set_force_build_on_target(false); options.set_force_build_on_target(false);
} else if local_args.get_flag("build-on-target") { } else if build_on_target {
options.set_force_build_on_target(true); options.set_force_build_on_target(true);
} }
@ -227,7 +222,7 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
deployment.set_options(options); deployment.set_options(options);
if local_args.get_flag("no-keys") && goal == Goal::UploadKeys { if no_keys && goal == Goal::UploadKeys {
log::error!("--no-keys cannot be used when the goal is to upload keys"); log::error!("--no-keys cannot be used when the goal is to upload keys");
quit::with_code(1); quit::with_code(1);
} }
@ -235,23 +230,17 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
let parallelism_limit = { let parallelism_limit = {
let mut limit = ParallelismLimit::default(); let mut limit = ParallelismLimit::default();
limit.set_apply_limit({ limit.set_apply_limit({
let limit = local_args.get_one::<usize>("parallel").unwrap().to_owned(); if parallel == 0 {
if limit == 0 {
n_targets n_targets
} else { } else {
limit parallel
} }
}); });
limit limit
}; };
let evaluation_node_limit = local_args
.get_one::<EvaluationNodeLimit>("eval-node-limit")
.unwrap()
.to_owned();
deployment.set_parallelism_limit(parallelism_limit); deployment.set_parallelism_limit(parallelism_limit);
deployment.set_evaluation_node_limit(evaluation_node_limit); deployment.set_evaluation_node_limit(eval_node_limit);
let (deployment, output) = tokio::join!(deployment.execute(), output.run_until_completion(),); let (deployment, output) = tokio::join!(deployment.execute(), output.run_until_completion(),);

View file

@ -1,9 +1,9 @@
use clap::{builder::PossibleValuesParser, Arg, Command as ClapCommand}; use clap::{builder::PossibleValuesParser, Arg, Args, Command as ClapCommand};
use crate::util; use crate::util;
use super::apply;
pub use super::apply::run; pub use super::apply::run;
use super::apply::DeployOpts;
pub fn subcommand() -> ClapCommand { pub fn subcommand() -> ClapCommand {
let command = ClapCommand::new("build") let command = ClapCommand::new("build")
@ -21,7 +21,5 @@ This subcommand behaves as if you invoked `apply` with the `build` goal."#,
.num_args(1), .num_args(1),
); );
let command = apply::register_deploy_args(command); util::register_selector_args(DeployOpts::augment_args_for_update(command))
util::register_selector_args(command)
} }

View file

@ -1,9 +1,9 @@
use clap::{builder::PossibleValuesParser, Arg, Command as ClapCommand}; use clap::{builder::PossibleValuesParser, Arg, Args, Command as ClapCommand};
use crate::util; use crate::util;
use super::apply;
pub use super::apply::run; pub use super::apply::run;
use super::apply::DeployOpts;
pub fn subcommand() -> ClapCommand { pub fn subcommand() -> ClapCommand {
let command = ClapCommand::new("upload-keys") let command = ClapCommand::new("upload-keys")
@ -21,7 +21,5 @@ This subcommand behaves as if you invoked `apply` with the pseudo `keys` goal."#
.num_args(1), .num_args(1),
); );
let command = apply::register_deploy_args(command); util::register_selector_args(DeployOpts::augment_args_for_update(command))
util::register_selector_args(command)
} }

View file

@ -1,5 +1,7 @@
//! Parallelism limits. //! Parallelism limits.
use std::str::FromStr;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
/// Amount of RAM reserved for the system, in MB. /// Amount of RAM reserved for the system, in MB.
@ -55,9 +57,10 @@ impl ParallelismLimit {
/// - A simple heuristic based on remaining memory in the system /// - A simple heuristic based on remaining memory in the system
/// - A supplied number /// - A supplied number
/// - No limit at all /// - No limit at all
#[derive(Copy, Clone, Debug)] #[derive(Copy, Clone, Debug, Default)]
pub enum EvaluationNodeLimit { pub enum EvaluationNodeLimit {
/// Use a naive heuristic based on available memory. /// Use a naive heuristic based on available memory.
#[default]
Heuristic, Heuristic,
/// Supply the maximum number of nodes. /// Supply the maximum number of nodes.
@ -67,9 +70,28 @@ pub enum EvaluationNodeLimit {
None, None,
} }
impl Default for EvaluationNodeLimit { impl FromStr for EvaluationNodeLimit {
fn default() -> Self { type Err = &'static str;
Self::Heuristic fn from_str(s: &str) -> Result<Self, Self::Err> {
if s == "auto" {
return Ok(EvaluationNodeLimit::Heuristic);
}
match s.parse::<usize>() {
Ok(0) => Ok(EvaluationNodeLimit::None),
Ok(n) => Ok(EvaluationNodeLimit::Manual(n)),
Err(_) => Err("The value must be a valid number or `auto`"),
}
}
}
impl std::fmt::Display for EvaluationNodeLimit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Heuristic => write!(f, "auto"),
Self::None => write!(f, "0"),
Self::Manual(n) => write!(f, "{n}"),
}
} }
} }

View file

@ -1,6 +1,6 @@
//! Deployment options. //! Deployment options.
use clap::{builder::PossibleValue, ValueEnum}; use clap::ValueEnum;
use crate::nix::CopyOptions; use crate::nix::CopyOptions;
@ -36,12 +36,22 @@ pub struct Options {
} }
/// Which evaluator to use. /// Which evaluator to use.
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq, Default, ValueEnum)]
pub enum EvaluatorType { pub enum EvaluatorType {
#[default]
Chunked, Chunked,
Streaming, Streaming,
} }
impl std::fmt::Display for EvaluatorType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
Self::Chunked => "chunked",
Self::Streaming => "streaming",
})
}
}
impl Options { impl Options {
pub fn set_substituters_push(&mut self, value: bool) { pub fn set_substituters_push(&mut self, value: bool) {
self.substituters_push = value; self.substituters_push = value;
@ -98,16 +108,3 @@ impl Default for Options {
} }
} }
} }
impl ValueEnum for EvaluatorType {
fn value_variants<'a>() -> &'a [Self] {
&[Self::Chunked, Self::Streaming]
}
fn to_possible_value<'a>(&self) -> Option<PossibleValue> {
match self {
Self::Chunked => Some(PossibleValue::new("chunked")),
Self::Streaming => Some(PossibleValue::new("streaming")),
}
}
}

View file

@ -5,12 +5,31 @@ use std::convert::AsRef;
use std::iter::{FromIterator, Iterator}; use std::iter::{FromIterator, Iterator};
use std::str::FromStr; use std::str::FromStr;
use clap::Args;
use glob::Pattern as GlobPattern; use glob::Pattern as GlobPattern;
use super::{ColmenaError, ColmenaResult, NodeConfig, NodeName}; use super::{ColmenaError, ColmenaResult, NodeConfig, NodeName};
#[derive(Debug, Args)]
pub struct NodeFilterOpts {
#[arg(
long,
value_name = "NODES",
help = "Node selector",
long_help = r#"Select a list of nodes to deploy to.
The list is comma-separated and globs are supported. To match tags, prepend the filter by @. Valid examples:
- host1,host2,host3
- edge-*
- edge-*,core-*
- @a-tag,@tags-can-have-*"#
)]
pub on: Option<NodeFilter>,
}
/// A node filter containing a list of rules. /// A node filter containing a list of rules.
#[derive(Clone)] #[derive(Clone, Debug)]
pub struct NodeFilter { pub struct NodeFilter {
rules: Vec<Rule>, rules: Vec<Rule>,
} }