From 6abbf025ed70f924073a487e302b056bc01c3d80 Mon Sep 17 00:00:00 2001 From: marshmallow Date: Thu, 12 Mar 2026 20:58:56 +1100 Subject: [PATCH] re: refactor how nodes are executed (#400) Signed-off-by: marshmallow Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- CHANGELOG.md | 5 + crates/cli/src/apply.rs | 35 +- crates/cli/src/cli.rs | 4 +- crates/cli/src/main.rs | 11 +- crates/core/src/commands/common.rs | 24 +- crates/core/src/commands/mod.rs | 35 +- crates/core/src/commands/noninteractive.rs | 15 +- crates/core/src/commands/pty/mod.rs | 52 +- crates/core/src/hive/executor.rs | 167 +++++ crates/core/src/hive/mod.rs | 33 +- crates/core/src/hive/node.rs | 700 ++++---------------- crates/core/src/hive/plan.rs | 735 +++++++++++++++++++++ crates/core/src/hive/steps/activate.rs | 172 +++-- crates/core/src/hive/steps/build.rs | 31 +- crates/core/src/hive/steps/cleanup.rs | 28 - crates/core/src/hive/steps/evaluate.rs | 11 +- crates/core/src/hive/steps/keys.rs | 175 +---- crates/core/src/hive/steps/mod.rs | 1 - crates/core/src/hive/steps/ping.rs | 34 +- crates/core/src/hive/steps/push.rs | 72 +- 20 files changed, 1311 insertions(+), 1029 deletions(-) create mode 100644 crates/core/src/hive/executor.rs create mode 100644 crates/core/src/hive/plan.rs delete mode 100644 crates/core/src/hive/steps/cleanup.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index cf7516a..ccd7608 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - The domain for documentation to be `wire.forall.systems`. The previous URL will continue to be available but may redirect in the future. +- Refactored node execution to be in two distinct phases, "planning" and + "execution". Previously, picking what steps would be run was done on the fly + during execution. +- Cases where there are no keys to deploy, such as having 0 keys or filtered + keys, the "Key" step will not be planned when it previously would have. ### Fixed diff --git a/crates/cli/src/apply.rs b/crates/cli/src/apply.rs index b83abbf..9a7aa16 100644 --- a/crates/cli/src/apply.rs +++ b/crates/cli/src/apply.rs @@ -11,7 +11,9 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use thiserror::Error; use tracing::{error, info}; -use wire_core::hive::node::{Context, GoalExecutor, Name, Node, Objective, StepState}; +use wire_core::hive::executor::execute; +use wire_core::hive::node::{Name, Node}; +use wire_core::hive::plan::{Goal, plan_for_node}; use wire_core::hive::{Hive, HiveLocation}; use wire_core::status::STATUS; use wire_core::{SubCommandModifiers, errors::HiveLibError}; @@ -96,15 +98,15 @@ where pub async fn apply( hive: &mut Hive, - should_shutdown: Arc, + should_quit: Arc, location: HiveLocation, args: CommonVerbArgs, partition: Partitions, - make_objective: F, + make_goal: F, mut modifiers: SubCommandModifiers, ) -> Result<()> where - F: Fn(&Name, &Node) -> Objective, + F: Fn(&Name, &Node) -> Goal, { let location = Arc::new(location); @@ -142,23 +144,17 @@ where .iter_mut() .filter(|(name, _)| partitioned_names.contains(name)) .map(|(name, node)| { - info!("Resolved {:?} to include {}", args.on, name); + let goal = make_goal(name, node); - let objective = make_objective(name, node); - - let context = Context { + let plan = plan_for_node( node, - name, - objective, - state: StepState::default(), - hive_location: location.clone(), - modifiers, - should_quit: should_shutdown.clone(), - }; - - GoalExecutor::new(context) - .execute() - .map(move |result| (name, result)) + name.clone(), + &goal, + location.clone(), + &modifiers, + should_quit.clone(), + ); + execute(plan).map(move |result| (name, result)) }) .peekable(); @@ -168,6 +164,7 @@ where let futures = futures::stream::iter(set).buffer_unordered(args.parallel); let result = futures.collect::>().await; + let (successful, errors): (Vec<_>, Vec<_>) = result .into_iter() diff --git a/crates/cli/src/cli.rs b/crates/cli/src/cli.rs index 9e49543..5e558ee 100644 --- a/crates/cli/src/cli.rs +++ b/crates/cli/src/cli.rs @@ -11,7 +11,9 @@ use clap_verbosity_flag::InfoLevel; use tokio::runtime::Handle; use wire_core::SubCommandModifiers; use wire_core::commands::common::get_hive_node_names; -use wire_core::hive::node::{Goal as HiveGoal, HandleUnreachable, Name, SwitchToConfigurationGoal}; +use wire_core::hive::node::{ + ApplyGoal as HiveGoal, HandleUnreachable, Name, SwitchToConfigurationGoal, +}; use wire_core::hive::{Hive, get_hive_location}; use std::io::IsTerminal; diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 7e618d8..8ddc06a 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -28,9 +28,9 @@ use wire_core::cache::InspectionCache; use wire_core::commands::common::get_hive_node_names; use wire_core::hive::Hive; use wire_core::hive::get_hive_location; -use wire_core::hive::node::ApplyObjective; -use wire_core::hive::node::Objective; use wire_core::hive::node::should_apply_locally; +use wire_core::hive::plan::ApplyGoalArgs; +use wire_core::hive::plan::Goal; #[macro_use] extern crate enum_display_derive; @@ -86,7 +86,7 @@ async fn main() -> Result<()> { match args.command { cli::Commands::Apply(apply_args) => { let mut hive = Hive::new_from_path(&location, cache.clone(), modifiers).await?; - let goal: wire_core::hive::node::Goal = apply_args.goal.clone().try_into().unwrap(); + let goal = apply_args.goal.clone().try_into().unwrap(); // Respect user's --always-build-local arg hive.force_always_local(apply_args.always_build_local)?; @@ -98,7 +98,7 @@ async fn main() -> Result<()> { apply_args.common, Partitions::default(), |name, node| { - Objective::Apply(ApplyObjective { + Goal::Apply(ApplyGoalArgs { goal, no_keys: apply_args.no_keys, reboot: apply_args.reboot, @@ -108,6 +108,7 @@ async fn main() -> Result<()> { &name.0, ), handle_unreachable: apply_args.handle_unreachable.clone().into(), + host_platform: node.host_platform.clone(), }) }, modifiers, @@ -123,7 +124,7 @@ async fn main() -> Result<()> { location, build_args.common, build_args.partition.unwrap_or_default(), - |_name, _node| Objective::BuildLocally, + |_name, _node| Goal::Build, modifiers, ) .await?; diff --git a/crates/core/src/commands/common.rs b/crates/core/src/commands/common.rs index 37eb39c..8d17b67 100644 --- a/crates/core/src/commands/common.rs +++ b/crates/core/src/commands/common.rs @@ -14,7 +14,7 @@ use crate::{ errors::{CommandError, HiveInitialisationError, HiveLibError}, hive::{ HiveLocation, - node::{Context, Objective, Push}, + node::{Context, Push, SharedTarget}, }, }; @@ -28,22 +28,24 @@ fn get_common_copy_path_help(error: &CommandError) -> Option { } } -pub async fn push(context: &Context<'_>, push: Push<'_>) -> Result<(), HiveLibError> { +pub async fn push( + context: &Context, + target: &SharedTarget, + push: Push<'_>, + substitute_on_destination: bool, +) -> Result<(), HiveLibError> { + let target = target.0.read().await; + let mut command_string = CommandStringBuilder::nix(); command_string.args(&["--extra-experimental-features", "nix-command", "copy"]); - if let Objective::Apply(apply_objective) = context.objective { - command_string.opt_arg( - apply_objective.substitute_on_destination, - "--substitute-on-destination", - ); - } + command_string.opt_arg(substitute_on_destination, "--substitute-on-destination"); command_string.arg("--to"); command_string.args(&[ format!( "ssh://{user}@{host}", - user = context.node.target.user, - host = context.node.target.get_preferred_host()?, + user = target.user, + host = target.get_preferred_host()?, ), match push { Push::Derivation(drv) => format!("{drv} --derivation"), @@ -56,7 +58,7 @@ pub async fn push(context: &Context<'_>, push: Push<'_>) -> Result<(), HiveLibEr .mode(crate::commands::ChildOutputMode::Nix), HashMap::from([( "NIX_SSHOPTS".into(), - context.node.target.create_ssh_opts(context.modifiers)?, + target.create_ssh_opts(context.modifiers)?, )]), ) .await?; diff --git a/crates/core/src/commands/mod.rs b/crates/core/src/commands/mod.rs index f3ab775..4e70b15 100644 --- a/crates/core/src/commands/mod.rs +++ b/crates/core/src/commands/mod.rs @@ -1,8 +1,15 @@ // SPDX-License-Identifier: AGPL-3.0-or-later // Copyright 2024-2025 wire Contributors -use crate::commands::pty::{InteractiveChildChip, interactive_command_with_env}; -use std::{collections::HashMap, str::from_utf8, sync::LazyLock}; +use crate::{ + commands::pty::{InteractiveChildChip, interactive_command_with_env}, + hive::node::SharedTarget, +}; +use std::{ + collections::HashMap, + str::from_utf8, + sync::{Arc, LazyLock}, +}; use aho_corasick::AhoCorasick; use gjson::Value; @@ -15,7 +22,6 @@ use crate::{ SubCommandModifiers, commands::noninteractive::{NonInteractiveChildChip, non_interactive_command_with_env}, errors::{CommandError, HiveLibError}, - hive::node::{Node, Target}, }; pub(crate) mod builder; @@ -37,9 +43,9 @@ pub enum Either { } #[derive(Debug)] -pub(crate) struct CommandArguments<'t, S: AsRef> { +pub(crate) struct CommandArguments> { modifiers: SubCommandModifiers, - target: Option<&'t Target>, + target: Option, output_mode: ChildOutputMode, command_string: S, keep_stdin_open: bool, @@ -55,7 +61,7 @@ static AHO_CORASICK: LazyLock = LazyLock::new(|| { .unwrap() }); -impl<'a, S: AsRef> CommandArguments<'a, S> { +impl> CommandArguments { pub(crate) const fn new(command_string: S, modifiers: SubCommandModifiers) -> Self { Self { command_string, @@ -68,7 +74,7 @@ impl<'a, S: AsRef> CommandArguments<'a, S> { } } - pub(crate) const fn execute_on_remote(mut self, target: Option<&'a Target>) -> Self { + pub(crate) fn execute_on_remote(mut self, target: Option) -> Self { self.target = target; self } @@ -83,9 +89,8 @@ impl<'a, S: AsRef> CommandArguments<'a, S> { self } - pub(crate) fn elevated(mut self, node: &Node) -> Self { - self.privilege_escalation_command = - Some(node.privilege_escalation_command.iter().join(" ")); + pub(crate) fn privileged(mut self, escalation_command: &[Arc]) -> Self { + self.privilege_escalation_command = Some(escalation_command.iter().join(" ")); self } @@ -100,13 +105,13 @@ impl<'a, S: AsRef> CommandArguments<'a, S> { } pub(crate) async fn run_command>( - arguments: &CommandArguments<'_, S>, + arguments: &CommandArguments, ) -> Result, HiveLibError> { run_command_with_env(arguments, HashMap::new()).await } pub(crate) async fn run_command_with_env>( - arguments: &CommandArguments<'_, S>, + arguments: &CommandArguments, envs: HashMap, ) -> Result, HiveLibError> { // use the non interactive command runner when forced @@ -114,9 +119,9 @@ pub(crate) async fn run_command_with_env>( if arguments.modifiers.non_interactive || (arguments.target.is_none() && !arguments.is_elevated()) { - return Ok(Either::Right(non_interactive_command_with_env( - arguments, envs, - )?)); + return Ok(Either::Right( + non_interactive_command_with_env(arguments, envs).await?, + )); } Ok(Either::Left( diff --git a/crates/core/src/commands/noninteractive.rs b/crates/core/src/commands/noninteractive.rs index b1d6a7e..267f2a2 100644 --- a/crates/core/src/commands/noninteractive.rs +++ b/crates/core/src/commands/noninteractive.rs @@ -11,7 +11,7 @@ use crate::{ SubCommandModifiers, commands::{ChildOutputMode, CommandArguments, WireCommandChip}, errors::{CommandError, HiveLibError}, - hive::node::Target, + hive::node::SharedTarget, }; use itertools::Itertools; use tokio::{ @@ -32,12 +32,12 @@ pub(crate) struct NonInteractiveChildChip { } #[instrument(skip_all, name = "run", fields(elevated = %arguments.is_elevated()))] -pub(crate) fn non_interactive_command_with_env>( +pub(crate) async fn non_interactive_command_with_env>( arguments: &CommandArguments, envs: HashMap, ) -> Result { - let mut command = if let Some(target) = arguments.target { - create_sync_ssh_command(target, arguments.modifiers)? + let mut command = if let Some(ref target) = arguments.target { + create_sync_ssh_command(target, arguments.modifiers).await? } else { let mut command = Command::new("sh"); @@ -187,12 +187,13 @@ pub async fn handle_io( debug!("io_handler: goodbye!"); } -fn create_sync_ssh_command( - target: &Target, +async fn create_sync_ssh_command( + target: &SharedTarget, modifiers: SubCommandModifiers, ) -> Result { + let target = target.0.read().await; let mut command = Command::new("ssh"); - command.args(target.create_ssh_args(modifiers, true)?); + command.args(target.create_ssh_args(modifiers)?); command.arg(target.get_preferred_host()?.to_string()); Ok(command) diff --git a/crates/core/src/commands/pty/mod.rs b/crates/core/src/commands/pty/mod.rs index e2e1879..bca4ca1 100644 --- a/crates/core/src/commands/pty/mod.rs +++ b/crates/core/src/commands/pty/mod.rs @@ -2,6 +2,7 @@ // Copyright 2024-2025 wire Contributors use crate::commands::pty::output::{WatchStdoutArguments, handle_pty_stdout}; +use crate::hive::node::SharedTarget; use crate::status::STATUS; use aho_corasick::PatternID; use itertools::Itertools; @@ -29,7 +30,6 @@ use crate::{SubCommandModifiers, acquire_stdin_lock}; use crate::{ commands::{ChildOutputMode, WireCommandChip}, errors::HiveLibError, - hive::node::Target, }; mod input; @@ -85,7 +85,7 @@ static FAILED_PATTERN: LazyLock = LazyLock::new(|| PatternID::must(2) const IO_SUBS: &str = "1> >(while IFS= read -r line; do echo \"#$line\"; done)"; fn create_ending_segment>( - arguments: &CommandArguments<'_, S>, + arguments: &CommandArguments, needles: &Needles, ) -> String { let Needles { @@ -110,7 +110,7 @@ fn create_ending_segment>( } fn create_starting_segment>( - arguments: &CommandArguments<'_, S>, + arguments: &CommandArguments, start_needle: &Arc>, ) -> String { if matches!(arguments.output_mode, ChildOutputMode::Interactive) { @@ -125,10 +125,10 @@ fn create_starting_segment>( #[instrument(skip_all, name = "run-int", fields(elevated = %arguments.is_elevated(), mode = ?arguments.output_mode))] pub(crate) async fn interactive_command_with_env>( - arguments: &CommandArguments<'_, S>, + arguments: &CommandArguments, envs: std::collections::HashMap, ) -> Result { - print_authenticate_warning(arguments)?; + print_authenticate_warning(arguments).await?; let needles = create_needles(); let pty_system = NativePtySystem::default(); @@ -148,7 +148,7 @@ pub(crate) async fn interactive_command_with_env>( debug!("{command_string}"); - let mut command = build_command(arguments, command_string)?; + let mut command = build_command(arguments, command_string).await?; // give command all env vars for (key, value) in envs { @@ -242,24 +242,29 @@ pub(crate) async fn interactive_command_with_env>( }) } -fn print_authenticate_warning>( +async fn print_authenticate_warning>( arguments: &CommandArguments, ) -> Result<(), HiveLibError> { if !arguments.is_elevated() { return Ok(()); } + let target_display = if let Some(ref target) = arguments.target { + let target = target.0.read().await; + + format!( + "{}@{}:{}", + target.user, + target.get_preferred_host()?, + target.port + ) + } else { + "localhost (!)".to_string() + }; + let _ = STATUS.lock().write_above_status( &format!( - "{} | Authenticate for \"sudo {}\":\n", - arguments - .target - .map_or(Ok("localhost (!)".to_string()), |target| Ok(format!( - "{}@{}:{}", - target.user, - target.get_preferred_host()?, - target.port - )))?, + "{target_display} | Authenticate for \"sudo {}\":\n", arguments.command_string.as_ref() ) .into_bytes(), @@ -306,12 +311,12 @@ fn setup_master(pty_pair: &PtyPair) -> Result<(), HiveLibError> { Ok(()) } -fn build_command>( - arguments: &CommandArguments<'_, S>, +async fn build_command>( + arguments: &CommandArguments, command_string: &String, ) -> Result { - let mut command = if let Some(target) = arguments.target { - let mut command = create_int_ssh_command(target, arguments.modifiers)?; + let mut command = if let Some(ref target) = arguments.target { + let mut command = create_int_ssh_command(target, arguments.modifiers).await?; // force ssh to use our pseudo terminal command.arg("-tt"); @@ -428,12 +433,13 @@ impl Drop for StdinTermiosAttrGuard { } } -fn create_int_ssh_command( - target: &Target, +async fn create_int_ssh_command( + target: &SharedTarget, modifiers: SubCommandModifiers, ) -> Result { + let target = target.0.read().await; let mut command = portable_pty::CommandBuilder::new("ssh"); - command.args(target.create_ssh_args(modifiers, false)?); + command.args(target.create_ssh_args(modifiers)?); command.arg(target.get_preferred_host()?.to_string()); Ok(command) } diff --git a/crates/core/src/hive/executor.rs b/crates/core/src/hive/executor.rs new file mode 100644 index 0000000..5d22298 --- /dev/null +++ b/crates/core/src/hive/executor.rs @@ -0,0 +1,167 @@ +use crate::hive::node::Step; +use std::{assert_matches::debug_assert_matches, sync::Arc}; + +use tracing::{Instrument, Span, debug, error, event, instrument}; + +use crate::{ + EvalGoal, SubCommandModifiers, + commands::common::evaluate_hive_attribute, + errors::HiveLibError, + hive::{ + HiveLocation, + node::{Context, Derivation, ExecuteStep, Name}, + plan::NodePlan, + }, + status::STATUS, +}; + +/// returns Err if the application should shut down. +fn app_shutdown_guard(context: &Context) -> Result<(), HiveLibError> { + if context + .should_quit + .load(std::sync::atomic::Ordering::Relaxed) + { + return Err(HiveLibError::Sigint); + } + + Ok(()) +} + +/// Task that evaluates the node. +#[instrument(skip_all, name = "eval")] +async fn evaluate_task( + tx: tokio::sync::oneshot::Sender>, + hive_location: Arc, + name: Name, + modifiers: SubCommandModifiers, +) { + let output = evaluate_hive_attribute(&hive_location, &EvalGoal::GetTopLevel(&name), modifiers) + .await + .and_then(|output| { + serde_json::from_str(&output).map_err(|e| { + HiveLibError::HiveInitialisationError( + crate::errors::HiveInitialisationError::ParseEvaluateError(e), + ) + }) + }); + + debug!(output = ?output, done = true); + + let _ = tx.send(output); +} + +/// Iterates and executes the steps in the plan. +/// Performs some optimisations such as greedily executing evaluation before +/// other steps independent of evaluation's result. +#[instrument(skip_all, fields(node = %plan.context.name))] +pub async fn execute(mut plan: NodePlan) -> Result<(), HiveLibError> { + app_shutdown_guard(&plan.context)?; + + let (tx, rx) = tokio::sync::oneshot::channel(); + plan.context.state.evaluation_rx = Some(rx); + + // The name of this span should never be changed without updating + // `wire/cli/tracing_setup.rs` + debug_assert_matches!(Span::current().metadata().unwrap().name(), "execute"); + // This span should always have a `node` field by the same file + debug_assert!( + Span::current() + .metadata() + .unwrap() + .fields() + .field("node") + .is_some() + ); + + if plan.greedy_evaluate { + tokio::spawn( + evaluate_task( + tx, + plan.context.hive_location.clone(), + plan.context.name.clone(), + plan.context.modifiers, + ) + .in_current_span(), + ); + } + + let length = plan.steps.len(); + + for (position, step) in plan.steps.iter().enumerate() { + app_shutdown_guard(&plan.context)?; + + event!( + tracing::Level::INFO, + step = step.to_string(), + progress = format!("{}/{length}", position + 1) + ); + + STATUS + .lock() + .set_node_step(&plan.context.name, step.to_string()); + + if let Err(err) = step.execute(&mut plan.context).await.inspect_err(|_| { + error!("Failed to execute `{step}`"); + }) { + if matches!(step, Step::Ping(..)) && plan.ignore_failed_ping { + return Ok(()); + } + + STATUS.lock().mark_node_failed(&plan.context.name); + + return Err(err); + } + } + + STATUS.lock().mark_node_succeeded(&plan.context.name); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::{ + SubCommandModifiers, + errors::HiveLibError, + function_name, get_test_path, + hive::{ + executor::execute, + node::{ApplyGoal, HandleUnreachable, Name, Node, SwitchToConfigurationGoal}, + plan::{ApplyGoalArgs, Goal, plan_for_node}, + }, + location, + }; + use std::{assert_matches::assert_matches, path::PathBuf}; + use std::{ + env, + sync::{Arc, atomic::AtomicBool}, + }; + + #[tokio::test] + async fn plan_executor_quits_sigint() { + let location = location!(get_test_path!()); + let node = Node::default(); + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(true)); + let plan = plan_for_node( + &node.clone(), + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), + should_apply_locally: true, + no_keys: true, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + let status = execute(plan).await; + + assert_matches!(status, Err(HiveLibError::Sigint)); + } +} diff --git a/crates/core/src/hive/mod.rs b/crates/core/src/hive/mod.rs index 85f9221..a4b0b82 100644 --- a/crates/core/src/hive/mod.rs +++ b/crates/core/src/hive/mod.rs @@ -22,7 +22,9 @@ use crate::commands::common::evaluate_hive_attribute; use crate::commands::{CommandArguments, Either, WireCommandChip, run_command}; use crate::errors::{HiveInitialisationError, HiveLocationError}; use crate::{EvalGoal, HiveLibError, SubCommandModifiers}; +pub mod executor; pub mod node; +pub mod plan; pub mod steps; #[derive(Serialize, Deserialize, Debug, PartialEq)] @@ -180,14 +182,14 @@ impl Display for Hive { } } -#[derive(Debug, PartialEq, Eq, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] pub struct FlakePrefetch { pub(crate) hash: String, #[serde(rename = "storePath")] pub(crate) store_path: String, } -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum HiveLocation { HiveNix(PathBuf), Flake { @@ -288,8 +290,6 @@ pub async fn get_hive_location( #[cfg(test)] mod tests { - use im::vector; - use crate::{ errors::CommandError, get_test_path, @@ -346,17 +346,20 @@ mod tests { let node = Node { target: node::Target::from_host("name"), - keys: vector![Key { - name: "different-than-a".into(), - dest_dir: "/run/keys/".into(), - path: "/run/keys/different-than-a".into(), - group: "root".into(), - user: "root".into(), - permissions: "0600".into(), - source: Source::String("hi".into()), - upload_at: UploadKeyAt::PreActivation, - environment: im::HashMap::new() - }], + keys: vec![ + Key { + name: "different-than-a".into(), + dest_dir: "/run/keys/".into(), + path: "/run/keys/different-than-a".into(), + group: "root".into(), + user: "root".into(), + permissions: "0600".into(), + source: Source::String("hi".into()), + upload_at: UploadKeyAt::PreActivation, + environment: im::HashMap::new(), + } + .into(), + ], build_remotely: true, ..Default::default() }; diff --git a/crates/core/src/hive/node.rs b/crates/core/src/hive/node.rs index 683ef70..2dcbbf8 100644 --- a/crates/core/src/hive/node.rs +++ b/crates/core/src/hive/node.rs @@ -5,26 +5,22 @@ use enum_dispatch::enum_dispatch; use gethostname::gethostname; use serde::{Deserialize, Serialize}; -use std::assert_matches::debug_assert_matches; use std::fmt::Display; use std::sync::Arc; use std::sync::atomic::AtomicBool; -use tokio::sync::oneshot; -use tracing::{Instrument, Level, Span, debug, error, event, instrument, trace}; +use tokio::sync::{RwLock, oneshot}; +use tracing::instrument; use crate::commands::builder::CommandStringBuilder; -use crate::commands::common::evaluate_hive_attribute; use crate::commands::{CommandArguments, WireCommandChip, run_command}; use crate::errors::NetworkError; use crate::hive::HiveLocation; use crate::hive::steps::build::Build; -use crate::hive::steps::cleanup::CleanUp; use crate::hive::steps::evaluate::Evaluate; -use crate::hive::steps::keys::{Key, Keys, PushKeyAgent, UploadKeyAt}; +use crate::hive::steps::keys::{Key, Keys, PushKeyAgent}; use crate::hive::steps::ping::Ping; use crate::hive::steps::push::{PushBuildOutput, PushEvaluatedOutput}; -use crate::status::STATUS; -use crate::{EvalGoal, StrictHostKeyChecking, SubCommandModifiers}; +use crate::{StrictHostKeyChecking, SubCommandModifiers}; use super::HiveLibError; use super::steps::activate::SwitchToConfiguration; @@ -44,17 +40,37 @@ pub struct Target { current_host: usize, } +#[derive(Clone, Debug)] +pub struct SharedTarget(pub Arc>); + +// Hack specifically for testing if two steps that have the same shared target +// are equal +#[cfg(test)] +impl PartialEq for SharedTarget { + fn eq(&self, other: &Self) -> bool { + let self_guard = self + .0 + .try_read() + .expect("failed to target read in test context"); + let other_guard = other + .0 + .try_read() + .expect("failed to target read in test context"); + + *self_guard == *other_guard + } +} + impl Target { #[instrument(ret(level = tracing::Level::DEBUG), skip_all)] pub fn create_ssh_opts(&self, modifiers: SubCommandModifiers) -> Result { - self.create_ssh_args(modifiers, false).map(|x| x.join(" ")) + self.create_ssh_args(modifiers).map(|x| x.join(" ")) } #[instrument(ret(level = tracing::Level::DEBUG))] pub fn create_ssh_args( &self, modifiers: SubCommandModifiers, - non_interactive_forced: bool, ) -> Result, HiveLibError> { let mut vector = vec![ "-l".to_string(), @@ -81,6 +97,32 @@ impl Target { Ok(vector) } + + /// Tests the connection to a node + pub async fn ping(&self, modifiers: SubCommandModifiers) -> Result<(), HiveLibError> { + let host = self.get_preferred_host()?; + + let mut command_string = CommandStringBuilder::new("ssh"); + command_string.arg(format!("{}@{host}", self.user)); + command_string.arg(self.create_ssh_opts(modifiers)?); + command_string.arg("exit"); + + let output = run_command( + &CommandArguments::new(command_string, modifiers) + .log_stdout() + .mode(crate::commands::ChildOutputMode::Interactive), + ) + .await?; + + output.wait_till_success().await.map_err(|source| { + HiveLibError::NetworkError(NetworkError::HostUnreachable { + host: host.to_string(), + source, + }) + })?; + + Ok(()) + } } #[cfg(test)] @@ -95,32 +137,6 @@ impl Default for Target { } } -#[cfg(test)] -impl<'a> Context<'a> { - fn create_test_context( - hive_location: HiveLocation, - name: &'a Name, - node: &'a mut Node, - ) -> Self { - Context { - name, - node, - hive_location: Arc::new(hive_location), - modifiers: SubCommandModifiers::default(), - objective: Objective::Apply(ApplyObjective { - goal: Goal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), - no_keys: false, - reboot: false, - should_apply_locally: false, - substitute_on_destination: false, - handle_unreachable: HandleUnreachable::default(), - }), - state: StepState::default(), - should_quit: Arc::new(AtomicBool::new(false)), - } - } -} - impl Target { pub fn get_preferred_host(&self) -> Result<&Arc, HiveLibError> { self.hosts @@ -171,7 +187,7 @@ pub struct Node { pub tags: im::HashSet, #[serde(rename(deserialize = "_keys", serialize = "keys"))] - pub keys: im::Vector, + pub keys: Vec>, #[serde(rename(deserialize = "_hostPlatform", serialize = "host_platform"))] pub host_platform: Arc, @@ -180,7 +196,7 @@ pub struct Node { deserialize = "privilegeEscalationCommand", serialize = "privilege_escalation_command" ))] - pub privilege_escalation_command: im::Vector>, + pub privilege_escalation_command: Arc>>, } #[cfg(test)] @@ -188,7 +204,7 @@ impl Default for Node { fn default() -> Self { Node { target: Target::default(), - keys: im::Vector::new(), + keys: Vec::new(), tags: im::HashSet::new(), privilege_escalation_command: vec!["sudo".into(), "--".into()].into(), allow_local_deployment: true, @@ -207,32 +223,6 @@ impl Node { ..Default::default() } } - - /// Tests the connection to a node - pub async fn ping(&self, modifiers: SubCommandModifiers) -> Result<(), HiveLibError> { - let host = self.target.get_preferred_host()?; - - let mut command_string = CommandStringBuilder::new("ssh"); - command_string.arg(format!("{}@{host}", self.target.user)); - command_string.arg(self.target.create_ssh_opts(modifiers)?); - command_string.arg("exit"); - - let output = run_command( - &CommandArguments::new(command_string, modifiers) - .log_stdout() - .mode(crate::commands::ChildOutputMode::Interactive), - ) - .await?; - - output.wait_till_success().await.map_err(|source| { - HiveLibError::NetworkError(NetworkError::HostUnreachable { - host: host.to_string(), - source, - }) - })?; - - Ok(()) - } } #[must_use] @@ -255,7 +245,7 @@ impl Display for Derivation { } } -#[derive(derive_more::Display, Debug, Clone, Copy)] +#[derive(derive_more::Display, Debug, Clone, Copy, PartialEq, Eq)] pub enum SwitchToConfigurationGoal { Switch, Boot, @@ -263,37 +253,17 @@ pub enum SwitchToConfigurationGoal { DryActivate, } -#[derive(derive_more::Display, Clone, Copy)] -pub enum Goal { +#[derive(derive_more::Display, Debug, Clone, Copy)] +pub enum ApplyGoal { SwitchToConfiguration(SwitchToConfigurationGoal), Build, Push, Keys, } -// TODO: Get rid of this allow and resolve it -#[allow(clippy::struct_excessive_bools)] -#[derive(Clone, Copy)] -pub struct ApplyObjective { - pub goal: Goal, - pub no_keys: bool, - pub reboot: bool, - pub should_apply_locally: bool, - pub substitute_on_destination: bool, - pub handle_unreachable: HandleUnreachable, -} - -#[derive(Clone, Copy)] -pub enum Objective { - Apply(ApplyObjective), - BuildLocally, -} - #[enum_dispatch] pub(crate) trait ExecuteStep: Send + Sync + Display + std::fmt::Debug { - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError>; - - fn should_execute(&self, context: &Context) -> bool; + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError>; } // may include other options such as FailAll in the future @@ -313,19 +283,18 @@ pub struct StepState { pub key_agent_directory: Option, } -pub struct Context<'a> { - pub name: &'a Name, - pub node: &'a mut Node, +pub struct Context { pub hive_location: Arc, pub modifiers: SubCommandModifiers, pub state: StepState, pub should_quit: Arc, - pub objective: Objective, + pub name: Name, } #[enum_dispatch(ExecuteStep)] -#[derive(Debug, PartialEq)] -enum Step { +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub enum Step { Ping, PushKeyAgent, Keys, @@ -334,7 +303,6 @@ enum Step { Build, PushBuildOutput, SwitchToConfiguration, - CleanUp, } impl Display for Step { @@ -348,158 +316,7 @@ impl Display for Step { Self::Build(step) => step.fmt(f), Self::PushBuildOutput(step) => step.fmt(f), Self::SwitchToConfiguration(step) => step.fmt(f), - Self::CleanUp(step) => step.fmt(f), - } - } -} - -pub struct GoalExecutor<'a> { - steps: Vec, - context: Context<'a>, -} - -/// returns Err if the application should shut down. -fn app_shutdown_guard(context: &Context) -> Result<(), HiveLibError> { - if context - .should_quit - .load(std::sync::atomic::Ordering::Relaxed) - { - return Err(HiveLibError::Sigint); - } - - Ok(()) -} - -impl<'a> GoalExecutor<'a> { - #[must_use] - pub fn new(context: Context<'a>) -> Self { - Self { - steps: vec![ - Step::Ping(Ping), - Step::PushKeyAgent(PushKeyAgent), - Step::Keys(Keys { - filter: UploadKeyAt::NoFilter, - }), - Step::Keys(Keys { - filter: UploadKeyAt::PreActivation, - }), - Step::Evaluate(super::steps::evaluate::Evaluate), - Step::PushEvaluatedOutput(super::steps::push::PushEvaluatedOutput), - Step::Build(super::steps::build::Build), - Step::PushBuildOutput(super::steps::push::PushBuildOutput), - Step::SwitchToConfiguration(SwitchToConfiguration), - Step::Keys(Keys { - filter: UploadKeyAt::PostActivation, - }), - ], - context, - } - } - - #[instrument(skip_all, name = "eval")] - async fn evaluate_task( - tx: oneshot::Sender>, - hive_location: Arc, - name: Name, - modifiers: SubCommandModifiers, - ) { - let output = - evaluate_hive_attribute(&hive_location, &EvalGoal::GetTopLevel(&name), modifiers) - .await - .map(|output| { - serde_json::from_str::(&output).expect("failed to parse derivation") - }); - - debug!(output = ?output, done = true); - - let _ = tx.send(output); - } - - #[instrument(skip_all, fields(node = %self.context.name))] - pub async fn execute(mut self) -> Result<(), HiveLibError> { - app_shutdown_guard(&self.context)?; - - let (tx, rx) = oneshot::channel(); - self.context.state.evaluation_rx = Some(rx); - - // The name of this span should never be changed without updating - // `wire/cli/tracing_setup.rs` - debug_assert_matches!(Span::current().metadata().unwrap().name(), "execute"); - // This span should always have a `node` field by the same file - debug_assert!( - Span::current() - .metadata() - .unwrap() - .fields() - .field("node") - .is_some() - ); - - let spawn_evaluator = match self.context.objective { - Objective::Apply(apply_objective) => !matches!(apply_objective.goal, Goal::Keys), - Objective::BuildLocally => true, - }; - - if spawn_evaluator { - tokio::spawn( - GoalExecutor::evaluate_task( - tx, - self.context.hive_location.clone(), - self.context.name.clone(), - self.context.modifiers, - ) - .in_current_span(), - ); - } - - let steps = self - .steps - .iter() - .filter(|step| step.should_execute(&self.context)) - .inspect(|step| { - trace!("Will execute step `{step}` for {}", self.context.name); - }) - .collect::>(); - let length = steps.len(); - - for (position, step) in steps.iter().enumerate() { - app_shutdown_guard(&self.context)?; - - event!( - Level::INFO, - step = step.to_string(), - progress = format!("{}/{length}", position + 1) - ); - - STATUS - .lock() - .set_node_step(self.context.name, step.to_string()); - - if let Err(err) = step.execute(&mut self.context).await.inspect_err(|_| { - error!("Failed to execute `{step}`"); - }) { - // discard error from cleanup - let _ = CleanUp.execute(&mut self.context).await; - - if let Objective::Apply(apply_objective) = self.context.objective - && matches!(step, Step::Ping(..)) - && matches!( - apply_objective.handle_unreachable, - HandleUnreachable::Ignore, - ) - { - return Ok(()); - } - - STATUS.lock().mark_node_failed(self.context.name); - - return Err(err); - } } - - STATUS.lock().mark_node_succeeded(self.context.name); - - Ok(()) } } @@ -508,318 +325,7 @@ mod tests { use rand::distr::Alphabetic; use super::*; - use crate::{ - function_name, get_test_path, - hive::{Hive, get_hive_location}, - location, - }; - use std::{assert_matches::assert_matches, path::PathBuf}; - use std::{collections::HashMap, env}; - - fn get_steps(goal_executor: GoalExecutor) -> std::vec::Vec { - goal_executor - .steps - .into_iter() - .filter(|step| step.should_execute(&goal_executor.context)) - .collect::>() - } - - #[tokio::test] - #[cfg_attr(feature = "no_web_tests", ignore)] - async fn default_values_match() { - let mut path = get_test_path!(); - - let location = - get_hive_location(path.display().to_string(), SubCommandModifiers::default()) - .await - .unwrap(); - let hive = Hive::new_from_path(&location, None, SubCommandModifiers::default()) - .await - .unwrap(); - - let node = Node::default(); - - let mut nodes = HashMap::new(); - nodes.insert(Name("NAME".into()), node); - - path.push("hive.nix"); - - assert_eq!( - hive, - Hive { - nodes, - schema: Hive::SCHEMA_VERSION - } - ); - } - - #[tokio::test] - async fn order_build_locally() { - let location = location!(get_test_path!()); - let mut node = Node { - build_remotely: false, - ..Default::default() - }; - let name = &Name(function_name!().into()); - let executor = GoalExecutor::new(Context::create_test_context(location, name, &mut node)); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - Ping.into(), - PushKeyAgent.into(), - Keys { - filter: UploadKeyAt::PreActivation - } - .into(), - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::build::Build.into(), - crate::hive::steps::push::PushBuildOutput.into(), - SwitchToConfiguration.into(), - Keys { - filter: UploadKeyAt::PostActivation - } - .into(), - ] - ); - } - - #[tokio::test] - async fn order_keys_only() { - let location = location!(get_test_path!()); - let mut node = Node::default(); - let name = &Name(function_name!().into()); - let mut context = Context::create_test_context(location, name, &mut node); - - let Objective::Apply(ref mut apply_objective) = context.objective else { - unreachable!() - }; - - apply_objective.goal = Goal::Keys; - - let executor = GoalExecutor::new(context); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - Ping.into(), - PushKeyAgent.into(), - Keys { - filter: UploadKeyAt::NoFilter - } - .into(), - ] - ); - } - - #[tokio::test] - async fn order_build() { - let location = location!(get_test_path!()); - let mut node = Node::default(); - let name = &Name(function_name!().into()); - let mut context = Context::create_test_context(location, name, &mut node); - - let Objective::Apply(ref mut apply_objective) = context.objective else { - unreachable!() - }; - apply_objective.goal = Goal::Build; - - let executor = GoalExecutor::new(context); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - Ping.into(), - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::build::Build.into(), - crate::hive::steps::push::PushBuildOutput.into(), - ] - ); - } - - #[tokio::test] - async fn order_push_only() { - let location = location!(get_test_path!()); - let mut node = Node::default(); - let name = &Name(function_name!().into()); - let mut context = Context::create_test_context(location, name, &mut node); - - let Objective::Apply(ref mut apply_objective) = context.objective else { - unreachable!() - }; - apply_objective.goal = Goal::Push; - - let executor = GoalExecutor::new(context); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - Ping.into(), - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::push::PushEvaluatedOutput.into(), - ] - ); - } - - #[tokio::test] - async fn order_remote_build() { - let location = location!(get_test_path!()); - let mut node = Node { - build_remotely: true, - ..Default::default() - }; - - let name = &Name(function_name!().into()); - let executor = GoalExecutor::new(Context::create_test_context(location, name, &mut node)); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - Ping.into(), - PushKeyAgent.into(), - Keys { - filter: UploadKeyAt::PreActivation - } - .into(), - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::push::PushEvaluatedOutput.into(), - crate::hive::steps::build::Build.into(), - SwitchToConfiguration.into(), - Keys { - filter: UploadKeyAt::PostActivation - } - .into(), - ] - ); - } - - #[tokio::test] - async fn order_nokeys() { - let location = location!(get_test_path!()); - let mut node = Node::default(); - - let name = &Name(function_name!().into()); - let mut context = Context::create_test_context(location, name, &mut node); - - let Objective::Apply(ref mut apply_objective) = context.objective else { - unreachable!() - }; - apply_objective.no_keys = true; - - let executor = GoalExecutor::new(context); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - Ping.into(), - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::build::Build.into(), - crate::hive::steps::push::PushBuildOutput.into(), - SwitchToConfiguration.into(), - ] - ); - } - - #[tokio::test] - async fn order_should_apply_locally() { - let location = location!(get_test_path!()); - let mut node = Node::default(); - - let name = &Name(function_name!().into()); - let mut context = Context::create_test_context(location, name, &mut node); - - let Objective::Apply(ref mut apply_objective) = context.objective else { - unreachable!() - }; - apply_objective.no_keys = true; - apply_objective.should_apply_locally = true; - - let executor = GoalExecutor::new(context); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::build::Build.into(), - SwitchToConfiguration.into(), - ] - ); - } - - #[tokio::test] - async fn order_build_only() { - let location = location!(get_test_path!()); - let mut node = Node::default(); - - let name = &Name(function_name!().into()); - let mut context = Context::create_test_context(location, name, &mut node); - - context.objective = Objective::BuildLocally; - - let executor = GoalExecutor::new(context); - let steps = get_steps(executor); - - assert_eq!( - steps, - vec![ - crate::hive::steps::evaluate::Evaluate.into(), - crate::hive::steps::build::Build.into() - ] - ); - } - - #[test] - fn target_fails_increments() { - let mut target = Target::from_host("localhost"); - - assert_eq!(target.current_host, 0); - - for i in 0..100 { - target.host_failed(); - assert_eq!(target.current_host, i + 1); - } - } - - #[test] - fn get_preferred_host_fails() { - let mut target = Target { - hosts: vec![ - "un.reachable.1".into(), - "un.reachable.2".into(), - "un.reachable.3".into(), - "un.reachable.4".into(), - "un.reachable.5".into(), - ], - ..Default::default() - }; - - assert_ne!( - target.get_preferred_host().unwrap().to_string(), - "un.reachable.5" - ); - - for i in 1..=5 { - assert_eq!( - target.get_preferred_host().unwrap().to_string(), - format!("un.reachable.{i}") - ); - target.host_failed(); - } - - for _ in 0..5 { - assert_matches!( - target.get_preferred_host(), - Err(HiveLibError::NetworkError(NetworkError::HostsExhausted)) - ); - } - } + use std::{assert_matches::assert_matches, env}; #[test] fn test_ssh_opts() { @@ -850,17 +356,14 @@ mod tests { "KbdInteractiveAuthentication=no".to_string(), ]; - assert_eq!( - target.create_ssh_args(subcommand_modifiers, false).unwrap(), - args - ); + assert_eq!(target.create_ssh_args(subcommand_modifiers).unwrap(), args); assert_eq!( target.create_ssh_opts(subcommand_modifiers).unwrap(), args.join(" ") ); assert_eq!( - target.create_ssh_args(subcommand_modifiers, false).unwrap(), + target.create_ssh_args(subcommand_modifiers).unwrap(), [ "-l".to_string(), target.user.to_string(), @@ -876,7 +379,7 @@ mod tests { ); assert_eq!( - target.create_ssh_args(subcommand_modifiers, true).unwrap(), + target.create_ssh_args(subcommand_modifiers).unwrap(), [ "-l".to_string(), target.user.to_string(), @@ -893,32 +396,59 @@ mod tests { // forced non interactive is the same as --non-interactive assert_eq!( - target.create_ssh_args(subcommand_modifiers, true).unwrap(), + target.create_ssh_args(subcommand_modifiers).unwrap(), target - .create_ssh_args( - SubCommandModifiers { - non_interactive: true, - ..Default::default() - }, - false - ) + .create_ssh_args(SubCommandModifiers { + non_interactive: true, + ..Default::default() + }) .unwrap() ); } - #[tokio::test] - async fn context_quits_sigint() { - let location = location!(get_test_path!()); - let mut node = Node::default(); + #[test] + fn target_fails_increments() { + let mut target = Target::from_host("localhost"); - let name = &Name(function_name!().into()); - let context = Context::create_test_context(location, name, &mut node); - context - .should_quit - .store(true, std::sync::atomic::Ordering::Relaxed); - let executor = GoalExecutor::new(context); - let status = executor.execute().await; + assert_eq!(target.current_host, 0); - assert_matches!(status, Err(HiveLibError::Sigint)); + for i in 0..100 { + target.host_failed(); + assert_eq!(target.current_host, i + 1); + } + } + + #[test] + fn get_preferred_host_fails() { + let mut target = Target { + hosts: vec![ + "un.reachable.1".into(), + "un.reachable.2".into(), + "un.reachable.3".into(), + "un.reachable.4".into(), + "un.reachable.5".into(), + ], + ..Default::default() + }; + + assert_ne!( + target.get_preferred_host().unwrap().to_string(), + "un.reachable.5" + ); + + for i in 1..=5 { + assert_eq!( + target.get_preferred_host().unwrap().to_string(), + format!("un.reachable.{i}") + ); + target.host_failed(); + } + + for _ in 0..5 { + assert_matches!( + target.get_preferred_host(), + Err(HiveLibError::NetworkError(NetworkError::HostsExhausted)) + ); + } } } diff --git a/crates/core/src/hive/plan.rs b/crates/core/src/hive/plan.rs new file mode 100644 index 0000000..291ab36 --- /dev/null +++ b/crates/core/src/hive/plan.rs @@ -0,0 +1,735 @@ +use std::sync::{Arc, atomic::AtomicBool}; + +use tokio::sync::RwLock; + +use crate::{ + SubCommandModifiers, + hive::{ + HiveLocation, + node::{ + ApplyGoal, Context, HandleUnreachable, Name, Node, SharedTarget, Step, StepState, + SwitchToConfigurationGoal, + }, + steps::{ + activate::SwitchToConfiguration, + build::Build, + evaluate::Evaluate, + keys::{Keys, PushKeyAgent, UploadKeyAt}, + ping::Ping, + push::{PushBuildOutput, PushEvaluatedOutput}, + }, + }, +}; + +pub struct NodePlan { + pub context: Context, + pub steps: Vec, + pub greedy_evaluate: bool, + pub ignore_failed_ping: bool, +} + +#[allow(clippy::struct_excessive_bools)] +pub struct ApplyGoalArgs { + pub goal: ApplyGoal, + pub should_apply_locally: bool, + pub no_keys: bool, + pub substitute_on_destination: bool, + pub reboot: bool, + pub host_platform: Arc, + pub handle_unreachable: HandleUnreachable, +} + +pub enum Goal { + Apply(ApplyGoalArgs), + Build, +} + +fn apply_plan_keys( + args: &ApplyGoalArgs, + node: &Node, + target: &SharedTarget, +) -> (Vec, Vec) { + let ApplyGoalArgs { + goal, + substitute_on_destination, + should_apply_locally, + host_platform, + .. + } = args; + let mut front_steps = Vec::new(); + let mut end_steps = Vec::new(); + + let (pre_keys, post_keys) = match goal { + ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) => node + .keys + .clone() + .into_iter() + .partition(|x| matches!(x.upload_at, UploadKeyAt::PreActivation)), + ApplyGoal::Keys => (node.keys.clone(), Vec::new()), + ApplyGoal::Build | ApplyGoal::Push | ApplyGoal::SwitchToConfiguration(_) => { + unreachable!("apply_plan_keys called with non-key goal: {:?}", goal) + } + }; + + // only push key agent if there are any keys at all + if !pre_keys.is_empty() || !post_keys.is_empty() { + front_steps.push(Step::PushKeyAgent(PushKeyAgent { + substitute_on_destination: *substitute_on_destination, + host_platform: host_platform.clone(), + target: if *should_apply_locally { + None + } else { + Some(target.clone()) + }, + })); + } + + if !pre_keys.is_empty() { + front_steps.push(Step::Keys(Keys { + keys: pre_keys, + target: if *should_apply_locally { + None + } else { + Some(target.clone()) + }, + privilege_escalation_command: node.privilege_escalation_command.clone(), + })); + } + + if !post_keys.is_empty() { + end_steps.push(Step::Keys(Keys { + keys: post_keys, + target: if *should_apply_locally { + None + } else { + Some(target.clone()) + }, + privilege_escalation_command: node.privilege_escalation_command.clone(), + })); + } + + (front_steps, end_steps) +} + +fn apply_plan( + args: &ApplyGoalArgs, + node: &Node, + name: &Name, + modifiers: SubCommandModifiers, + hive_location: Arc, + should_quit: Arc, +) -> NodePlan { + let ApplyGoalArgs { + goal, + should_apply_locally, + no_keys, + substitute_on_destination, + reboot, + handle_unreachable, + .. + } = args; + + let mut steps: Vec = Vec::new(); + let mut end: Vec = Vec::new(); + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); + + if !*should_apply_locally { + steps.push(Step::Ping(Ping { + target: target.clone(), + })); + } + + if !*no_keys + && matches!( + &goal, + ApplyGoal::Keys | ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) + ) + { + let (pre, post) = apply_plan_keys(args, node, &target); + steps.extend(pre); + end.extend(post); + } + + if !matches!(goal, ApplyGoal::Keys) { + steps.push(Step::Evaluate(Evaluate)); + } + + if !matches!(goal, ApplyGoal::Keys) + && !should_apply_locally + && (node.build_remotely || matches!(goal, ApplyGoal::Push)) + { + steps.push(Step::PushEvaluatedOutput(PushEvaluatedOutput { + substitute_on_destination: *substitute_on_destination, + target: target.clone(), + })); + } + + if !matches!(goal, ApplyGoal::Keys | ApplyGoal::Push) { + steps.push(Step::Build(Build { + target: if node.build_remotely && !*should_apply_locally { + Some(target.clone()) + } else { + None + }, + })); + } + + if !node.build_remotely + && !should_apply_locally + && !matches!(goal, ApplyGoal::Keys | ApplyGoal::Push) + { + steps.push(Step::PushBuildOutput(PushBuildOutput { + substitute_on_destination: *substitute_on_destination, + target: target.clone(), + })); + } + + if let ApplyGoal::SwitchToConfiguration(goal) = goal { + steps.push(Step::SwitchToConfiguration(SwitchToConfiguration { + goal: *goal, + reboot: *reboot, + target: if *should_apply_locally { + None + } else { + Some(target.clone()) + }, + privilege_escalation_command: node.privilege_escalation_command.clone(), + })); + } + + steps.extend(end); + + NodePlan { + context: Context { + state: StepState::default(), + name: name.clone(), + hive_location, + modifiers, + should_quit, + }, + steps, + greedy_evaluate: !matches!(&goal, ApplyGoal::Keys), + ignore_failed_ping: matches!(handle_unreachable, HandleUnreachable::Ignore), + } +} + +#[allow(clippy::too_many_lines)] +pub fn plan_for_node( + node: &Node, + name: Name, + goal: &'_ Goal, + hive_location: Arc, + modifiers: &SubCommandModifiers, + should_quit: Arc, +) -> NodePlan { + match goal { + Goal::Build => NodePlan { + context: Context { + state: StepState::default(), + modifiers: *modifiers, + hive_location, + should_quit, + name, + }, + steps: vec![ + Step::Evaluate(Evaluate), + Step::Build(Build { target: None }), + ], + greedy_evaluate: true, + ignore_failed_ping: false, + }, + Goal::Apply(args) => apply_plan(args, node, &name, *modifiers, hive_location, should_quit), + } +} + +#[cfg(test)] +mod tests { + use tokio::sync::RwLock; + + use crate::{ + SubCommandModifiers, function_name, get_test_path, + hive::{ + node::{ + ApplyGoal, HandleUnreachable, Name, Node, SharedTarget, Step, + SwitchToConfigurationGoal, + }, + plan::{ApplyGoalArgs, Goal, plan_for_node}, + steps::{ + activate::SwitchToConfiguration, + build::Build, + evaluate::Evaluate, + keys::{Key, Keys, PushKeyAgent, Source, UploadKeyAt}, + ping::Ping, + push::PushEvaluatedOutput, + }, + }, + location, + }; + use std::path::PathBuf; + use std::{ + env, + sync::{Arc, atomic::AtomicBool}, + }; + + fn new_key(upload_at: &UploadKeyAt) -> Key { + Key { + upload_at: upload_at.clone(), + source: Source::String(match upload_at { + UploadKeyAt::PreActivation => "pre".into(), + UploadKeyAt::PostActivation => "post".into(), + UploadKeyAt::NoFilter => "none".into(), + }), + ..Default::default() + } + } + + #[tokio::test] + async fn order_build() { + let location = location!(get_test_path!()); + let node = Node { + build_remotely: false, + ..Default::default() + }; + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let plan = plan_for_node( + &node, + name.clone(), + &Goal::Build, + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Evaluate.into(), + Build { target: None }.into() // TODO: this was previously used in an old test, may lose + // coverage by deleting it. + // Ping { }.into(), + // PushKeyAgent { host_platform: "x86_64-linux".into(), substitute_on_destination: true, target: Target::default() }.into(), + // Keys { .. }.into(), + // crate::hive::steps::evaluate::Evaluate.into(), + // crate::hive::steps::build::Build { .. }.into(), + // crate::hive::steps::push::PushBuildOutput { .. }.into(), + // SwitchToConfiguration { .. }.into(), + // Keys { + // filter: UploadKeyAt::PostActivation + // } + // .into(), + ] + ); + } + + #[tokio::test] + async fn order_apply_build() { + let location = location!(get_test_path!()); + let node = Node { + build_remotely: true, + ..Default::default() + }; + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); + let plan = plan_for_node( + &node, + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::Build, + should_apply_locally: false, + no_keys: true, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Ping { + target: target.clone() + } + .into(), + crate::hive::steps::evaluate::Evaluate.into(), + crate::hive::steps::push::PushEvaluatedOutput { + substitute_on_destination: true, + target: target.clone() + } + .into(), + crate::hive::steps::build::Build { + target: Some(target.clone()) + } + .into(), + ] + ); + + let node = Node { + build_remotely: false, + ..Default::default() + }; + let plan = plan_for_node( + &node, + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::Build, + should_apply_locally: false, + no_keys: true, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Ping { + target: target.clone() + } + .into(), + crate::hive::steps::evaluate::Evaluate.into(), + crate::hive::steps::build::Build { target: None }.into(), + crate::hive::steps::push::PushBuildOutput { + substitute_on_destination: true, + target: target.clone() + } + .into(), + ] + ); + } + + #[tokio::test] + async fn order_keys_only() { + let location = location!(get_test_path!()); + let node = Node { + keys: vec![ + new_key(&UploadKeyAt::PreActivation).into(), + new_key(&UploadKeyAt::PostActivation).into(), + new_key(&UploadKeyAt::PreActivation).into(), + new_key(&UploadKeyAt::PostActivation).into(), + ], + ..Default::default() + }; + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); + let plan_apply_keys = plan_for_node( + &node.clone(), + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::Keys, + should_apply_locally: false, + no_keys: false, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan_apply_keys.steps, + vec![ + Ping { + target: target.clone() + } + .into(), + PushKeyAgent { + substitute_on_destination: true, + target: Some(target.clone()), + host_platform: node.host_platform.clone() + } + .into(), + Keys { + target: Some(target.clone()), + // test that all keys are included + keys: node.keys.clone(), + privilege_escalation_command: node.privilege_escalation_command.clone() + } + .into(), + ] + ); + } + + #[tokio::test] + async fn order_key_split() { + let location = location!(get_test_path!()); + let node = Node { + keys: vec![ + new_key(&UploadKeyAt::PreActivation).into(), + new_key(&UploadKeyAt::PostActivation).into(), + new_key(&UploadKeyAt::PreActivation).into(), + new_key(&UploadKeyAt::PostActivation).into(), + ], + ..Default::default() + }; + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + + // Test that keys are split by their `upload_at`, also tests that key + // step's `target` abides by should_apply_locally + let plan_activate_with_keys = plan_for_node( + &node, + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::SwitchToConfiguration( + crate::hive::node::SwitchToConfigurationGoal::Switch, + ), + should_apply_locally: true, + no_keys: false, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan_activate_with_keys + .steps + .into_iter() + .filter(|x| matches!( + x, + Step::Keys(Keys { .. }) | Step::PushKeyAgent(PushKeyAgent { .. }) + )) + .collect::>(), + vec![ + PushKeyAgent { + substitute_on_destination: true, + target: None, + host_platform: node.host_platform.clone() + } + .into(), + Keys { + target: None, + keys: node + .keys + .iter() + .filter(|key| matches!(key.upload_at, UploadKeyAt::PreActivation)) + .cloned() + .collect::>(), + privilege_escalation_command: node.privilege_escalation_command.clone() + } + .into(), + Keys { + target: None, + keys: node + .keys + .iter() + .filter(|key| matches!(key.upload_at, UploadKeyAt::PostActivation)) + .cloned() + .collect::>(), + privilege_escalation_command: node.privilege_escalation_command.clone() + } + .into(), + ] + ); + } + + #[tokio::test] + async fn order_push_only() { + let location = location!(get_test_path!()); + let node = Node::default(); + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); + let plan = plan_for_node( + &node.clone(), + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::Push, + should_apply_locally: false, + no_keys: false, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Ping { + target: target.clone() + } + .into(), + Evaluate.into(), + PushEvaluatedOutput { + substitute_on_destination: true, + target: target.clone() + } + .into() + ] + ); + } + + #[tokio::test] + async fn order_remote_build() { + let location = location!(get_test_path!()); + let node = Node { + build_remotely: true, + ..Default::default() + }; + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); + let plan = plan_for_node( + &node.clone(), + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), + should_apply_locally: false, + no_keys: false, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Ping { + target: target.clone() + } + .into(), + Evaluate.into(), + PushEvaluatedOutput { + substitute_on_destination: true, + target: target.clone() + } + .into(), + Build { + target: Some(target.clone()) + } + .into(), + SwitchToConfiguration { + goal: SwitchToConfigurationGoal::Switch, + reboot: false, + target: Some(target.clone()), + privilege_escalation_command: node.privilege_escalation_command, + } + .into(), + ] + ); + } + + #[tokio::test] + async fn order_nokeys() { + let location = location!(get_test_path!()); + let node = Node { + keys: vec![Key::default().into(), Key::default().into()], + build_remotely: true, + ..Default::default() + }; + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); + let plan = plan_for_node( + &node.clone(), + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), + should_apply_locally: false, + no_keys: true, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Ping { + target: target.clone() + } + .into(), + Evaluate.into(), + PushEvaluatedOutput { + substitute_on_destination: true, + target: target.clone() + } + .into(), + Build { + target: Some(target.clone()) + } + .into(), + SwitchToConfiguration { + goal: SwitchToConfigurationGoal::Switch, + reboot: false, + target: Some(target.clone()), + privilege_escalation_command: node.privilege_escalation_command, + } + .into(), + ] + ); + } + + #[tokio::test] + async fn order_should_apply_locally() { + let location = location!(get_test_path!()); + let node = Node::default(); + let name = &Name(function_name!().into()); + let should_quit = Arc::new(AtomicBool::new(false)); + let plan = plan_for_node( + &node.clone(), + name.clone(), + &Goal::Apply(ApplyGoalArgs { + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), + should_apply_locally: true, + no_keys: true, + substitute_on_destination: true, + reboot: false, + host_platform: "x86_64-linux".into(), + handle_unreachable: HandleUnreachable::default(), + }), + location.clone().into(), + &SubCommandModifiers::default(), + should_quit.clone(), + ); + + assert_eq!( + plan.steps, + vec![ + Evaluate.into(), + Build { target: None }.into(), + SwitchToConfiguration { + goal: SwitchToConfigurationGoal::Switch, + reboot: false, + target: None, + privilege_escalation_command: node.privilege_escalation_command, + } + .into(), + ] + ); + } +} diff --git a/crates/core/src/hive/steps/activate.rs b/crates/core/src/hive/steps/activate.rs index 9483fdd..99fbfb9 100644 --- a/crates/core/src/hive/steps/activate.rs +++ b/crates/core/src/hive/steps/activate.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-or-later // Copyright 2024-2025 wire Contributors -use std::fmt::Display; +use std::{fmt::Display, sync::Arc}; use tracing::{error, info, instrument, warn}; @@ -9,11 +9,17 @@ use crate::{ HiveLibError, commands::{CommandArguments, WireCommandChip, builder::CommandStringBuilder, run_command}, errors::{ActivationError, NetworkError}, - hive::node::{Context, ExecuteStep, Goal, Objective, SwitchToConfigurationGoal}, + hive::node::{Context, ExecuteStep, SharedTarget, SwitchToConfigurationGoal}, }; -#[derive(Debug, PartialEq)] -pub struct SwitchToConfiguration; +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub struct SwitchToConfiguration { + pub goal: SwitchToConfigurationGoal, + pub reboot: bool, + pub target: Option, + pub privilege_escalation_command: Arc>>, +} impl Display for SwitchToConfiguration { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -21,98 +27,75 @@ impl Display for SwitchToConfiguration { } } -async fn wait_for_ping(ctx: &Context<'_>) -> Result<(), HiveLibError> { - let host = ctx.node.target.get_preferred_host()?; - let mut result = ctx.node.ping(ctx.modifiers).await; +async fn wait_for_ping(target: &SharedTarget, ctx: &Context) -> Result<(), HiveLibError> { + let target = target.0.read().await; + let host = target.get_preferred_host()?; - for num in 0..2 { + for num in 0..3 { warn!("Trying to ping {host} (attempt {}/3)", num + 1); - result = ctx.node.ping(ctx.modifiers).await; + let result = target.ping(ctx.modifiers).await; if result.is_ok() { info!("Regained connection to {} via {host}", ctx.name); - break; + return Ok(()); } } - result + Err(HiveLibError::NetworkError(NetworkError::HostsExhausted)) } -async fn set_profile( - goal: SwitchToConfigurationGoal, - built_path: &String, - ctx: &Context<'_>, -) -> Result<(), HiveLibError> { - info!("Setting profiles in anticipation for switch-to-configuration {goal}"); - - let mut command_string = CommandStringBuilder::new("nix-env"); - command_string.args(&["-p", "/nix/var/nix/profiles/system", "--set"]); - command_string.arg(built_path); - - let Objective::Apply(apply_objective) = ctx.objective else { - unreachable!() - }; - - let child = run_command( - &CommandArguments::new(command_string, ctx.modifiers) - .mode(crate::commands::ChildOutputMode::Nix) - .execute_on_remote(if apply_objective.should_apply_locally { - None - } else { - Some(&ctx.node.target) - }) - .elevated(ctx.node), - ) - .await?; - - let _ = child - .wait_till_success() - .await - .map_err(HiveLibError::CommandError)?; - - info!("Set system profile"); - - Ok(()) -} +impl SwitchToConfiguration { + async fn set_profile(&self, built_path: &String, ctx: &Context) -> Result<(), HiveLibError> { + info!( + "Setting profiles in anticipation for switch-to-configuration {}", + self.goal + ); -impl ExecuteStep for SwitchToConfiguration { - fn should_execute(&self, ctx: &Context) -> bool { - let Objective::Apply(apply_objective) = ctx.objective else { - return false; - }; + let mut command_string = CommandStringBuilder::new("nix-env"); + command_string.args(&["-p", "/nix/var/nix/profiles/system", "--set"]); + command_string.arg(built_path); + + let child = run_command( + &CommandArguments::new(command_string, ctx.modifiers) + .mode(crate::commands::ChildOutputMode::Nix) + .execute_on_remote(self.target.clone()) + .privileged(&self.privilege_escalation_command), + ) + .await?; + + let _ = child + .wait_till_success() + .await + .map_err(HiveLibError::CommandError)?; - matches!(apply_objective.goal, Goal::SwitchToConfiguration(..)) + info!("Set system profile"); + + Ok(()) } +} +impl ExecuteStep for SwitchToConfiguration { #[allow(clippy::too_many_lines)] #[instrument(skip_all, name = "activate")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let built_path = ctx.state.build.as_ref().unwrap(); - let Objective::Apply(apply_objective) = ctx.objective else { - unreachable!() - }; - - let Goal::SwitchToConfiguration(goal) = &apply_objective.goal else { - unreachable!("Cannot reach as guarded by should_execute") - }; - if matches!( - goal, + self.goal, // switch profile if switch or boot // https://github.com/NixOS/nixpkgs/blob/a2c92aa34735a04010671e3378e2aa2d109b2a72/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/services.py#L224 SwitchToConfigurationGoal::Switch | SwitchToConfigurationGoal::Boot ) { - set_profile(*goal, built_path, ctx).await?; + self.set_profile(built_path, ctx).await?; } - info!("Running switch-to-configuration {goal}"); + info!("Running switch-to-configuration {}", self.goal); let mut command_string = CommandStringBuilder::new(format!("{built_path}/bin/switch-to-configuration")); - command_string.arg(match goal { + command_string.arg(match self.goal { SwitchToConfigurationGoal::Switch => "switch", SwitchToConfigurationGoal::Boot => "boot", SwitchToConfigurationGoal::Test => "test", @@ -121,12 +104,8 @@ impl ExecuteStep for SwitchToConfiguration { let child = run_command( &CommandArguments::new(command_string, ctx.modifiers) - .execute_on_remote(if apply_objective.should_apply_locally { - None - } else { - Some(&ctx.node.target) - }) - .elevated(ctx.node) + .execute_on_remote(self.target.clone()) + .privileged(&self.privilege_escalation_command) .log_stdout(), ) .await?; @@ -135,23 +114,23 @@ impl ExecuteStep for SwitchToConfiguration { match result { Ok(_) => { - if !apply_objective.reboot { + if !self.reboot { return Ok(()); } - if apply_objective.should_apply_locally { + let Some(ref target) = self.target else { error!("Refusing to reboot local machine!"); return Ok(()); - } + }; warn!("Rebooting {name}!", name = ctx.name); let reboot = run_command( &CommandArguments::new("reboot now", ctx.modifiers) .log_stdout() - .execute_on_remote(Some(&ctx.node.target)) - .elevated(ctx.node), + .execute_on_remote(Some(target.clone())) + .privileged(&self.privilege_escalation_command), ) .await?; @@ -164,19 +143,21 @@ impl ExecuteStep for SwitchToConfiguration { info!("Rebooted {name}, waiting to reconnect...", name = ctx.name); - if wait_for_ping(ctx).await.is_ok() { + if wait_for_ping(target, ctx).await.is_ok() { return Ok(()); } + let target = target.0.read().await; + error!( "Failed to get regain connection to {name} via {host} after reboot.", name = ctx.name, - host = ctx.node.target.get_preferred_host()? + host = target.get_preferred_host()? ); return Err(HiveLibError::NetworkError( NetworkError::HostUnreachableAfterReboot( - ctx.node.target.get_preferred_host()?.to_string(), + target.get_preferred_host()?.to_string(), ), )); } @@ -188,29 +169,42 @@ impl ExecuteStep for SwitchToConfiguration { // Bail if the command couldn't of broken the system // and don't try to regain connection to localhost - if matches!(goal, SwitchToConfigurationGoal::DryActivate) - || apply_objective.should_apply_locally - { + let Some(target) = self + .target + .as_ref() + .filter(|_| !matches!(self.goal, SwitchToConfigurationGoal::DryActivate)) + else { return Err(HiveLibError::ActivationError( - ActivationError::SwitchToConfigurationError(*goal, ctx.name.clone(), error), + ActivationError::SwitchToConfigurationError( + self.goal, + ctx.name.clone(), + error, + ), )); - } + }; - if wait_for_ping(ctx).await.is_ok() { + if wait_for_ping(target, ctx).await.is_ok() { return Err(HiveLibError::ActivationError( - ActivationError::SwitchToConfigurationError(*goal, ctx.name.clone(), error), + ActivationError::SwitchToConfigurationError( + self.goal, + ctx.name.clone(), + error, + ), )); } + let target = target.0.read().await; + error!( "Failed to get regain connection to {name} via {host} after {goal} activation.", name = ctx.name, - host = ctx.node.target.get_preferred_host()? + host = target.get_preferred_host()?, + goal = self.goal ); return Err(HiveLibError::NetworkError( NetworkError::HostUnreachableAfterReboot( - ctx.node.target.get_preferred_host()?.to_string(), + target.get_preferred_host()?.to_string(), ), )); } diff --git a/crates/core/src/hive/steps/build.rs b/crates/core/src/hive/steps/build.rs index 1a9a356..1068cbf 100644 --- a/crates/core/src/hive/steps/build.rs +++ b/crates/core/src/hive/steps/build.rs @@ -11,11 +11,14 @@ use crate::{ CommandArguments, Either, WireCommandChip, builder::CommandStringBuilder, run_command_with_env, }, - hive::node::{Context, ExecuteStep, Goal, Objective}, + hive::node::{Context, ExecuteStep, SharedTarget}, }; -#[derive(Debug, PartialEq)] -pub struct Build; +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub struct Build { + pub(crate) target: Option, +} impl Display for Build { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -24,17 +27,8 @@ impl Display for Build { } impl ExecuteStep for Build { - fn should_execute(&self, ctx: &Context) -> bool { - match ctx.objective { - Objective::Apply(apply_objective) => { - !matches!(apply_objective.goal, Goal::Keys | Goal::Push) - } - Objective::BuildLocally => true, - } - } - #[instrument(skip_all, name = "build")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let top_level = ctx.state.evaluation.as_ref().unwrap(); let mut command_string = CommandStringBuilder::nix(); @@ -51,16 +45,7 @@ impl ExecuteStep for Build { let status = run_command_with_env( &CommandArguments::new(command_string, ctx.modifiers) // build remotely if asked for AND we arent applying locally - .execute_on_remote( - if ctx.node.build_remotely - && let Objective::Apply(apply_objective) = ctx.objective - && !apply_objective.should_apply_locally - { - Some(&ctx.node.target) - } else { - None - }, - ) + .execute_on_remote(self.target.clone()) .mode(crate::commands::ChildOutputMode::Nix) .log_stdout(), std::collections::HashMap::new(), diff --git a/crates/core/src/hive/steps/cleanup.rs b/crates/core/src/hive/steps/cleanup.rs deleted file mode 100644 index f8964f0..0000000 --- a/crates/core/src/hive/steps/cleanup.rs +++ /dev/null @@ -1,28 +0,0 @@ -// SPDX-License-Identifier: AGPL-3.0-or-later -// Copyright 2024-2025 wire Contributors - -use std::fmt::Display; - -use crate::{ - errors::HiveLibError, - hive::node::{Context, ExecuteStep}, -}; - -#[derive(PartialEq, Debug)] -pub(crate) struct CleanUp; - -impl Display for CleanUp { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Clean up") - } -} - -impl ExecuteStep for CleanUp { - fn should_execute(&self, _ctx: &Context) -> bool { - false - } - - async fn execute(&self, _ctx: &mut Context<'_>) -> Result<(), HiveLibError> { - Ok(()) - } -} diff --git a/crates/core/src/hive/steps/evaluate.rs b/crates/core/src/hive/steps/evaluate.rs index 72b4764..2573b8a 100644 --- a/crates/core/src/hive/steps/evaluate.rs +++ b/crates/core/src/hive/steps/evaluate.rs @@ -7,7 +7,7 @@ use tracing::instrument; use crate::{ HiveLibError, - hive::node::{Context, ExecuteStep, Goal, Objective}, + hive::node::{Context, ExecuteStep}, }; #[derive(Debug, PartialEq)] @@ -20,15 +20,8 @@ impl Display for Evaluate { } impl ExecuteStep for Evaluate { - fn should_execute(&self, ctx: &Context) -> bool { - match ctx.objective { - Objective::Apply(apply_objective) => !matches!(apply_objective.goal, Goal::Keys), - Objective::BuildLocally => true, - } - } - #[instrument(skip_all, name = "eval")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let rx = ctx.state.evaluation_rx.take().unwrap(); ctx.state.evaluation = Some(rx.await.unwrap()?); diff --git a/crates/core/src/hive/steps/keys.rs b/crates/core/src/hive/steps/keys.rs index e3d896b..0a40568 100644 --- a/crates/core/src/hive/steps/keys.rs +++ b/crates/core/src/hive/steps/keys.rs @@ -4,7 +4,6 @@ use base64::Engine; use base64::prelude::BASE64_STANDARD; use futures::future::join_all; -use im::Vector; use itertools::{Itertools, Position}; use owo_colors::OwoColorize; use prost::Message; @@ -19,6 +18,7 @@ use std::path::PathBuf; use std::pin::Pin; use std::process::Stdio; use std::str::from_utf8; +use std::sync::Arc; use std::vec::IntoIter; use tokio::io::AsyncReadExt as _; use tokio::process::Command; @@ -31,7 +31,7 @@ use crate::commands::builder::CommandStringBuilder; use crate::commands::common::push; use crate::commands::{CommandArguments, WireCommandChip, run_command}; use crate::errors::KeyError; -use crate::hive::node::{Context, ExecuteStep, Goal, Objective, Push, SwitchToConfigurationGoal}; +use crate::hive::node::{Context, ExecuteStep, Push, SharedTarget}; #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Hash)] #[serde(tag = "t", content = "c")] @@ -179,16 +179,25 @@ async fn process_key(key: &Key) -> Result<(wire_key_agent::keys::KeySpec, Vec>, + pub target: Option, + pub privilege_escalation_command: Arc>>, +} + +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub struct PushKeyAgent { + pub substitute_on_destination: bool, + pub host_platform: Arc, + pub target: Option, } -#[derive(Debug, PartialEq)] -pub struct PushKeyAgent; impl Display for Keys { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Upload key @ {:?}", self.filter) + write!(f, "Upload {} key(s)", self.keys.len()) } } @@ -225,32 +234,11 @@ where } impl ExecuteStep for Keys { - fn should_execute(&self, ctx: &Context) -> bool { - let Objective::Apply(apply_objective) = ctx.objective else { - return false; - }; - - if apply_objective.no_keys { - return false; - } - - // should execute if no filter, and the goal is keys. - // otherwise, only execute if the goal is switch and non-nofilter - matches!( - (&self.filter, &apply_objective.goal), - (UploadKeyAt::NoFilter, Goal::Keys) - | ( - UploadKeyAt::PreActivation | UploadKeyAt::PostActivation, - Goal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) - ) - ) - } - #[instrument(skip_all, name = "keys")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let agent_directory = ctx.state.key_agent_directory.as_ref().unwrap(); - let mut keys = self.select_keys(&ctx.node.keys).await?; + let mut keys = self.select_keys(&self.keys).await?; if keys.peek().is_none() { debug!("Had no keys to push, ending KeyStep early."); @@ -260,18 +248,10 @@ impl ExecuteStep for Keys { let command_string = CommandStringBuilder::new(format!("{agent_directory}/bin/wire-key-agent")); - let Objective::Apply(apply_objective) = ctx.objective else { - unreachable!() - }; - let mut child = run_command( &CommandArguments::new(command_string, ctx.modifiers) - .execute_on_remote(if apply_objective.should_apply_locally { - None - } else { - Some(&ctx.node.target) - }) - .elevated(ctx.node) + .execute_on_remote(self.target.clone()) + .privileged(&self.privilege_escalation_command) .keep_stdin_open() .log_stdout(), ) @@ -306,17 +286,14 @@ impl ExecuteStep for Keys { impl Keys { async fn select_keys( &self, - keys: &Vector, + keys: &[Arc], ) -> Result)>>, HiveLibError> { - let futures = keys - .iter() - .filter(|key| self.filter == UploadKeyAt::NoFilter || (key.upload_at == self.filter)) - .map(|key| async move { - process_key(key) - .await - .map_err(|err| HiveLibError::KeyError(key.name.clone(), err)) - }); + let futures = keys.iter().map(|key| async move { + process_key(key) + .await + .map_err(|err| HiveLibError::KeyError(key.name.clone(), err)) + }); Ok(join_all(futures) .await @@ -328,26 +305,11 @@ impl Keys { } impl ExecuteStep for PushKeyAgent { - fn should_execute(&self, ctx: &Context) -> bool { - let Objective::Apply(apply_objective) = ctx.objective else { - return false; - }; - - if apply_objective.no_keys { - return false; - } - - matches!( - &apply_objective.goal, - Goal::Keys | Goal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) - ) - } - #[instrument(skip_all, name = "push_agent")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let arg_name = format!( "WIRE_KEY_AGENT_{platform}", - platform = ctx.node.host_platform.replace('-', "_") + platform = self.host_platform.replace('-', "_") ); let agent_directory = match env::var_os(&arg_name) { @@ -359,12 +321,14 @@ impl ExecuteStep for PushKeyAgent { ), }; - let Objective::Apply(apply_objective) = ctx.objective else { - unreachable!() - }; - - if !apply_objective.should_apply_locally { - push(ctx, Push::Path(&agent_directory)).await?; + if let Some(ref target) = self.target { + push( + ctx, + target, + Push::Path(&agent_directory), + self.substitute_on_destination, + ) + .await?; } ctx.state.key_agent_directory = Some(agent_directory); @@ -372,70 +336,3 @@ impl ExecuteStep for PushKeyAgent { Ok(()) } } - -#[cfg(test)] -mod tests { - use im::Vector; - - use crate::hive::steps::keys::{Key, Keys, UploadKeyAt, process_key}; - - fn new_key(upload_at: &UploadKeyAt) -> Key { - Key { - upload_at: upload_at.clone(), - source: super::Source::String(match upload_at { - UploadKeyAt::PreActivation => "pre".into(), - UploadKeyAt::PostActivation => "post".into(), - UploadKeyAt::NoFilter => "none".into(), - }), - ..Default::default() - } - } - - #[tokio::test] - async fn key_filtering() { - let keys = Vector::from(vec![ - new_key(&UploadKeyAt::PreActivation), - new_key(&UploadKeyAt::PostActivation), - new_key(&UploadKeyAt::PreActivation), - new_key(&UploadKeyAt::PostActivation), - ]); - - for (_, buf) in (Keys { - filter: crate::hive::steps::keys::UploadKeyAt::PreActivation, - }) - .select_keys(&keys) - .await - .unwrap() - { - assert_eq!(String::from_utf8_lossy(&buf), "pre"); - } - - for (_, buf) in (Keys { - filter: crate::hive::steps::keys::UploadKeyAt::PostActivation, - }) - .select_keys(&keys) - .await - .unwrap() - { - assert_eq!(String::from_utf8_lossy(&buf), "post"); - } - - // test that NoFilter processes all keys. - let processed_all = - futures::future::join_all(keys.iter().map(async |x| process_key(x).await)) - .await - .iter() - .flatten() - .cloned() - .collect::>(); - let no_filter = (Keys { - filter: crate::hive::steps::keys::UploadKeyAt::NoFilter, - }) - .select_keys(&keys) - .await - .unwrap() - .collect::>(); - - assert_eq!(processed_all, no_filter); - } -} diff --git a/crates/core/src/hive/steps/mod.rs b/crates/core/src/hive/steps/mod.rs index 3fbc77e..a62b5d4 100644 --- a/crates/core/src/hive/steps/mod.rs +++ b/crates/core/src/hive/steps/mod.rs @@ -3,7 +3,6 @@ pub mod activate; pub mod build; -pub mod cleanup; pub mod evaluate; pub mod keys; pub mod ping; diff --git a/crates/core/src/hive/steps/ping.rs b/crates/core/src/hive/steps/ping.rs index fcf31f6..ddda6f1 100644 --- a/crates/core/src/hive/steps/ping.rs +++ b/crates/core/src/hive/steps/ping.rs @@ -7,11 +7,14 @@ use tracing::{Level, event, instrument}; use crate::{ HiveLibError, - hive::node::{Context, ExecuteStep, Objective}, + hive::node::{Context, ExecuteStep, SharedTarget}, }; -#[derive(Debug, PartialEq)] -pub struct Ping; +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub struct Ping { + pub target: SharedTarget, +} impl Display for Ping { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -20,28 +23,22 @@ impl Display for Ping { } impl ExecuteStep for Ping { - fn should_execute(&self, ctx: &Context) -> bool { - let Objective::Apply(apply_objective) = ctx.objective else { - return false; - }; - - !apply_objective.should_apply_locally - } - #[instrument(skip_all, name = "ping")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { loop { + let target = self.target.0.read().await; + event!( Level::INFO, status = "attempting", - host = ctx.node.target.get_preferred_host()?.to_string() + host = target.get_preferred_host()?.to_string() ); - if ctx.node.ping(ctx.modifiers).await.is_ok() { + if target.ping(ctx.modifiers).await.is_ok() { event!( Level::INFO, status = "success", - host = ctx.node.target.get_preferred_host()?.to_string() + host = target.get_preferred_host()?.to_string() ); return Ok(()); } @@ -50,9 +47,12 @@ impl ExecuteStep for Ping { event!( Level::WARN, status = "failed to ping", - host = ctx.node.target.get_preferred_host()?.to_string() + host = target.get_preferred_host()?.to_string() ); - ctx.node.target.host_failed(); + + drop(target); + + self.target.0.write().await.host_failed(); } } } diff --git a/crates/core/src/hive/steps/push.rs b/crates/core/src/hive/steps/push.rs index 06cfc0f..760f1a1 100644 --- a/crates/core/src/hive/steps/push.rs +++ b/crates/core/src/hive/steps/push.rs @@ -8,13 +8,22 @@ use tracing::instrument; use crate::{ HiveLibError, commands::common::push, - hive::node::{Context, ExecuteStep, Goal, Objective}, + hive::node::{Context, ExecuteStep, SharedTarget}, }; -#[derive(Debug, PartialEq)] -pub struct PushEvaluatedOutput; -#[derive(Debug, PartialEq)] -pub struct PushBuildOutput; +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub struct PushEvaluatedOutput { + pub substitute_on_destination: bool, + pub target: SharedTarget, +} + +#[derive(Debug)] +#[cfg_attr(test, derive(PartialEq))] +pub struct PushBuildOutput { + pub substitute_on_destination: bool, + pub target: SharedTarget, +} impl Display for PushEvaluatedOutput { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -29,55 +38,34 @@ impl Display for PushBuildOutput { } impl ExecuteStep for PushEvaluatedOutput { - fn should_execute(&self, ctx: &Context) -> bool { - let Objective::Apply(apply_objective) = ctx.objective else { - return false; - }; - - !matches!(apply_objective.goal, Goal::Keys) - && !apply_objective.should_apply_locally - && (ctx.node.build_remotely | matches!(apply_objective.goal, Goal::Push)) - } - #[instrument(skip_all, name = "push_eval")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let top_level = ctx.state.evaluation.as_ref().unwrap(); - push(ctx, crate::hive::node::Push::Derivation(top_level)).await?; + push( + ctx, + &self.target, + crate::hive::node::Push::Derivation(top_level), + self.substitute_on_destination, + ) + .await?; Ok(()) } } impl ExecuteStep for PushBuildOutput { - fn should_execute(&self, ctx: &Context) -> bool { - let Objective::Apply(apply_objective) = ctx.objective else { - return false; - }; - - if matches!(apply_objective.goal, Goal::Keys | Goal::Push) { - // skip if we are not building - return false; - } - - if ctx.node.build_remotely { - // skip if we are building remotely - return false; - } - - if apply_objective.should_apply_locally { - // skip step if we are applying locally - return false; - } - - true - } - #[instrument(skip_all, name = "push_build")] - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { let built_path = ctx.state.build.as_ref().unwrap(); - push(ctx, crate::hive::node::Push::Path(built_path)).await?; + push( + ctx, + &self.target, + crate::hive::node::Push::Path(built_path), + self.substitute_on_destination, + ) + .await?; Ok(()) } -- 2.51.2