From f779c26f9056fb58773a37e2d3c634f38c08f104 Mon Sep 17 00:00:00 2001 From: Yaroslav Bolyukin Date: Wed, 24 Jul 2024 11:01:44 +0000 Subject: [PATCH] feat: ability to select specialisation to activate --- --- a/Cargo.lock +++ b/Cargo.lock @@ -784,7 +784,7 @@ "itertools", "nix-eval", "nixlike", - "once_cell", + "nom", "openssh", "owo-colors", "peg", --- a/cmds/fleet/Cargo.toml +++ b/cmds/fleet/Cargo.toml @@ -19,7 +19,6 @@ serde_json.workspace = true tempfile.workspace = true time = { version = "0.3", features = ["serde"] } -once_cell = "1.19" hostname = "0.4.0" age-core = "0.10" peg = "0.8" @@ -45,6 +44,7 @@ human-repr = { version = "1.1", optional = true } indicatif = { version = "0.17", optional = true } nix-eval.workspace = true +nom = "7.1.3" [features] # Not quite stable --- a/cmds/fleet/src/cmds/build_systems.rs +++ b/cmds/fleet/src/cmds/build_systems.rs @@ -126,6 +126,7 @@ action: DeployAction, host: &ConfigHost, built: PathBuf, + specialisation: Option, disable_rollback: bool, ) -> Result<()> { let mut failed = false; @@ -190,9 +191,14 @@ if action.should_activate() && !failed { let _span = info_span!("activating").entered(); info!("executing activation script"); - let mut switch_script = built.clone(); - switch_script.push("bin"); - switch_script.push("switch-to-configuration"); + let specialised = if let Some(specialisation) = specialisation { + let mut specialised = built.join("specialisation"); + specialised.push(specialisation); + specialised + } else { + built.clone() + }; + let switch_script = specialised.join("bin/switch-to-configuration"); let mut cmd = host.cmd(switch_script).in_current_span().await?; cmd.arg(action.name().expect("upload.should_activate == false")); if let Err(e) = cmd.sudo().run().in_current_span().await { @@ -255,12 +261,11 @@ .system .build[{ build_attr }] ); - let outputs = drv.build().await.map_err(|e| { + let outputs = drv.build().await.inspect_err(|_| { if build_attr == "sdImage" { info!("sd-image build failed"); info!("Make sure you have imported modulesPath/installer/sd-card/sd-image-[-installer].nix (For installer, you may want to check config)"); } - e })?; let out_output = outputs .get("out") @@ -275,7 +280,7 @@ let set = LocalSet::new(); let build_attr = self.build_attr.clone(); for host in hosts.into_iter() { - if config.should_skip(&host.name) { + if config.should_skip(&host).await? { continue; } let config = config.clone(); @@ -324,7 +329,7 @@ let hosts = config.list_hosts().await?; let set = LocalSet::new(); for host in hosts.into_iter() { - if config.should_skip(&host.name) { + if config.should_skip(&host).await? { continue; } let config = config.clone(); @@ -379,8 +384,19 @@ } } } - if let Err(e) = - deploy_task(self.action, &host, built, self.disable_rollback).await + if let Err(e) = deploy_task( + self.action, + &host, + built, + if let Ok(v) = config.action_attr(&host, "specialisation").await { + v + } else { + error!("unreachable? failed to get specialization"); + return; + }, + self.disable_rollback, + ) + .await { error!("activation failed: {e}"); } --- a/cmds/fleet/src/cmds/secrets/mod.rs +++ b/cmds/fleet/src/cmds/secrets/mod.rs @@ -436,7 +436,7 @@ match self { Secret::ForceKeys => { for host in config.list_hosts().await? { - if config.should_skip(&host.name) { + if config.should_skip(&host).await? { continue; } config.key(&host.name).await?; @@ -639,7 +639,7 @@ } } for host in config.list_hosts().await? { - if config.should_skip(&host.name) { + if config.should_skip(&host).await? { continue; } --- a/cmds/fleet/src/host.rs +++ b/cmds/fleet/src/host.rs @@ -1,4 +1,6 @@ use std::{ + cell::OnceCell, + collections::BTreeMap, env::current_dir, ffi::{OsStr, OsString}, fmt::Display, @@ -10,9 +12,16 @@ }; use anyhow::{anyhow, bail, ensure, Context, Result}; -use clap::{ArgGroup, Parser}; +use clap::Parser; use fleet_shared::SecretData; use nix_eval::{nix_go, nix_go_json, NixSessionPool, Value}; +use nom::{ + bytes::complete::take_while1, + character::complete::char, + combinator::{map, opt}, + multi::separated_list1, + sequence::{preceded, separated_pair}, +}; use openssh::SessionBuilder; use serde::de::DeserializeOwned; use tempfile::NamedTempFile; @@ -53,10 +62,26 @@ pub name: String, pub local: bool, pub session: OnceLock>, + groups: OnceCell>, pub nixos_config: Option, } impl ConfigHost { + pub async fn tags(&self) -> Result> { + if let Some(v) = self.groups.get() { + return Ok(v.clone()); + } + // TOCTOU is possible here in case if config is changed, but this case is not handled anywhere anyway, + // assuming getting tags always returns the same value. + let Some(nixos_config) = &self.nixos_config else { + return Ok(vec![]); + }; + let tags: Vec = nix_go_json!(nixos_config.tags); + + let _ = self.groups.set(tags.clone()); + + Ok(tags) + } async fn open_session(&self) -> Result> { assert!(!self.local, "do not open ssh connection to local session"); // FIXME: TOCTOU @@ -217,15 +242,71 @@ } impl Config { - pub fn should_skip(&self, host: &str) -> bool { - if !self.opts.skip.is_empty() { - self.opts.skip.iter().any(|h| h as &str == host) - } else if !self.opts.only.is_empty() { - !self.opts.only.iter().any(|h| h as &str == host) - } else { - false + pub async fn should_skip(&self, host: &ConfigHost) -> Result { + if !self.opts.skip.is_empty() && self.opts.skip.iter().any(|h| h as &str == host.name) { + return Ok(true); + } + if self.opts.only.is_empty() { + return Ok(false); + } + let mut have_group_matches = false; + for item in self.opts.only.iter() { + match item { + HostItem::Host { name, .. } if *name == host.name => { + return Ok(false); + } + HostItem::Tag { .. } => { + have_group_matches = true; + } + _ => {} + } } + if have_group_matches { + let host_tags = host.tags().await?; + for item in self.opts.only.iter() { + match item { + HostItem::Tag { name, .. } if host_tags.contains(name) => { + return Ok(false); + } + _ => {} + } + } + } + Ok(true) } + pub async fn action_attr(&self, host: &ConfigHost, attr: &str) -> Result> { + if self.opts.only.is_empty() { + return Ok(None); + } + let mut have_group_matches = false; + for item in self.opts.only.iter() { + match item { + HostItem::Host { name, attrs } + if *name == host.name && attrs.contains_key(attr) => + { + return Ok(attrs.get(attr).cloned()); + } + HostItem::Tag { attrs, .. } if attrs.contains_key(attr) => { + have_group_matches = true; + } + _ => {} + } + } + if have_group_matches { + let host_tags = host.tags().await?; + for item in self.opts.only.iter() { + match item { + HostItem::Tag { name, attrs } + if host_tags.contains(name) && attrs.contains_key(attr) => + { + return Ok(attrs.get(attr).cloned()); + } + _ => {} + } + } + } + Ok(None) + } pub fn is_local(&self, host: &str) -> bool { self.opts.localhost.as_ref().map(|s| s as &str) == Some(host) } @@ -237,6 +318,11 @@ local: true, session: OnceLock::new(), nixos_config: None, + groups: { + let cell = OnceCell::new(); + let _ = cell.set(vec![]); + cell + }, } } @@ -249,6 +335,7 @@ local: self.is_local(name), session: OnceLock::new(), nixos_config: Some(nixos_config), + groups: OnceCell::new(), }) } pub async fn list_hosts(&self) -> Result> { @@ -356,15 +443,59 @@ } } +#[derive(Clone)] +enum HostItem { + Host { + name: String, + attrs: BTreeMap, + }, + Tag { + name: String, + attrs: BTreeMap, + }, +} +fn host_item_parser(input: &str) -> Result { + fn err_to_string(err: nom::Err>) -> String { + err.to_string() + } + + let (input, is_tag) = map(opt(char('@')), |c| c.is_some())(input).map_err(err_to_string)?; + let (input, name) = map( + take_while1(|v| v != ',' && v != '?' && v != '@'), + str::to_owned, + )(input) + .map_err(err_to_string)?; + + let kw_item = separated_pair( + map(take_while1(|v| v != '&' && v != '='), str::to_owned), + char('='), + map(take_while1(|v| v != '&'), str::to_owned), + ); + let kw = map(separated_list1(char('&'), kw_item), |vec| { + vec.into_iter().collect::>() + }); + let mut opt_kw = map(opt(preceded(char('?'), kw)), Option::unwrap_or_default); + + let (input, attrs) = opt_kw(input).map_err(err_to_string)?; + + if !input.is_empty() { + return Err(format!("unexpected trailing input: {input:?}")); + } + Ok(if is_tag { + HostItem::Tag { name, attrs } + } else { + HostItem::Host { name, attrs } + }) +} + #[derive(Parser, Clone)] -#[clap(group = ArgGroup::new("target_hosts"))] pub struct FleetOpts { /// All hosts except those would be skipped - #[clap(long, number_of_values = 1, group = "target_hosts")] - only: Vec, + #[clap(long, number_of_values = 1, value_parser = host_item_parser)] + only: Vec, /// Hosts to skip - #[clap(long, number_of_values = 1, group = "target_hosts")] + #[clap(long, number_of_values = 1)] skip: Vec, /// Host, which should be threaten as current machine --- a/flake.lock +++ b/flake.lock @@ -7,11 +7,11 @@ ] }, "locked": { - "lastModified": 1720226507, - "narHash": "sha256-yHVvNsgrpyNTXZBEokL8uyB2J6gB1wEx0KOJzoeZi1A=", + "lastModified": 1721699339, + "narHash": "sha256-UqtSwU13vpzzM6w8tGghEbA7ObM3NCDzSpz19QQo9XE=", "owner": "ipetkov", "repo": "crane", - "rev": "0aed560c5c0a61c9385bddff471a13036203e11c", + "rev": "0081e9c447f3b70822c142908f08ceeb436982b8", "type": "github" }, "original": { @@ -40,11 +40,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1720525988, - "narHash": "sha256-6Vvrwl2rKrRt5gAYTFlM/pihCwHw8SY2o81TBm7KhIQ=", + "lastModified": 1721814637, + "narHash": "sha256-L3QkCvxeByJfW45wLkdZ9pL5h9PezOwwfx7G2sRfjiU=", "owner": "nixos", "repo": "nixpkgs", - "rev": "a630e7a8476e51b116f1ca7444dbad20701823d7", + "rev": "e0c444a0b8413a31df199052f5714d409dc4c1d0", "type": "github" }, "original": { @@ -68,11 +68,11 @@ }, "nixpkgs-stable-for-tests": { "locked": { - "lastModified": 1720386169, - "narHash": "sha256-NGKVY4PjzwAa4upkGtAMz1npHGoRzWotlSnVlqI40mo=", + "lastModified": 1721548954, + "narHash": "sha256-7cCC8+Tdq1+3OPyc3+gVo9dzUNkNIQfwSDJ2HSi2u3o=", "owner": "nixos", "repo": "nixpkgs", - "rev": "194846768975b7ad2c4988bdb82572c00222c0d7", + "rev": "63d37ccd2d178d54e7fb691d7ec76000740ea24a", "type": "github" }, "original": { @@ -98,11 +98,11 @@ ] }, "locked": { - "lastModified": 1720491570, - "narHash": "sha256-PHS2BcQ9kxBpu9GKlDg3uAlrX/ahQOoAiVmwGl6BjD4=", + "lastModified": 1721810656, + "narHash": "sha256-33UCMmgPL+sz06+iupNkl99hcBABP56ENcxSoKqr0TY=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "b970af40fdc4bd80fd764796c5f97c15e2b564eb", + "rev": "a6afdaab4a47d6ecf647a74968e92a51c4a18e5a", "type": "github" }, "original": { -- gitstuff