+5

CHANGELOG.md

reviewed

··· 17 17 18 18 - The domain for documentation to be `wire.forall.systems`. The previous URL 19 19 will continue to be available but may redirect in the future. 20 20 + - Refactored node execution to be in two distinct phases, "planning" and 21 21 + "execution". Previously, picking what steps would be run was done on the fly 22 22 + during execution. 23 23 + - Cases where there are no keys to deploy, such as having 0 keys or filtered 24 24 + keys, the "Key" step will not be planned when it previously would have. 20 25 21 26 ### Fixed 22 27

+16 -19

crates/cli/src/apply.rs

reviewed

··· 11 11 use std::sync::atomic::AtomicBool; 12 12 use thiserror::Error; 13 13 use tracing::{error, info}; 14 14 - use wire_core::hive::node::{Context, GoalExecutor, Name, Node, Objective, StepState}; 14 14 + use wire_core::hive::executor::execute; 15 15 + use wire_core::hive::node::{Name, Node}; 16 16 + use wire_core::hive::plan::{Goal, plan_for_node}; 15 17 use wire_core::hive::{Hive, HiveLocation}; 16 18 use wire_core::status::STATUS; 17 19 use wire_core::{SubCommandModifiers, errors::HiveLibError}; ··· 96 98 97 99 pub async fn apply<F>( 98 100 hive: &mut Hive, 99 99 - should_shutdown: Arc<AtomicBool>, 101 101 + should_quit: Arc<AtomicBool>, 100 102 location: HiveLocation, 101 103 args: CommonVerbArgs, 102 104 partition: Partitions, 103 103 - make_objective: F, 105 105 + make_goal: F, 104 106 mut modifiers: SubCommandModifiers, 105 107 ) -> Result<()> 106 108 where 107 107 - F: Fn(&Name, &Node) -> Objective, 109 109 + F: Fn(&Name, &Node) -> Goal, 108 110 { 109 111 let location = Arc::new(location); 110 112 ··· 142 144 .iter_mut() 143 145 .filter(|(name, _)| partitioned_names.contains(name)) 144 146 .map(|(name, node)| { 145 145 - info!("Resolved {:?} to include {}", args.on, name); 147 147 + let goal = make_goal(name, node); 146 148 147 147 - let objective = make_objective(name, node); 148 148 - 149 149 - let context = Context { 149 149 + let plan = plan_for_node( 150 150 node, 151 151 - name, 152 152 - objective, 153 153 - state: StepState::default(), 154 154 - hive_location: location.clone(), 155 155 - modifiers, 156 156 - should_quit: should_shutdown.clone(), 157 157 - }; 158 158 - 159 159 - GoalExecutor::new(context) 160 160 - .execute() 161 161 - .map(move |result| (name, result)) 151 151 + name.clone(), 152 152 + &goal, 153 153 + location.clone(), 154 154 + &modifiers, 155 155 + should_quit.clone(), 156 156 + ); 157 157 + execute(plan).map(move |result| (name, result)) 162 158 }) 163 159 .peekable(); 164 160 ··· 168 164 169 165 let futures = futures::stream::iter(set).buffer_unordered(args.parallel); 170 166 let result = futures.collect::<Vec<_>>().await; 167 167 + 171 168 let (successful, errors): (Vec<_>, Vec<_>) = 172 169 result 173 170 .into_iter()

+3 -1

crates/cli/src/cli.rs

reviewed

··· 11 11 use tokio::runtime::Handle; 12 12 use wire_core::SubCommandModifiers; 13 13 use wire_core::commands::common::get_hive_node_names; 14 14 - use wire_core::hive::node::{Goal as HiveGoal, HandleUnreachable, Name, SwitchToConfigurationGoal}; 14 14 + use wire_core::hive::node::{ 15 15 + ApplyGoal as HiveGoal, HandleUnreachable, Name, SwitchToConfigurationGoal, 16 16 + }; 15 17 use wire_core::hive::{Hive, get_hive_location}; 16 18 17 19 use std::io::IsTerminal;

+6 -5

crates/cli/src/main.rs

reviewed

··· 28 28 use wire_core::commands::common::get_hive_node_names; 29 29 use wire_core::hive::Hive; 30 30 use wire_core::hive::get_hive_location; 31 31 - use wire_core::hive::node::ApplyObjective; 32 32 - use wire_core::hive::node::Objective; 33 31 use wire_core::hive::node::should_apply_locally; 32 32 + use wire_core::hive::plan::ApplyGoalArgs; 33 33 + use wire_core::hive::plan::Goal; 34 34 35 35 #[macro_use] 36 36 extern crate enum_display_derive; ··· 86 86 match args.command { 87 87 cli::Commands::Apply(apply_args) => { 88 88 let mut hive = Hive::new_from_path(&location, cache.clone(), modifiers).await?; 89 89 - let goal: wire_core::hive::node::Goal = apply_args.goal.clone().try_into().unwrap(); 89 89 + let goal = apply_args.goal.clone().try_into().unwrap(); 90 90 91 91 // Respect user's --always-build-local arg 92 92 hive.force_always_local(apply_args.always_build_local)?; ··· 98 98 apply_args.common, 99 99 Partitions::default(), 100 100 |name, node| { 101 101 - Objective::Apply(ApplyObjective { 101 101 + Goal::Apply(ApplyGoalArgs { 102 102 goal, 103 103 no_keys: apply_args.no_keys, 104 104 reboot: apply_args.reboot, ··· 108 108 &name.0, 109 109 ), 110 110 handle_unreachable: apply_args.handle_unreachable.clone().into(), 111 111 + host_platform: node.host_platform.clone(), 111 112 }) 112 113 }, 113 114 modifiers, ··· 123 124 location, 124 125 build_args.common, 125 126 build_args.partition.unwrap_or_default(), 126 126 - |_name, _node| Objective::BuildLocally, 127 127 + |_name, _node| Goal::Build, 127 128 modifiers, 128 129 ) 129 130 .await?;

+13 -11

crates/core/src/commands/common.rs

reviewed

··· 14 14 errors::{CommandError, HiveInitialisationError, HiveLibError}, 15 15 hive::{ 16 16 HiveLocation, 17 17 - node::{Context, Objective, Push}, 17 17 + node::{Context, Push, SharedTarget}, 18 18 }, 19 19 }; 20 20 ··· 28 28 } 29 29 } 30 30 31 31 - pub async fn push(context: &Context<'_>, push: Push<'_>) -> Result<(), HiveLibError> { 31 31 + pub async fn push( 32 32 + context: &Context, 33 33 + target: &SharedTarget, 34 34 + push: Push<'_>, 35 35 + substitute_on_destination: bool, 36 36 + ) -> Result<(), HiveLibError> { 37 37 + let target = target.0.read().await; 38 38 + 32 39 let mut command_string = CommandStringBuilder::nix(); 33 40 34 41 command_string.args(&["--extra-experimental-features", "nix-command", "copy"]); 35 35 - if let Objective::Apply(apply_objective) = context.objective { 36 36 - command_string.opt_arg( 37 37 - apply_objective.substitute_on_destination, 38 38 - "--substitute-on-destination", 39 39 - ); 40 40 - } 42 42 + command_string.opt_arg(substitute_on_destination, "--substitute-on-destination"); 41 43 command_string.arg("--to"); 42 44 command_string.args(&[ 43 45 format!( 44 46 "ssh://{user}@{host}", 45 45 - user = context.node.target.user, 46 46 - host = context.node.target.get_preferred_host()?, 47 47 + user = target.user, 48 48 + host = target.get_preferred_host()?, 47 49 ), 48 50 match push { 49 51 Push::Derivation(drv) => format!("{drv} --derivation"), ··· 56 58 .mode(crate::commands::ChildOutputMode::Nix), 57 59 HashMap::from([( 58 60 "NIX_SSHOPTS".into(), 59 59 - context.node.target.create_ssh_opts(context.modifiers)?, 61 61 + target.create_ssh_opts(context.modifiers)?, 60 62 )]), 61 63 ) 62 64 .await?;

+20 -15

crates/core/src/commands/mod.rs

reviewed

··· 1 1 // SPDX-License-Identifier: AGPL-3.0-or-later 2 2 // Copyright 2024-2025 wire Contributors 3 3 4 4 - use crate::commands::pty::{InteractiveChildChip, interactive_command_with_env}; 5 5 - use std::{collections::HashMap, str::from_utf8, sync::LazyLock}; 4 4 + use crate::{ 5 5 + commands::pty::{InteractiveChildChip, interactive_command_with_env}, 6 6 + hive::node::SharedTarget, 7 7 + }; 8 8 + use std::{ 9 9 + collections::HashMap, 10 10 + str::from_utf8, 11 11 + sync::{Arc, LazyLock}, 12 12 + }; 6 13 7 14 use aho_corasick::AhoCorasick; 8 15 use gjson::Value; ··· 15 22 SubCommandModifiers, 16 23 commands::noninteractive::{NonInteractiveChildChip, non_interactive_command_with_env}, 17 24 errors::{CommandError, HiveLibError}, 18 18 - hive::node::{Node, Target}, 19 25 }; 20 26 21 27 pub(crate) mod builder; ··· 37 43 } 38 44 39 45 #[derive(Debug)] 40 40 - pub(crate) struct CommandArguments<'t, S: AsRef<str>> { 46 46 + pub(crate) struct CommandArguments<S: AsRef<str>> { 41 47 modifiers: SubCommandModifiers, 42 42 - target: Option<&'t Target>, 48 48 + target: Option<SharedTarget>, 43 49 output_mode: ChildOutputMode, 44 50 command_string: S, 45 51 keep_stdin_open: bool, ··· 55 61 .unwrap() 56 62 }); 57 63 58 58 - impl<'a, S: AsRef<str>> CommandArguments<'a, S> { 64 64 + impl<S: AsRef<str>> CommandArguments<S> { 59 65 pub(crate) const fn new(command_string: S, modifiers: SubCommandModifiers) -> Self { 60 66 Self { 61 67 command_string, ··· 68 74 } 69 75 } 70 76 71 71 - pub(crate) const fn execute_on_remote(mut self, target: Option<&'a Target>) -> Self { 77 77 + pub(crate) fn execute_on_remote(mut self, target: Option<SharedTarget>) -> Self { 72 78 self.target = target; 73 79 self 74 80 } ··· 83 89 self 84 90 } 85 91 86 86 - pub(crate) fn elevated(mut self, node: &Node) -> Self { 87 87 - self.privilege_escalation_command = 88 88 - Some(node.privilege_escalation_command.iter().join(" ")); 92 92 + pub(crate) fn privileged(mut self, escalation_command: &[Arc<str>]) -> Self { 93 93 + self.privilege_escalation_command = Some(escalation_command.iter().join(" ")); 89 94 self 90 95 } 91 96 ··· 100 105 } 101 106 102 107 pub(crate) async fn run_command<S: AsRef<str>>( 103 103 - arguments: &CommandArguments<'_, S>, 108 108 + arguments: &CommandArguments<S>, 104 109 ) -> Result<Either<InteractiveChildChip, NonInteractiveChildChip>, HiveLibError> { 105 110 run_command_with_env(arguments, HashMap::new()).await 106 111 } 107 112 108 113 pub(crate) async fn run_command_with_env<S: AsRef<str>>( 109 109 - arguments: &CommandArguments<'_, S>, 114 114 + arguments: &CommandArguments<S>, 110 115 envs: HashMap<String, String>, 111 116 ) -> Result<Either<InteractiveChildChip, NonInteractiveChildChip>, HiveLibError> { 112 117 // use the non interactive command runner when forced ··· 114 119 if arguments.modifiers.non_interactive 115 120 || (arguments.target.is_none() && !arguments.is_elevated()) 116 121 { 117 117 - return Ok(Either::Right(non_interactive_command_with_env( 118 118 - arguments, envs, 119 119 - )?)); 122 122 + return Ok(Either::Right( 123 123 + non_interactive_command_with_env(arguments, envs).await?, 124 124 + )); 120 125 } 121 126 122 127 Ok(Either::Left(

+8 -7

crates/core/src/commands/noninteractive.rs

reviewed

··· 11 11 SubCommandModifiers, 12 12 commands::{ChildOutputMode, CommandArguments, WireCommandChip}, 13 13 errors::{CommandError, HiveLibError}, 14 14 - hive::node::Target, 14 14 + hive::node::SharedTarget, 15 15 }; 16 16 use itertools::Itertools; 17 17 use tokio::{ ··· 32 32 } 33 33 34 34 #[instrument(skip_all, name = "run", fields(elevated = %arguments.is_elevated()))] 35 35 - pub(crate) fn non_interactive_command_with_env<S: AsRef<str>>( 35 35 + pub(crate) async fn non_interactive_command_with_env<S: AsRef<str>>( 36 36 arguments: &CommandArguments<S>, 37 37 envs: HashMap<String, String>, 38 38 ) -> Result<NonInteractiveChildChip, HiveLibError> { 39 39 - let mut command = if let Some(target) = arguments.target { 40 40 - create_sync_ssh_command(target, arguments.modifiers)? 39 39 + let mut command = if let Some(ref target) = arguments.target { 40 40 + create_sync_ssh_command(target, arguments.modifiers).await? 41 41 } else { 42 42 let mut command = Command::new("sh"); 43 43 ··· 187 187 debug!("io_handler: goodbye!"); 188 188 } 189 189 190 190 - fn create_sync_ssh_command( 191 191 - target: &Target, 190 190 + async fn create_sync_ssh_command( 191 191 + target: &SharedTarget, 192 192 modifiers: SubCommandModifiers, 193 193 ) -> Result<Command, HiveLibError> { 194 194 + let target = target.0.read().await; 194 195 let mut command = Command::new("ssh"); 195 195 - command.args(target.create_ssh_args(modifiers, true)?); 196 196 + command.args(target.create_ssh_args(modifiers)?); 196 197 command.arg(target.get_preferred_host()?.to_string()); 197 198 198 199 Ok(command)

+29 -23

crates/core/src/commands/pty/mod.rs

reviewed

··· 2 2 // Copyright 2024-2025 wire Contributors 3 3 4 4 use crate::commands::pty::output::{WatchStdoutArguments, handle_pty_stdout}; 5 5 + use crate::hive::node::SharedTarget; 5 6 use crate::status::STATUS; 6 7 use aho_corasick::PatternID; 7 8 use itertools::Itertools; ··· 29 30 use crate::{ 30 31 commands::{ChildOutputMode, WireCommandChip}, 31 32 errors::HiveLibError, 32 32 - hive::node::Target, 33 33 }; 34 34 35 35 mod input; ··· 85 85 const IO_SUBS: &str = "1> >(while IFS= read -r line; do echo \"#$line\"; done)"; 86 86 87 87 fn create_ending_segment<S: AsRef<str>>( 88 88 - arguments: &CommandArguments<'_, S>, 88 88 + arguments: &CommandArguments<S>, 89 89 needles: &Needles, 90 90 ) -> String { 91 91 let Needles { ··· 110 110 } 111 111 112 112 fn create_starting_segment<S: AsRef<str>>( 113 113 - arguments: &CommandArguments<'_, S>, 113 113 + arguments: &CommandArguments<S>, 114 114 start_needle: &Arc<Vec<u8>>, 115 115 ) -> String { 116 116 if matches!(arguments.output_mode, ChildOutputMode::Interactive) { ··· 125 125 126 126 #[instrument(skip_all, name = "run-int", fields(elevated = %arguments.is_elevated(), mode = ?arguments.output_mode))] 127 127 pub(crate) async fn interactive_command_with_env<S: AsRef<str>>( 128 128 - arguments: &CommandArguments<'_, S>, 128 128 + arguments: &CommandArguments<S>, 129 129 envs: std::collections::HashMap<String, String>, 130 130 ) -> Result<InteractiveChildChip, HiveLibError> { 131 131 - print_authenticate_warning(arguments)?; 131 131 + print_authenticate_warning(arguments).await?; 132 132 133 133 let needles = create_needles(); 134 134 let pty_system = NativePtySystem::default(); ··· 148 148 149 149 debug!("{command_string}"); 150 150 151 151 - let mut command = build_command(arguments, command_string)?; 151 151 + let mut command = build_command(arguments, command_string).await?; 152 152 153 153 // give command all env vars 154 154 for (key, value) in envs { ··· 242 242 }) 243 243 } 244 244 245 245 - fn print_authenticate_warning<S: AsRef<str>>( 245 245 + async fn print_authenticate_warning<S: AsRef<str>>( 246 246 arguments: &CommandArguments<S>, 247 247 ) -> Result<(), HiveLibError> { 248 248 if !arguments.is_elevated() { 249 249 return Ok(()); 250 250 } 251 251 252 252 + let target_display = if let Some(ref target) = arguments.target { 253 253 + let target = target.0.read().await; 254 254 + 255 255 + format!( 256 256 + "{}@{}:{}", 257 257 + target.user, 258 258 + target.get_preferred_host()?, 259 259 + target.port 260 260 + ) 261 261 + } else { 262 262 + "localhost (!)".to_string() 263 263 + }; 264 264 + 252 265 let _ = STATUS.lock().write_above_status( 253 266 &format!( 254 254 - "{} | Authenticate for \"sudo {}\":\n", 255 255 - arguments 256 256 - .target 257 257 - .map_or(Ok("localhost (!)".to_string()), |target| Ok(format!( 258 258 - "{}@{}:{}", 259 259 - target.user, 260 260 - target.get_preferred_host()?, 261 261 - target.port 262 262 - )))?, 267 267 + "{target_display} | Authenticate for \"sudo {}\":\n", 263 268 arguments.command_string.as_ref() 264 269 ) 265 270 .into_bytes(), ··· 306 311 Ok(()) 307 312 } 308 313 309 309 - fn build_command<S: AsRef<str>>( 310 310 - arguments: &CommandArguments<'_, S>, 314 314 + async fn build_command<S: AsRef<str>>( 315 315 + arguments: &CommandArguments<S>, 311 316 command_string: &String, 312 317 ) -> Result<CommandBuilder, HiveLibError> { 313 313 - let mut command = if let Some(target) = arguments.target { 314 314 - let mut command = create_int_ssh_command(target, arguments.modifiers)?; 318 318 + let mut command = if let Some(ref target) = arguments.target { 319 319 + let mut command = create_int_ssh_command(target, arguments.modifiers).await?; 315 320 316 321 // force ssh to use our pseudo terminal 317 322 command.arg("-tt"); ··· 428 433 } 429 434 } 430 435 431 431 - fn create_int_ssh_command( 432 432 - target: &Target, 436 436 + async fn create_int_ssh_command( 437 437 + target: &SharedTarget, 433 438 modifiers: SubCommandModifiers, 434 439 ) -> Result<portable_pty::CommandBuilder, HiveLibError> { 440 440 + let target = target.0.read().await; 435 441 let mut command = portable_pty::CommandBuilder::new("ssh"); 436 436 - command.args(target.create_ssh_args(modifiers, false)?); 442 442 + command.args(target.create_ssh_args(modifiers)?); 437 443 command.arg(target.get_preferred_host()?.to_string()); 438 444 Ok(command) 439 445 }

+167

crates/core/src/hive/executor.rs

reviewed

··· 1 1 + use crate::hive::node::Step; 2 2 + use std::{assert_matches::debug_assert_matches, sync::Arc}; 3 3 + 4 4 + use tracing::{Instrument, Span, debug, error, event, instrument}; 5 5 + 6 6 + use crate::{ 7 7 + EvalGoal, SubCommandModifiers, 8 8 + commands::common::evaluate_hive_attribute, 9 9 + errors::HiveLibError, 10 10 + hive::{ 11 11 + HiveLocation, 12 12 + node::{Context, Derivation, ExecuteStep, Name}, 13 13 + plan::NodePlan, 14 14 + }, 15 15 + status::STATUS, 16 16 + }; 17 17 + 18 18 + /// returns Err if the application should shut down. 19 19 + fn app_shutdown_guard(context: &Context) -> Result<(), HiveLibError> { 20 20 + if context 21 21 + .should_quit 22 22 + .load(std::sync::atomic::Ordering::Relaxed) 23 23 + { 24 24 + return Err(HiveLibError::Sigint); 25 25 + } 26 26 + 27 27 + Ok(()) 28 28 + } 29 29 + 30 30 + /// Task that evaluates the node. 31 31 + #[instrument(skip_all, name = "eval")] 32 32 + async fn evaluate_task( 33 33 + tx: tokio::sync::oneshot::Sender<Result<Derivation, HiveLibError>>, 34 34 + hive_location: Arc<HiveLocation>, 35 35 + name: Name, 36 36 + modifiers: SubCommandModifiers, 37 37 + ) { 38 38 + let output = evaluate_hive_attribute(&hive_location, &EvalGoal::GetTopLevel(&name), modifiers) 39 39 + .await 40 40 + .and_then(|output| { 41 41 + serde_json::from_str(&output).map_err(|e| { 42 42 + HiveLibError::HiveInitialisationError( 43 43 + crate::errors::HiveInitialisationError::ParseEvaluateError(e), 44 44 + ) 45 45 + }) 46 46 + }); 47 47 + 48 48 + debug!(output = ?output, done = true); 49 49 + 50 50 + let _ = tx.send(output); 51 51 + } 52 52 + 53 53 + /// Iterates and executes the steps in the plan. 54 54 + /// Performs some optimisations such as greedily executing evaluation before 55 55 + /// other steps independent of evaluation's result. 56 56 + #[instrument(skip_all, fields(node = %plan.context.name))] 57 57 + pub async fn execute(mut plan: NodePlan) -> Result<(), HiveLibError> { 58 58 + app_shutdown_guard(&plan.context)?; 59 59 + 60 60 + let (tx, rx) = tokio::sync::oneshot::channel(); 61 61 + plan.context.state.evaluation_rx = Some(rx); 62 62 + 63 63 + // The name of this span should never be changed without updating 64 64 + // `wire/cli/tracing_setup.rs` 65 65 + debug_assert_matches!(Span::current().metadata().unwrap().name(), "execute"); 66 66 + // This span should always have a `node` field by the same file 67 67 + debug_assert!( 68 68 + Span::current() 69 69 + .metadata() 70 70 + .unwrap() 71 71 + .fields() 72 72 + .field("node") 73 73 + .is_some() 74 74 + ); 75 75 + 76 76 + if plan.greedy_evaluate { 77 77 + tokio::spawn( 78 78 + evaluate_task( 79 79 + tx, 80 80 + plan.context.hive_location.clone(), 81 81 + plan.context.name.clone(), 82 82 + plan.context.modifiers, 83 83 + ) 84 84 + .in_current_span(), 85 85 + ); 86 86 + } 87 87 + 88 88 + let length = plan.steps.len(); 89 89 + 90 90 + for (position, step) in plan.steps.iter().enumerate() { 91 91 + app_shutdown_guard(&plan.context)?; 92 92 + 93 93 + event!( 94 94 + tracing::Level::INFO, 95 95 + step = step.to_string(), 96 96 + progress = format!("{}/{length}", position + 1) 97 97 + ); 98 98 + 99 99 + STATUS 100 100 + .lock() 101 101 + .set_node_step(&plan.context.name, step.to_string()); 102 102 + 103 103 + if let Err(err) = step.execute(&mut plan.context).await.inspect_err(|_| { 104 104 + error!("Failed to execute `{step}`"); 105 105 + }) { 106 106 + if matches!(step, Step::Ping(..)) && plan.ignore_failed_ping { 107 107 + return Ok(()); 108 108 + } 109 109 + 110 110 + STATUS.lock().mark_node_failed(&plan.context.name); 111 111 + 112 112 + return Err(err); 113 113 + } 114 114 + } 115 115 + 116 116 + STATUS.lock().mark_node_succeeded(&plan.context.name); 117 117 + 118 118 + Ok(()) 119 119 + } 120 120 + 121 121 + #[cfg(test)] 122 122 + mod tests { 123 123 + use crate::{ 124 124 + SubCommandModifiers, 125 125 + errors::HiveLibError, 126 126 + function_name, get_test_path, 127 127 + hive::{ 128 128 + executor::execute, 129 129 + node::{ApplyGoal, HandleUnreachable, Name, Node, SwitchToConfigurationGoal}, 130 130 + plan::{ApplyGoalArgs, Goal, plan_for_node}, 131 131 + }, 132 132 + location, 133 133 + }; 134 134 + use std::{assert_matches::assert_matches, path::PathBuf}; 135 135 + use std::{ 136 136 + env, 137 137 + sync::{Arc, atomic::AtomicBool}, 138 138 + }; 139 139 + 140 140 + #[tokio::test] 141 141 + async fn plan_executor_quits_sigint() { 142 142 + let location = location!(get_test_path!()); 143 143 + let node = Node::default(); 144 144 + let name = &Name(function_name!().into()); 145 145 + let should_quit = Arc::new(AtomicBool::new(true)); 146 146 + let plan = plan_for_node( 147 147 + &node.clone(), 148 148 + name.clone(), 149 149 + &Goal::Apply(ApplyGoalArgs { 150 150 + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), 151 151 + should_apply_locally: true, 152 152 + no_keys: true, 153 153 + substitute_on_destination: true, 154 154 + reboot: false, 155 155 + host_platform: "x86_64-linux".into(), 156 156 + handle_unreachable: HandleUnreachable::default(), 157 157 + }), 158 158 + location.clone().into(), 159 159 + &SubCommandModifiers::default(), 160 160 + should_quit.clone(), 161 161 + ); 162 162 + 163 163 + let status = execute(plan).await; 164 164 + 165 165 + assert_matches!(status, Err(HiveLibError::Sigint)); 166 166 + } 167 167 + }

+18 -15

crates/core/src/hive/mod.rs

reviewed

··· 22 22 use crate::commands::{CommandArguments, Either, WireCommandChip, run_command}; 23 23 use crate::errors::{HiveInitialisationError, HiveLocationError}; 24 24 use crate::{EvalGoal, HiveLibError, SubCommandModifiers}; 25 25 + pub mod executor; 25 26 pub mod node; 27 27 + pub mod plan; 26 28 pub mod steps; 27 29 28 30 #[derive(Serialize, Deserialize, Debug, PartialEq)] ··· 180 182 } 181 183 } 182 184 183 183 - #[derive(Debug, PartialEq, Eq, Deserialize)] 185 185 + #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] 184 186 pub struct FlakePrefetch { 185 187 pub(crate) hash: String, 186 188 #[serde(rename = "storePath")] 187 189 pub(crate) store_path: String, 188 190 } 189 191 190 190 - #[derive(Debug, PartialEq, Eq)] 192 192 + #[derive(Debug, Clone, PartialEq, Eq)] 191 193 pub enum HiveLocation { 192 194 HiveNix(PathBuf), 193 195 Flake { ··· 288 290 289 291 #[cfg(test)] 290 292 mod tests { 291 291 - use im::vector; 292 292 - 293 293 use crate::{ 294 294 errors::CommandError, 295 295 get_test_path, ··· 346 346 347 347 let node = Node { 348 348 target: node::Target::from_host("name"), 349 349 - keys: vector![Key { 350 350 - name: "different-than-a".into(), 351 351 - dest_dir: "/run/keys/".into(), 352 352 - path: "/run/keys/different-than-a".into(), 353 353 - group: "root".into(), 354 354 - user: "root".into(), 355 355 - permissions: "0600".into(), 356 356 - source: Source::String("hi".into()), 357 357 - upload_at: UploadKeyAt::PreActivation, 358 358 - environment: im::HashMap::new() 359 359 - }], 349 349 + keys: vec![ 350 350 + Key { 351 351 + name: "different-than-a".into(), 352 352 + dest_dir: "/run/keys/".into(), 353 353 + path: "/run/keys/different-than-a".into(), 354 354 + group: "root".into(), 355 355 + user: "root".into(), 356 356 + permissions: "0600".into(), 357 357 + source: Source::String("hi".into()), 358 358 + upload_at: UploadKeyAt::PreActivation, 359 359 + environment: im::HashMap::new(), 360 360 + } 361 361 + .into(), 362 362 + ], 360 363 build_remotely: true, 361 364 ..Default::default() 362 365 };

+115 -585

crates/core/src/hive/node.rs

reviewed

··· 5 5 use enum_dispatch::enum_dispatch; 6 6 use gethostname::gethostname; 7 7 use serde::{Deserialize, Serialize}; 8 8 - use std::assert_matches::debug_assert_matches; 9 8 use std::fmt::Display; 10 9 use std::sync::Arc; 11 10 use std::sync::atomic::AtomicBool; 12 12 - use tokio::sync::oneshot; 13 13 - use tracing::{Instrument, Level, Span, debug, error, event, instrument, trace}; 11 11 + use tokio::sync::{RwLock, oneshot}; 12 12 + use tracing::instrument; 14 13 15 14 use crate::commands::builder::CommandStringBuilder; 16 16 - use crate::commands::common::evaluate_hive_attribute; 17 15 use crate::commands::{CommandArguments, WireCommandChip, run_command}; 18 16 use crate::errors::NetworkError; 19 17 use crate::hive::HiveLocation; 20 18 use crate::hive::steps::build::Build; 21 21 - use crate::hive::steps::cleanup::CleanUp; 22 19 use crate::hive::steps::evaluate::Evaluate; 23 23 - use crate::hive::steps::keys::{Key, Keys, PushKeyAgent, UploadKeyAt}; 20 20 + use crate::hive::steps::keys::{Key, Keys, PushKeyAgent}; 24 21 use crate::hive::steps::ping::Ping; 25 22 use crate::hive::steps::push::{PushBuildOutput, PushEvaluatedOutput}; 26 26 - use crate::status::STATUS; 27 27 - use crate::{EvalGoal, StrictHostKeyChecking, SubCommandModifiers}; 23 23 + use crate::{StrictHostKeyChecking, SubCommandModifiers}; 28 24 29 25 use super::HiveLibError; 30 26 use super::steps::activate::SwitchToConfiguration; ··· 44 40 current_host: usize, 45 41 } 46 42 43 43 + #[derive(Clone, Debug)] 44 44 + pub struct SharedTarget(pub Arc<RwLock<Target>>); 45 45 + 46 46 + // Hack specifically for testing if two steps that have the same shared target 47 47 + // are equal 48 48 + #[cfg(test)] 49 49 + impl PartialEq for SharedTarget { 50 50 + fn eq(&self, other: &Self) -> bool { 51 51 + let self_guard = self 52 52 + .0 53 53 + .try_read() 54 54 + .expect("failed to target read in test context"); 55 55 + let other_guard = other 56 56 + .0 57 57 + .try_read() 58 58 + .expect("failed to target read in test context"); 59 59 + 60 60 + *self_guard == *other_guard 61 61 + } 62 62 + } 63 63 + 47 64 impl Target { 48 65 #[instrument(ret(level = tracing::Level::DEBUG), skip_all)] 49 66 pub fn create_ssh_opts(&self, modifiers: SubCommandModifiers) -> Result<String, HiveLibError> { 50 50 - self.create_ssh_args(modifiers, false).map(|x| x.join(" ")) 67 67 + self.create_ssh_args(modifiers).map(|x| x.join(" ")) 51 68 } 52 69 53 70 #[instrument(ret(level = tracing::Level::DEBUG))] 54 71 pub fn create_ssh_args( 55 72 &self, 56 73 modifiers: SubCommandModifiers, 57 57 - non_interactive_forced: bool, 58 74 ) -> Result<Vec<String>, HiveLibError> { 59 75 let mut vector = vec![ 60 76 "-l".to_string(), ··· 81 97 82 98 Ok(vector) 83 99 } 100 100 + 101 101 + /// Tests the connection to a node 102 102 + pub async fn ping(&self, modifiers: SubCommandModifiers) -> Result<(), HiveLibError> { 103 103 + let host = self.get_preferred_host()?; 104 104 + 105 105 + let mut command_string = CommandStringBuilder::new("ssh"); 106 106 + command_string.arg(format!("{}@{host}", self.user)); 107 107 + command_string.arg(self.create_ssh_opts(modifiers)?); 108 108 + command_string.arg("exit"); 109 109 + 110 110 + let output = run_command( 111 111 + &CommandArguments::new(command_string, modifiers) 112 112 + .log_stdout() 113 113 + .mode(crate::commands::ChildOutputMode::Interactive), 114 114 + ) 115 115 + .await?; 116 116 + 117 117 + output.wait_till_success().await.map_err(|source| { 118 118 + HiveLibError::NetworkError(NetworkError::HostUnreachable { 119 119 + host: host.to_string(), 120 120 + source, 121 121 + }) 122 122 + })?; 123 123 + 124 124 + Ok(()) 125 125 + } 84 126 } 85 127 86 128 #[cfg(test)] ··· 91 133 user: "root".into(), 92 134 port: 22, 93 135 current_host: 0, 94 94 - } 95 95 - } 96 96 - } 97 97 - 98 98 - #[cfg(test)] 99 99 - impl<'a> Context<'a> { 100 100 - fn create_test_context( 101 101 - hive_location: HiveLocation, 102 102 - name: &'a Name, 103 103 - node: &'a mut Node, 104 104 - ) -> Self { 105 105 - Context { 106 106 - name, 107 107 - node, 108 108 - hive_location: Arc::new(hive_location), 109 109 - modifiers: SubCommandModifiers::default(), 110 110 - objective: Objective::Apply(ApplyObjective { 111 111 - goal: Goal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), 112 112 - no_keys: false, 113 113 - reboot: false, 114 114 - should_apply_locally: false, 115 115 - substitute_on_destination: false, 116 116 - handle_unreachable: HandleUnreachable::default(), 117 117 - }), 118 118 - state: StepState::default(), 119 119 - should_quit: Arc::new(AtomicBool::new(false)), 120 136 } 121 137 } 122 138 } ··· 171 187 pub tags: im::HashSet<String>, 172 188 173 189 #[serde(rename(deserialize = "_keys", serialize = "keys"))] 174 174 - pub keys: im::Vector<Key>, 190 190 + pub keys: Vec<Arc<Key>>, 175 191 176 192 #[serde(rename(deserialize = "_hostPlatform", serialize = "host_platform"))] 177 193 pub host_platform: Arc<str>, ··· 180 196 deserialize = "privilegeEscalationCommand", 181 197 serialize = "privilege_escalation_command" 182 198 ))] 183 183 - pub privilege_escalation_command: im::Vector<Arc<str>>, 199 199 + pub privilege_escalation_command: Arc<Vec<Arc<str>>>, 184 200 } 185 201 186 202 #[cfg(test)] ··· 188 204 fn default() -> Self { 189 205 Node { 190 206 target: Target::default(), 191 191 - keys: im::Vector::new(), 207 207 + keys: Vec::new(), 192 208 tags: im::HashSet::new(), 193 209 privilege_escalation_command: vec!["sudo".into(), "--".into()].into(), 194 210 allow_local_deployment: true, ··· 207 223 ..Default::default() 208 224 } 209 225 } 210 210 - 211 211 - /// Tests the connection to a node 212 212 - pub async fn ping(&self, modifiers: SubCommandModifiers) -> Result<(), HiveLibError> { 213 213 - let host = self.target.get_preferred_host()?; 214 214 - 215 215 - let mut command_string = CommandStringBuilder::new("ssh"); 216 216 - command_string.arg(format!("{}@{host}", self.target.user)); 217 217 - command_string.arg(self.target.create_ssh_opts(modifiers)?); 218 218 - command_string.arg("exit"); 219 219 - 220 220 - let output = run_command( 221 221 - &CommandArguments::new(command_string, modifiers) 222 222 - .log_stdout() 223 223 - .mode(crate::commands::ChildOutputMode::Interactive), 224 224 - ) 225 225 - .await?; 226 226 - 227 227 - output.wait_till_success().await.map_err(|source| { 228 228 - HiveLibError::NetworkError(NetworkError::HostUnreachable { 229 229 - host: host.to_string(), 230 230 - source, 231 231 - }) 232 232 - })?; 233 233 - 234 234 - Ok(()) 235 235 - } 236 226 } 237 227 238 228 #[must_use] ··· 255 245 } 256 246 } 257 247 258 258 - #[derive(derive_more::Display, Debug, Clone, Copy)] 248 248 + #[derive(derive_more::Display, Debug, Clone, Copy, PartialEq, Eq)] 259 249 pub enum SwitchToConfigurationGoal { 260 250 Switch, 261 251 Boot, ··· 263 253 DryActivate, 264 254 } 265 255 266 266 - #[derive(derive_more::Display, Clone, Copy)] 267 267 - pub enum Goal { 256 256 + #[derive(derive_more::Display, Debug, Clone, Copy)] 257 257 + pub enum ApplyGoal { 268 258 SwitchToConfiguration(SwitchToConfigurationGoal), 269 259 Build, 270 260 Push, 271 261 Keys, 272 262 } 273 263 274 274 - // TODO: Get rid of this allow and resolve it 275 275 - #[allow(clippy::struct_excessive_bools)] 276 276 - #[derive(Clone, Copy)] 277 277 - pub struct ApplyObjective { 278 278 - pub goal: Goal, 279 279 - pub no_keys: bool, 280 280 - pub reboot: bool, 281 281 - pub should_apply_locally: bool, 282 282 - pub substitute_on_destination: bool, 283 283 - pub handle_unreachable: HandleUnreachable, 284 284 - } 285 285 - 286 286 - #[derive(Clone, Copy)] 287 287 - pub enum Objective { 288 288 - Apply(ApplyObjective), 289 289 - BuildLocally, 290 290 - } 291 291 - 292 264 #[enum_dispatch] 293 265 pub(crate) trait ExecuteStep: Send + Sync + Display + std::fmt::Debug { 294 294 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError>; 295 295 - 296 296 - fn should_execute(&self, context: &Context) -> bool; 266 266 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError>; 297 267 } 298 268 299 269 // may include other options such as FailAll in the future ··· 313 283 pub key_agent_directory: Option<String>, 314 284 } 315 285 316 316 - pub struct Context<'a> { 317 317 - pub name: &'a Name, 318 318 - pub node: &'a mut Node, 286 286 + pub struct Context { 319 287 pub hive_location: Arc<HiveLocation>, 320 288 pub modifiers: SubCommandModifiers, 321 289 pub state: StepState, 322 290 pub should_quit: Arc<AtomicBool>, 323 323 - pub objective: Objective, 291 291 + pub name: Name, 324 292 } 325 293 326 294 #[enum_dispatch(ExecuteStep)] 327 327 - #[derive(Debug, PartialEq)] 328 328 - enum Step { 295 295 + #[derive(Debug)] 296 296 + #[cfg_attr(test, derive(PartialEq))] 297 297 + pub enum Step { 329 298 Ping, 330 299 PushKeyAgent, 331 300 Keys, ··· 334 303 Build, 335 304 PushBuildOutput, 336 305 SwitchToConfiguration, 337 337 - CleanUp, 338 306 } 339 307 340 308 impl Display for Step { ··· 348 316 Self::Build(step) => step.fmt(f), 349 317 Self::PushBuildOutput(step) => step.fmt(f), 350 318 Self::SwitchToConfiguration(step) => step.fmt(f), 351 351 - Self::CleanUp(step) => step.fmt(f), 352 319 } 353 320 } 354 321 } 355 322 356 356 - pub struct GoalExecutor<'a> { 357 357 - steps: Vec<Step>, 358 358 - context: Context<'a>, 359 359 - } 360 360 - 361 361 - /// returns Err if the application should shut down. 362 362 - fn app_shutdown_guard(context: &Context) -> Result<(), HiveLibError> { 363 363 - if context 364 364 - .should_quit 365 365 - .load(std::sync::atomic::Ordering::Relaxed) 366 366 - { 367 367 - return Err(HiveLibError::Sigint); 368 368 - } 369 369 - 370 370 - Ok(()) 371 371 - } 372 372 - 373 373 - impl<'a> GoalExecutor<'a> { 374 374 - #[must_use] 375 375 - pub fn new(context: Context<'a>) -> Self { 376 376 - Self { 377 377 - steps: vec![ 378 378 - Step::Ping(Ping), 379 379 - Step::PushKeyAgent(PushKeyAgent), 380 380 - Step::Keys(Keys { 381 381 - filter: UploadKeyAt::NoFilter, 382 382 - }), 383 383 - Step::Keys(Keys { 384 384 - filter: UploadKeyAt::PreActivation, 385 385 - }), 386 386 - Step::Evaluate(super::steps::evaluate::Evaluate), 387 387 - Step::PushEvaluatedOutput(super::steps::push::PushEvaluatedOutput), 388 388 - Step::Build(super::steps::build::Build), 389 389 - Step::PushBuildOutput(super::steps::push::PushBuildOutput), 390 390 - Step::SwitchToConfiguration(SwitchToConfiguration), 391 391 - Step::Keys(Keys { 392 392 - filter: UploadKeyAt::PostActivation, 393 393 - }), 394 394 - ], 395 395 - context, 396 396 - } 397 397 - } 398 398 - 399 399 - #[instrument(skip_all, name = "eval")] 400 400 - async fn evaluate_task( 401 401 - tx: oneshot::Sender<Result<Derivation, HiveLibError>>, 402 402 - hive_location: Arc<HiveLocation>, 403 403 - name: Name, 404 404 - modifiers: SubCommandModifiers, 405 405 - ) { 406 406 - let output = 407 407 - evaluate_hive_attribute(&hive_location, &EvalGoal::GetTopLevel(&name), modifiers) 408 408 - .await 409 409 - .map(|output| { 410 410 - serde_json::from_str::<Derivation>(&output).expect("failed to parse derivation") 411 411 - }); 412 412 - 413 413 - debug!(output = ?output, done = true); 414 414 - 415 415 - let _ = tx.send(output); 416 416 - } 417 417 - 418 418 - #[instrument(skip_all, fields(node = %self.context.name))] 419 419 - pub async fn execute(mut self) -> Result<(), HiveLibError> { 420 420 - app_shutdown_guard(&self.context)?; 421 421 - 422 422 - let (tx, rx) = oneshot::channel(); 423 423 - self.context.state.evaluation_rx = Some(rx); 424 424 - 425 425 - // The name of this span should never be changed without updating 426 426 - // `wire/cli/tracing_setup.rs` 427 427 - debug_assert_matches!(Span::current().metadata().unwrap().name(), "execute"); 428 428 - // This span should always have a `node` field by the same file 429 429 - debug_assert!( 430 430 - Span::current() 431 431 - .metadata() 432 432 - .unwrap() 433 433 - .fields() 434 434 - .field("node") 435 435 - .is_some() 436 436 - ); 437 437 - 438 438 - let spawn_evaluator = match self.context.objective { 439 439 - Objective::Apply(apply_objective) => !matches!(apply_objective.goal, Goal::Keys), 440 440 - Objective::BuildLocally => true, 441 441 - }; 442 442 - 443 443 - if spawn_evaluator { 444 444 - tokio::spawn( 445 445 - GoalExecutor::evaluate_task( 446 446 - tx, 447 447 - self.context.hive_location.clone(), 448 448 - self.context.name.clone(), 449 449 - self.context.modifiers, 450 450 - ) 451 451 - .in_current_span(), 452 452 - ); 453 453 - } 454 454 - 455 455 - let steps = self 456 456 - .steps 457 457 - .iter() 458 458 - .filter(|step| step.should_execute(&self.context)) 459 459 - .inspect(|step| { 460 460 - trace!("Will execute step `{step}` for {}", self.context.name); 461 461 - }) 462 462 - .collect::<Vec<_>>(); 463 463 - let length = steps.len(); 464 464 - 465 465 - for (position, step) in steps.iter().enumerate() { 466 466 - app_shutdown_guard(&self.context)?; 467 467 - 468 468 - event!( 469 469 - Level::INFO, 470 470 - step = step.to_string(), 471 471 - progress = format!("{}/{length}", position + 1) 472 472 - ); 473 473 - 474 474 - STATUS 475 475 - .lock() 476 476 - .set_node_step(self.context.name, step.to_string()); 477 477 - 478 478 - if let Err(err) = step.execute(&mut self.context).await.inspect_err(|_| { 479 479 - error!("Failed to execute `{step}`"); 480 480 - }) { 481 481 - // discard error from cleanup 482 482 - let _ = CleanUp.execute(&mut self.context).await; 483 483 - 484 484 - if let Objective::Apply(apply_objective) = self.context.objective 485 485 - && matches!(step, Step::Ping(..)) 486 486 - && matches!( 487 487 - apply_objective.handle_unreachable, 488 488 - HandleUnreachable::Ignore, 489 489 - ) 490 490 - { 491 491 - return Ok(()); 492 492 - } 493 493 - 494 494 - STATUS.lock().mark_node_failed(self.context.name); 495 495 - 496 496 - return Err(err); 497 497 - } 498 498 - } 499 499 - 500 500 - STATUS.lock().mark_node_succeeded(self.context.name); 501 501 - 502 502 - Ok(()) 503 503 - } 504 504 - } 505 505 - 506 323 #[cfg(test)] 507 324 mod tests { 508 325 use rand::distr::Alphabetic; 509 326 510 327 use super::*; 511 511 - use crate::{ 512 512 - function_name, get_test_path, 513 513 - hive::{Hive, get_hive_location}, 514 514 - location, 515 515 - }; 516 516 - use std::{assert_matches::assert_matches, path::PathBuf}; 517 517 - use std::{collections::HashMap, env}; 518 518 - 519 519 - fn get_steps(goal_executor: GoalExecutor) -> std::vec::Vec<Step> { 520 520 - goal_executor 521 521 - .steps 522 522 - .into_iter() 523 523 - .filter(|step| step.should_execute(&goal_executor.context)) 524 524 - .collect::<Vec<_>>() 525 525 - } 526 526 - 527 527 - #[tokio::test] 528 528 - #[cfg_attr(feature = "no_web_tests", ignore)] 529 529 - async fn default_values_match() { 530 530 - let mut path = get_test_path!(); 531 531 - 532 532 - let location = 533 533 - get_hive_location(path.display().to_string(), SubCommandModifiers::default()) 534 534 - .await 535 535 - .unwrap(); 536 536 - let hive = Hive::new_from_path(&location, None, SubCommandModifiers::default()) 537 537 - .await 538 538 - .unwrap(); 539 539 - 540 540 - let node = Node::default(); 541 541 - 542 542 - let mut nodes = HashMap::new(); 543 543 - nodes.insert(Name("NAME".into()), node); 544 544 - 545 545 - path.push("hive.nix"); 546 546 - 547 547 - assert_eq!( 548 548 - hive, 549 549 - Hive { 550 550 - nodes, 551 551 - schema: Hive::SCHEMA_VERSION 552 552 - } 553 553 - ); 554 554 - } 555 555 - 556 556 - #[tokio::test] 557 557 - async fn order_build_locally() { 558 558 - let location = location!(get_test_path!()); 559 559 - let mut node = Node { 560 560 - build_remotely: false, 561 561 - ..Default::default() 562 562 - }; 563 563 - let name = &Name(function_name!().into()); 564 564 - let executor = GoalExecutor::new(Context::create_test_context(location, name, &mut node)); 565 565 - let steps = get_steps(executor); 566 566 - 567 567 - assert_eq!( 568 568 - steps, 569 569 - vec![ 570 570 - Ping.into(), 571 571 - PushKeyAgent.into(), 572 572 - Keys { 573 573 - filter: UploadKeyAt::PreActivation 574 574 - } 575 575 - .into(), 576 576 - crate::hive::steps::evaluate::Evaluate.into(), 577 577 - crate::hive::steps::build::Build.into(), 578 578 - crate::hive::steps::push::PushBuildOutput.into(), 579 579 - SwitchToConfiguration.into(), 580 580 - Keys { 581 581 - filter: UploadKeyAt::PostActivation 582 582 - } 583 583 - .into(), 584 584 - ] 585 585 - ); 586 586 - } 587 587 - 588 588 - #[tokio::test] 589 589 - async fn order_keys_only() { 590 590 - let location = location!(get_test_path!()); 591 591 - let mut node = Node::default(); 592 592 - let name = &Name(function_name!().into()); 593 593 - let mut context = Context::create_test_context(location, name, &mut node); 594 594 - 595 595 - let Objective::Apply(ref mut apply_objective) = context.objective else { 596 596 - unreachable!() 597 597 - }; 598 598 - 599 599 - apply_objective.goal = Goal::Keys; 600 600 - 601 601 - let executor = GoalExecutor::new(context); 602 602 - let steps = get_steps(executor); 603 603 - 604 604 - assert_eq!( 605 605 - steps, 606 606 - vec![ 607 607 - Ping.into(), 608 608 - PushKeyAgent.into(), 609 609 - Keys { 610 610 - filter: UploadKeyAt::NoFilter 611 611 - } 612 612 - .into(), 613 613 - ] 614 614 - ); 615 615 - } 616 616 - 617 617 - #[tokio::test] 618 618 - async fn order_build() { 619 619 - let location = location!(get_test_path!()); 620 620 - let mut node = Node::default(); 621 621 - let name = &Name(function_name!().into()); 622 622 - let mut context = Context::create_test_context(location, name, &mut node); 623 623 - 624 624 - let Objective::Apply(ref mut apply_objective) = context.objective else { 625 625 - unreachable!() 626 626 - }; 627 627 - apply_objective.goal = Goal::Build; 628 628 - 629 629 - let executor = GoalExecutor::new(context); 630 630 - let steps = get_steps(executor); 631 631 - 632 632 - assert_eq!( 633 633 - steps, 634 634 - vec![ 635 635 - Ping.into(), 636 636 - crate::hive::steps::evaluate::Evaluate.into(), 637 637 - crate::hive::steps::build::Build.into(), 638 638 - crate::hive::steps::push::PushBuildOutput.into(), 639 639 - ] 640 640 - ); 641 641 - } 642 642 - 643 643 - #[tokio::test] 644 644 - async fn order_push_only() { 645 645 - let location = location!(get_test_path!()); 646 646 - let mut node = Node::default(); 647 647 - let name = &Name(function_name!().into()); 648 648 - let mut context = Context::create_test_context(location, name, &mut node); 649 649 - 650 650 - let Objective::Apply(ref mut apply_objective) = context.objective else { 651 651 - unreachable!() 652 652 - }; 653 653 - apply_objective.goal = Goal::Push; 654 654 - 655 655 - let executor = GoalExecutor::new(context); 656 656 - let steps = get_steps(executor); 657 657 - 658 658 - assert_eq!( 659 659 - steps, 660 660 - vec![ 661 661 - Ping.into(), 662 662 - crate::hive::steps::evaluate::Evaluate.into(), 663 663 - crate::hive::steps::push::PushEvaluatedOutput.into(), 664 664 - ] 665 665 - ); 666 666 - } 667 667 - 668 668 - #[tokio::test] 669 669 - async fn order_remote_build() { 670 670 - let location = location!(get_test_path!()); 671 671 - let mut node = Node { 672 672 - build_remotely: true, 673 673 - ..Default::default() 674 674 - }; 675 675 - 676 676 - let name = &Name(function_name!().into()); 677 677 - let executor = GoalExecutor::new(Context::create_test_context(location, name, &mut node)); 678 678 - let steps = get_steps(executor); 679 679 - 680 680 - assert_eq!( 681 681 - steps, 682 682 - vec![ 683 683 - Ping.into(), 684 684 - PushKeyAgent.into(), 685 685 - Keys { 686 686 - filter: UploadKeyAt::PreActivation 687 687 - } 688 688 - .into(), 689 689 - crate::hive::steps::evaluate::Evaluate.into(), 690 690 - crate::hive::steps::push::PushEvaluatedOutput.into(), 691 691 - crate::hive::steps::build::Build.into(), 692 692 - SwitchToConfiguration.into(), 693 693 - Keys { 694 694 - filter: UploadKeyAt::PostActivation 695 695 - } 696 696 - .into(), 697 697 - ] 698 698 - ); 699 699 - } 700 700 - 701 701 - #[tokio::test] 702 702 - async fn order_nokeys() { 703 703 - let location = location!(get_test_path!()); 704 704 - let mut node = Node::default(); 705 705 - 706 706 - let name = &Name(function_name!().into()); 707 707 - let mut context = Context::create_test_context(location, name, &mut node); 708 708 - 709 709 - let Objective::Apply(ref mut apply_objective) = context.objective else { 710 710 - unreachable!() 711 711 - }; 712 712 - apply_objective.no_keys = true; 713 713 - 714 714 - let executor = GoalExecutor::new(context); 715 715 - let steps = get_steps(executor); 716 716 - 717 717 - assert_eq!( 718 718 - steps, 719 719 - vec![ 720 720 - Ping.into(), 721 721 - crate::hive::steps::evaluate::Evaluate.into(), 722 722 - crate::hive::steps::build::Build.into(), 723 723 - crate::hive::steps::push::PushBuildOutput.into(), 724 724 - SwitchToConfiguration.into(), 725 725 - ] 726 726 - ); 727 727 - } 728 728 - 729 729 - #[tokio::test] 730 730 - async fn order_should_apply_locally() { 731 731 - let location = location!(get_test_path!()); 732 732 - let mut node = Node::default(); 733 733 - 734 734 - let name = &Name(function_name!().into()); 735 735 - let mut context = Context::create_test_context(location, name, &mut node); 736 736 - 737 737 - let Objective::Apply(ref mut apply_objective) = context.objective else { 738 738 - unreachable!() 739 739 - }; 740 740 - apply_objective.no_keys = true; 741 741 - apply_objective.should_apply_locally = true; 742 742 - 743 743 - let executor = GoalExecutor::new(context); 744 744 - let steps = get_steps(executor); 745 745 - 746 746 - assert_eq!( 747 747 - steps, 748 748 - vec![ 749 749 - crate::hive::steps::evaluate::Evaluate.into(), 750 750 - crate::hive::steps::build::Build.into(), 751 751 - SwitchToConfiguration.into(), 752 752 - ] 753 753 - ); 754 754 - } 755 755 - 756 756 - #[tokio::test] 757 757 - async fn order_build_only() { 758 758 - let location = location!(get_test_path!()); 759 759 - let mut node = Node::default(); 760 760 - 761 761 - let name = &Name(function_name!().into()); 762 762 - let mut context = Context::create_test_context(location, name, &mut node); 763 763 - 764 764 - context.objective = Objective::BuildLocally; 765 765 - 766 766 - let executor = GoalExecutor::new(context); 767 767 - let steps = get_steps(executor); 768 768 - 769 769 - assert_eq!( 770 770 - steps, 771 771 - vec![ 772 772 - crate::hive::steps::evaluate::Evaluate.into(), 773 773 - crate::hive::steps::build::Build.into() 774 774 - ] 775 775 - ); 776 776 - } 777 777 - 778 778 - #[test] 779 779 - fn target_fails_increments() { 780 780 - let mut target = Target::from_host("localhost"); 781 781 - 782 782 - assert_eq!(target.current_host, 0); 783 783 - 784 784 - for i in 0..100 { 785 785 - target.host_failed(); 786 786 - assert_eq!(target.current_host, i + 1); 787 787 - } 788 788 - } 789 789 - 790 790 - #[test] 791 791 - fn get_preferred_host_fails() { 792 792 - let mut target = Target { 793 793 - hosts: vec![ 794 794 - "un.reachable.1".into(), 795 795 - "un.reachable.2".into(), 796 796 - "un.reachable.3".into(), 797 797 - "un.reachable.4".into(), 798 798 - "un.reachable.5".into(), 799 799 - ], 800 800 - ..Default::default() 801 801 - }; 802 802 - 803 803 - assert_ne!( 804 804 - target.get_preferred_host().unwrap().to_string(), 805 805 - "un.reachable.5" 806 806 - ); 807 807 - 808 808 - for i in 1..=5 { 809 809 - assert_eq!( 810 810 - target.get_preferred_host().unwrap().to_string(), 811 811 - format!("un.reachable.{i}") 812 812 - ); 813 813 - target.host_failed(); 814 814 - } 815 815 - 816 816 - for _ in 0..5 { 817 817 - assert_matches!( 818 818 - target.get_preferred_host(), 819 819 - Err(HiveLibError::NetworkError(NetworkError::HostsExhausted)) 820 820 - ); 821 821 - } 822 822 - } 328 328 + use std::{assert_matches::assert_matches, env}; 823 329 824 330 #[test] 825 331 fn test_ssh_opts() { ··· 850 356 "KbdInteractiveAuthentication=no".to_string(), 851 357 ]; 852 358 853 853 - assert_eq!( 854 854 - target.create_ssh_args(subcommand_modifiers, false).unwrap(), 855 855 - args 856 856 - ); 359 359 + assert_eq!(target.create_ssh_args(subcommand_modifiers).unwrap(), args); 857 360 assert_eq!( 858 361 target.create_ssh_opts(subcommand_modifiers).unwrap(), 859 362 args.join(" ") 860 363 ); 861 364 862 365 assert_eq!( 863 863 - target.create_ssh_args(subcommand_modifiers, false).unwrap(), 366 366 + target.create_ssh_args(subcommand_modifiers).unwrap(), 864 367 [ 865 368 "-l".to_string(), 866 369 target.user.to_string(), ··· 876 379 ); 877 380 878 381 assert_eq!( 879 879 - target.create_ssh_args(subcommand_modifiers, true).unwrap(), 382 382 + target.create_ssh_args(subcommand_modifiers).unwrap(), 880 383 [ 881 384 "-l".to_string(), 882 385 target.user.to_string(), ··· 893 396 894 397 // forced non interactive is the same as --non-interactive 895 398 assert_eq!( 896 896 - target.create_ssh_args(subcommand_modifiers, true).unwrap(), 399 399 + target.create_ssh_args(subcommand_modifiers).unwrap(), 897 400 target 898 898 - .create_ssh_args( 899 899 - SubCommandModifiers { 900 900 - non_interactive: true, 901 901 - ..Default::default() 902 902 - }, 903 903 - false 904 904 - ) 401 401 + .create_ssh_args(SubCommandModifiers { 402 402 + non_interactive: true, 403 403 + ..Default::default() 404 404 + }) 905 405 .unwrap() 906 406 ); 907 407 } 908 408 909 909 - #[tokio::test] 910 910 - async fn context_quits_sigint() { 911 911 - let location = location!(get_test_path!()); 912 912 - let mut node = Node::default(); 409 409 + #[test] 410 410 + fn target_fails_increments() { 411 411 + let mut target = Target::from_host("localhost"); 412 412 + 413 413 + assert_eq!(target.current_host, 0); 414 414 + 415 415 + for i in 0..100 { 416 416 + target.host_failed(); 417 417 + assert_eq!(target.current_host, i + 1); 418 418 + } 419 419 + } 420 420 + 421 421 + #[test] 422 422 + fn get_preferred_host_fails() { 423 423 + let mut target = Target { 424 424 + hosts: vec![ 425 425 + "un.reachable.1".into(), 426 426 + "un.reachable.2".into(), 427 427 + "un.reachable.3".into(), 428 428 + "un.reachable.4".into(), 429 429 + "un.reachable.5".into(), 430 430 + ], 431 431 + ..Default::default() 432 432 + }; 433 433 + 434 434 + assert_ne!( 435 435 + target.get_preferred_host().unwrap().to_string(), 436 436 + "un.reachable.5" 437 437 + ); 913 438 914 914 - let name = &Name(function_name!().into()); 915 915 - let context = Context::create_test_context(location, name, &mut node); 916 916 - context 917 917 - .should_quit 918 918 - .store(true, std::sync::atomic::Ordering::Relaxed); 919 919 - let executor = GoalExecutor::new(context); 920 920 - let status = executor.execute().await; 439 439 + for i in 1..=5 { 440 440 + assert_eq!( 441 441 + target.get_preferred_host().unwrap().to_string(), 442 442 + format!("un.reachable.{i}") 443 443 + ); 444 444 + target.host_failed(); 445 445 + } 921 446 922 922 - assert_matches!(status, Err(HiveLibError::Sigint)); 447 447 + for _ in 0..5 { 448 448 + assert_matches!( 449 449 + target.get_preferred_host(), 450 450 + Err(HiveLibError::NetworkError(NetworkError::HostsExhausted)) 451 451 + ); 452 452 + } 923 453 } 924 454 }

+735

crates/core/src/hive/plan.rs

reviewed

··· 1 1 + use std::sync::{Arc, atomic::AtomicBool}; 2 2 + 3 3 + use tokio::sync::RwLock; 4 4 + 5 5 + use crate::{ 6 6 + SubCommandModifiers, 7 7 + hive::{ 8 8 + HiveLocation, 9 9 + node::{ 10 10 + ApplyGoal, Context, HandleUnreachable, Name, Node, SharedTarget, Step, StepState, 11 11 + SwitchToConfigurationGoal, 12 12 + }, 13 13 + steps::{ 14 14 + activate::SwitchToConfiguration, 15 15 + build::Build, 16 16 + evaluate::Evaluate, 17 17 + keys::{Keys, PushKeyAgent, UploadKeyAt}, 18 18 + ping::Ping, 19 19 + push::{PushBuildOutput, PushEvaluatedOutput}, 20 20 + }, 21 21 + }, 22 22 + }; 23 23 + 24 24 + pub struct NodePlan { 25 25 + pub context: Context, 26 26 + pub steps: Vec<Step>, 27 27 + pub greedy_evaluate: bool, 28 28 + pub ignore_failed_ping: bool, 29 29 + } 30 30 + 31 31 + #[allow(clippy::struct_excessive_bools)] 32 32 + pub struct ApplyGoalArgs { 33 33 + pub goal: ApplyGoal, 34 34 + pub should_apply_locally: bool, 35 35 + pub no_keys: bool, 36 36 + pub substitute_on_destination: bool, 37 37 + pub reboot: bool, 38 38 + pub host_platform: Arc<str>, 39 39 + pub handle_unreachable: HandleUnreachable, 40 40 + } 41 41 + 42 42 + pub enum Goal { 43 43 + Apply(ApplyGoalArgs), 44 44 + Build, 45 45 + } 46 46 + 47 47 + fn apply_plan_keys( 48 48 + args: &ApplyGoalArgs, 49 49 + node: &Node, 50 50 + target: &SharedTarget, 51 51 + ) -> (Vec<Step>, Vec<Step>) { 52 52 + let ApplyGoalArgs { 53 53 + goal, 54 54 + substitute_on_destination, 55 55 + should_apply_locally, 56 56 + host_platform, 57 57 + .. 58 58 + } = args; 59 59 + let mut front_steps = Vec::new(); 60 60 + let mut end_steps = Vec::new(); 61 61 + 62 62 + let (pre_keys, post_keys) = match goal { 63 63 + ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) => node 64 64 + .keys 65 65 + .clone() 66 66 + .into_iter() 67 67 + .partition(|x| matches!(x.upload_at, UploadKeyAt::PreActivation)), 68 68 + ApplyGoal::Keys => (node.keys.clone(), Vec::new()), 69 69 + ApplyGoal::Build | ApplyGoal::Push | ApplyGoal::SwitchToConfiguration(_) => { 70 70 + unreachable!("apply_plan_keys called with non-key goal: {:?}", goal) 71 71 + } 72 72 + }; 73 73 + 74 74 + // only push key agent if there are any keys at all 75 75 + if !pre_keys.is_empty() || !post_keys.is_empty() { 76 76 + front_steps.push(Step::PushKeyAgent(PushKeyAgent { 77 77 + substitute_on_destination: *substitute_on_destination, 78 78 + host_platform: host_platform.clone(), 79 79 + target: if *should_apply_locally { 80 80 + None 81 81 + } else { 82 82 + Some(target.clone()) 83 83 + }, 84 84 + })); 85 85 + } 86 86 + 87 87 + if !pre_keys.is_empty() { 88 88 + front_steps.push(Step::Keys(Keys { 89 89 + keys: pre_keys, 90 90 + target: if *should_apply_locally { 91 91 + None 92 92 + } else { 93 93 + Some(target.clone()) 94 94 + }, 95 95 + privilege_escalation_command: node.privilege_escalation_command.clone(), 96 96 + })); 97 97 + } 98 98 + 99 99 + if !post_keys.is_empty() { 100 100 + end_steps.push(Step::Keys(Keys { 101 101 + keys: post_keys, 102 102 + target: if *should_apply_locally { 103 103 + None 104 104 + } else { 105 105 + Some(target.clone()) 106 106 + }, 107 107 + privilege_escalation_command: node.privilege_escalation_command.clone(), 108 108 + })); 109 109 + } 110 110 + 111 111 + (front_steps, end_steps) 112 112 + } 113 113 + 114 114 + fn apply_plan( 115 115 + args: &ApplyGoalArgs, 116 116 + node: &Node, 117 117 + name: &Name, 118 118 + modifiers: SubCommandModifiers, 119 119 + hive_location: Arc<HiveLocation>, 120 120 + should_quit: Arc<AtomicBool>, 121 121 + ) -> NodePlan { 122 122 + let ApplyGoalArgs { 123 123 + goal, 124 124 + should_apply_locally, 125 125 + no_keys, 126 126 + substitute_on_destination, 127 127 + reboot, 128 128 + handle_unreachable, 129 129 + .. 130 130 + } = args; 131 131 + 132 132 + let mut steps: Vec<Step> = Vec::new(); 133 133 + let mut end: Vec<Step> = Vec::new(); 134 134 + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); 135 135 + 136 136 + if !*should_apply_locally { 137 137 + steps.push(Step::Ping(Ping { 138 138 + target: target.clone(), 139 139 + })); 140 140 + } 141 141 + 142 142 + if !*no_keys 143 143 + && matches!( 144 144 + &goal, 145 145 + ApplyGoal::Keys | ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) 146 146 + ) 147 147 + { 148 148 + let (pre, post) = apply_plan_keys(args, node, &target); 149 149 + steps.extend(pre); 150 150 + end.extend(post); 151 151 + } 152 152 + 153 153 + if !matches!(goal, ApplyGoal::Keys) { 154 154 + steps.push(Step::Evaluate(Evaluate)); 155 155 + } 156 156 + 157 157 + if !matches!(goal, ApplyGoal::Keys) 158 158 + && !should_apply_locally 159 159 + && (node.build_remotely || matches!(goal, ApplyGoal::Push)) 160 160 + { 161 161 + steps.push(Step::PushEvaluatedOutput(PushEvaluatedOutput { 162 162 + substitute_on_destination: *substitute_on_destination, 163 163 + target: target.clone(), 164 164 + })); 165 165 + } 166 166 + 167 167 + if !matches!(goal, ApplyGoal::Keys | ApplyGoal::Push) { 168 168 + steps.push(Step::Build(Build { 169 169 + target: if node.build_remotely && !*should_apply_locally { 170 170 + Some(target.clone()) 171 171 + } else { 172 172 + None 173 173 + }, 174 174 + })); 175 175 + } 176 176 + 177 177 + if !node.build_remotely 178 178 + && !should_apply_locally 179 179 + && !matches!(goal, ApplyGoal::Keys | ApplyGoal::Push) 180 180 + { 181 181 + steps.push(Step::PushBuildOutput(PushBuildOutput { 182 182 + substitute_on_destination: *substitute_on_destination, 183 183 + target: target.clone(), 184 184 + })); 185 185 + } 186 186 + 187 187 + if let ApplyGoal::SwitchToConfiguration(goal) = goal { 188 188 + steps.push(Step::SwitchToConfiguration(SwitchToConfiguration { 189 189 + goal: *goal, 190 190 + reboot: *reboot, 191 191 + target: if *should_apply_locally { 192 192 + None 193 193 + } else { 194 194 + Some(target.clone()) 195 195 + }, 196 196 + privilege_escalation_command: node.privilege_escalation_command.clone(), 197 197 + })); 198 198 + } 199 199 + 200 200 + steps.extend(end); 201 201 + 202 202 + NodePlan { 203 203 + context: Context { 204 204 + state: StepState::default(), 205 205 + name: name.clone(), 206 206 + hive_location, 207 207 + modifiers, 208 208 + should_quit, 209 209 + }, 210 210 + steps, 211 211 + greedy_evaluate: !matches!(&goal, ApplyGoal::Keys), 212 212 + ignore_failed_ping: matches!(handle_unreachable, HandleUnreachable::Ignore), 213 213 + } 214 214 + } 215 215 + 216 216 + #[allow(clippy::too_many_lines)] 217 217 + pub fn plan_for_node( 218 218 + node: &Node, 219 219 + name: Name, 220 220 + goal: &'_ Goal, 221 221 + hive_location: Arc<HiveLocation>, 222 222 + modifiers: &SubCommandModifiers, 223 223 + should_quit: Arc<AtomicBool>, 224 224 + ) -> NodePlan { 225 225 + match goal { 226 226 + Goal::Build => NodePlan { 227 227 + context: Context { 228 228 + state: StepState::default(), 229 229 + modifiers: *modifiers, 230 230 + hive_location, 231 231 + should_quit, 232 232 + name, 233 233 + }, 234 234 + steps: vec![ 235 235 + Step::Evaluate(Evaluate), 236 236 + Step::Build(Build { target: None }), 237 237 + ], 238 238 + greedy_evaluate: true, 239 239 + ignore_failed_ping: false, 240 240 + }, 241 241 + Goal::Apply(args) => apply_plan(args, node, &name, *modifiers, hive_location, should_quit), 242 242 + } 243 243 + } 244 244 + 245 245 + #[cfg(test)] 246 246 + mod tests { 247 247 + use tokio::sync::RwLock; 248 248 + 249 249 + use crate::{ 250 250 + SubCommandModifiers, function_name, get_test_path, 251 251 + hive::{ 252 252 + node::{ 253 253 + ApplyGoal, HandleUnreachable, Name, Node, SharedTarget, Step, 254 254 + SwitchToConfigurationGoal, 255 255 + }, 256 256 + plan::{ApplyGoalArgs, Goal, plan_for_node}, 257 257 + steps::{ 258 258 + activate::SwitchToConfiguration, 259 259 + build::Build, 260 260 + evaluate::Evaluate, 261 261 + keys::{Key, Keys, PushKeyAgent, Source, UploadKeyAt}, 262 262 + ping::Ping, 263 263 + push::PushEvaluatedOutput, 264 264 + }, 265 265 + }, 266 266 + location, 267 267 + }; 268 268 + use std::path::PathBuf; 269 269 + use std::{ 270 270 + env, 271 271 + sync::{Arc, atomic::AtomicBool}, 272 272 + }; 273 273 + 274 274 + fn new_key(upload_at: &UploadKeyAt) -> Key { 275 275 + Key { 276 276 + upload_at: upload_at.clone(), 277 277 + source: Source::String(match upload_at { 278 278 + UploadKeyAt::PreActivation => "pre".into(), 279 279 + UploadKeyAt::PostActivation => "post".into(), 280 280 + UploadKeyAt::NoFilter => "none".into(), 281 281 + }), 282 282 + ..Default::default() 283 283 + } 284 284 + } 285 285 + 286 286 + #[tokio::test] 287 287 + async fn order_build() { 288 288 + let location = location!(get_test_path!()); 289 289 + let node = Node { 290 290 + build_remotely: false, 291 291 + ..Default::default() 292 292 + }; 293 293 + let name = &Name(function_name!().into()); 294 294 + let should_quit = Arc::new(AtomicBool::new(false)); 295 295 + let plan = plan_for_node( 296 296 + &node, 297 297 + name.clone(), 298 298 + &Goal::Build, 299 299 + location.clone().into(), 300 300 + &SubCommandModifiers::default(), 301 301 + should_quit.clone(), 302 302 + ); 303 303 + 304 304 + assert_eq!( 305 305 + plan.steps, 306 306 + vec![ 307 307 + Evaluate.into(), 308 308 + Build { target: None }.into() // TODO: this was previously used in an old test, may lose 309 309 + // coverage by deleting it. 310 310 + // Ping { }.into(), 311 311 + // PushKeyAgent { host_platform: "x86_64-linux".into(), substitute_on_destination: true, target: Target::default() }.into(), 312 312 + // Keys { .. }.into(), 313 313 + // crate::hive::steps::evaluate::Evaluate.into(), 314 314 + // crate::hive::steps::build::Build { .. }.into(), 315 315 + // crate::hive::steps::push::PushBuildOutput { .. }.into(), 316 316 + // SwitchToConfiguration { .. }.into(), 317 317 + // Keys { 318 318 + // filter: UploadKeyAt::PostActivation 319 319 + // } 320 320 + // .into(), 321 321 + ] 322 322 + ); 323 323 + } 324 324 + 325 325 + #[tokio::test] 326 326 + async fn order_apply_build() { 327 327 + let location = location!(get_test_path!()); 328 328 + let node = Node { 329 329 + build_remotely: true, 330 330 + ..Default::default() 331 331 + }; 332 332 + let name = &Name(function_name!().into()); 333 333 + let should_quit = Arc::new(AtomicBool::new(false)); 334 334 + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); 335 335 + let plan = plan_for_node( 336 336 + &node, 337 337 + name.clone(), 338 338 + &Goal::Apply(ApplyGoalArgs { 339 339 + goal: ApplyGoal::Build, 340 340 + should_apply_locally: false, 341 341 + no_keys: true, 342 342 + substitute_on_destination: true, 343 343 + reboot: false, 344 344 + host_platform: "x86_64-linux".into(), 345 345 + handle_unreachable: HandleUnreachable::default(), 346 346 + }), 347 347 + location.clone().into(), 348 348 + &SubCommandModifiers::default(), 349 349 + should_quit.clone(), 350 350 + ); 351 351 + 352 352 + assert_eq!( 353 353 + plan.steps, 354 354 + vec![ 355 355 + Ping { 356 356 + target: target.clone() 357 357 + } 358 358 + .into(), 359 359 + crate::hive::steps::evaluate::Evaluate.into(), 360 360 + crate::hive::steps::push::PushEvaluatedOutput { 361 361 + substitute_on_destination: true, 362 362 + target: target.clone() 363 363 + } 364 364 + .into(), 365 365 + crate::hive::steps::build::Build { 366 366 + target: Some(target.clone()) 367 367 + } 368 368 + .into(), 369 369 + ] 370 370 + ); 371 371 + 372 372 + let node = Node { 373 373 + build_remotely: false, 374 374 + ..Default::default() 375 375 + }; 376 376 + let plan = plan_for_node( 377 377 + &node, 378 378 + name.clone(), 379 379 + &Goal::Apply(ApplyGoalArgs { 380 380 + goal: ApplyGoal::Build, 381 381 + should_apply_locally: false, 382 382 + no_keys: true, 383 383 + substitute_on_destination: true, 384 384 + reboot: false, 385 385 + host_platform: "x86_64-linux".into(), 386 386 + handle_unreachable: HandleUnreachable::default(), 387 387 + }), 388 388 + location.clone().into(), 389 389 + &SubCommandModifiers::default(), 390 390 + should_quit.clone(), 391 391 + ); 392 392 + 393 393 + assert_eq!( 394 394 + plan.steps, 395 395 + vec![ 396 396 + Ping { 397 397 + target: target.clone() 398 398 + } 399 399 + .into(), 400 400 + crate::hive::steps::evaluate::Evaluate.into(), 401 401 + crate::hive::steps::build::Build { target: None }.into(), 402 402 + crate::hive::steps::push::PushBuildOutput { 403 403 + substitute_on_destination: true, 404 404 + target: target.clone() 405 405 + } 406 406 + .into(), 407 407 + ] 408 408 + ); 409 409 + } 410 410 + 411 411 + #[tokio::test] 412 412 + async fn order_keys_only() { 413 413 + let location = location!(get_test_path!()); 414 414 + let node = Node { 415 415 + keys: vec![ 416 416 + new_key(&UploadKeyAt::PreActivation).into(), 417 417 + new_key(&UploadKeyAt::PostActivation).into(), 418 418 + new_key(&UploadKeyAt::PreActivation).into(), 419 419 + new_key(&UploadKeyAt::PostActivation).into(), 420 420 + ], 421 421 + ..Default::default() 422 422 + }; 423 423 + let name = &Name(function_name!().into()); 424 424 + let should_quit = Arc::new(AtomicBool::new(false)); 425 425 + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); 426 426 + let plan_apply_keys = plan_for_node( 427 427 + &node.clone(), 428 428 + name.clone(), 429 429 + &Goal::Apply(ApplyGoalArgs { 430 430 + goal: ApplyGoal::Keys, 431 431 + should_apply_locally: false, 432 432 + no_keys: false, 433 433 + substitute_on_destination: true, 434 434 + reboot: false, 435 435 + host_platform: "x86_64-linux".into(), 436 436 + handle_unreachable: HandleUnreachable::default(), 437 437 + }), 438 438 + location.clone().into(), 439 439 + &SubCommandModifiers::default(), 440 440 + should_quit.clone(), 441 441 + ); 442 442 + 443 443 + assert_eq!( 444 444 + plan_apply_keys.steps, 445 445 + vec![ 446 446 + Ping { 447 447 + target: target.clone() 448 448 + } 449 449 + .into(), 450 450 + PushKeyAgent { 451 451 + substitute_on_destination: true, 452 452 + target: Some(target.clone()), 453 453 + host_platform: node.host_platform.clone() 454 454 + } 455 455 + .into(), 456 456 + Keys { 457 457 + target: Some(target.clone()), 458 458 + // test that all keys are included 459 459 + keys: node.keys.clone(), 460 460 + privilege_escalation_command: node.privilege_escalation_command.clone() 461 461 + } 462 462 + .into(), 463 463 + ] 464 464 + ); 465 465 + } 466 466 + 467 467 + #[tokio::test] 468 468 + async fn order_key_split() { 469 469 + let location = location!(get_test_path!()); 470 470 + let node = Node { 471 471 + keys: vec![ 472 472 + new_key(&UploadKeyAt::PreActivation).into(), 473 473 + new_key(&UploadKeyAt::PostActivation).into(), 474 474 + new_key(&UploadKeyAt::PreActivation).into(), 475 475 + new_key(&UploadKeyAt::PostActivation).into(), 476 476 + ], 477 477 + ..Default::default() 478 478 + }; 479 479 + let name = &Name(function_name!().into()); 480 480 + let should_quit = Arc::new(AtomicBool::new(false)); 481 481 + 482 482 + // Test that keys are split by their `upload_at`, also tests that key 483 483 + // step's `target` abides by should_apply_locally 484 484 + let plan_activate_with_keys = plan_for_node( 485 485 + &node, 486 486 + name.clone(), 487 487 + &Goal::Apply(ApplyGoalArgs { 488 488 + goal: ApplyGoal::SwitchToConfiguration( 489 489 + crate::hive::node::SwitchToConfigurationGoal::Switch, 490 490 + ), 491 491 + should_apply_locally: true, 492 492 + no_keys: false, 493 493 + substitute_on_destination: true, 494 494 + reboot: false, 495 495 + host_platform: "x86_64-linux".into(), 496 496 + handle_unreachable: HandleUnreachable::default(), 497 497 + }), 498 498 + location.clone().into(), 499 499 + &SubCommandModifiers::default(), 500 500 + should_quit.clone(), 501 501 + ); 502 502 + 503 503 + assert_eq!( 504 504 + plan_activate_with_keys 505 505 + .steps 506 506 + .into_iter() 507 507 + .filter(|x| matches!( 508 508 + x, 509 509 + Step::Keys(Keys { .. }) | Step::PushKeyAgent(PushKeyAgent { .. }) 510 510 + )) 511 511 + .collect::<Vec<Step>>(), 512 512 + vec![ 513 513 + PushKeyAgent { 514 514 + substitute_on_destination: true, 515 515 + target: None, 516 516 + host_platform: node.host_platform.clone() 517 517 + } 518 518 + .into(), 519 519 + Keys { 520 520 + target: None, 521 521 + keys: node 522 522 + .keys 523 523 + .iter() 524 524 + .filter(|key| matches!(key.upload_at, UploadKeyAt::PreActivation)) 525 525 + .cloned() 526 526 + .collect::<Vec<_>>(), 527 527 + privilege_escalation_command: node.privilege_escalation_command.clone() 528 528 + } 529 529 + .into(), 530 530 + Keys { 531 531 + target: None, 532 532 + keys: node 533 533 + .keys 534 534 + .iter() 535 535 + .filter(|key| matches!(key.upload_at, UploadKeyAt::PostActivation)) 536 536 + .cloned() 537 537 + .collect::<Vec<_>>(), 538 538 + privilege_escalation_command: node.privilege_escalation_command.clone() 539 539 + } 540 540 + .into(), 541 541 + ] 542 542 + ); 543 543 + } 544 544 + 545 545 + #[tokio::test] 546 546 + async fn order_push_only() { 547 547 + let location = location!(get_test_path!()); 548 548 + let node = Node::default(); 549 549 + let name = &Name(function_name!().into()); 550 550 + let should_quit = Arc::new(AtomicBool::new(false)); 551 551 + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); 552 552 + let plan = plan_for_node( 553 553 + &node.clone(), 554 554 + name.clone(), 555 555 + &Goal::Apply(ApplyGoalArgs { 556 556 + goal: ApplyGoal::Push, 557 557 + should_apply_locally: false, 558 558 + no_keys: false, 559 559 + substitute_on_destination: true, 560 560 + reboot: false, 561 561 + host_platform: "x86_64-linux".into(), 562 562 + handle_unreachable: HandleUnreachable::default(), 563 563 + }), 564 564 + location.clone().into(), 565 565 + &SubCommandModifiers::default(), 566 566 + should_quit.clone(), 567 567 + ); 568 568 + 569 569 + assert_eq!( 570 570 + plan.steps, 571 571 + vec![ 572 572 + Ping { 573 573 + target: target.clone() 574 574 + } 575 575 + .into(), 576 576 + Evaluate.into(), 577 577 + PushEvaluatedOutput { 578 578 + substitute_on_destination: true, 579 579 + target: target.clone() 580 580 + } 581 581 + .into() 582 582 + ] 583 583 + ); 584 584 + } 585 585 + 586 586 + #[tokio::test] 587 587 + async fn order_remote_build() { 588 588 + let location = location!(get_test_path!()); 589 589 + let node = Node { 590 590 + build_remotely: true, 591 591 + ..Default::default() 592 592 + }; 593 593 + let name = &Name(function_name!().into()); 594 594 + let should_quit = Arc::new(AtomicBool::new(false)); 595 595 + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); 596 596 + let plan = plan_for_node( 597 597 + &node.clone(), 598 598 + name.clone(), 599 599 + &Goal::Apply(ApplyGoalArgs { 600 600 + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), 601 601 + should_apply_locally: false, 602 602 + no_keys: false, 603 603 + substitute_on_destination: true, 604 604 + reboot: false, 605 605 + host_platform: "x86_64-linux".into(), 606 606 + handle_unreachable: HandleUnreachable::default(), 607 607 + }), 608 608 + location.clone().into(), 609 609 + &SubCommandModifiers::default(), 610 610 + should_quit.clone(), 611 611 + ); 612 612 + 613 613 + assert_eq!( 614 614 + plan.steps, 615 615 + vec![ 616 616 + Ping { 617 617 + target: target.clone() 618 618 + } 619 619 + .into(), 620 620 + Evaluate.into(), 621 621 + PushEvaluatedOutput { 622 622 + substitute_on_destination: true, 623 623 + target: target.clone() 624 624 + } 625 625 + .into(), 626 626 + Build { 627 627 + target: Some(target.clone()) 628 628 + } 629 629 + .into(), 630 630 + SwitchToConfiguration { 631 631 + goal: SwitchToConfigurationGoal::Switch, 632 632 + reboot: false, 633 633 + target: Some(target.clone()), 634 634 + privilege_escalation_command: node.privilege_escalation_command, 635 635 + } 636 636 + .into(), 637 637 + ] 638 638 + ); 639 639 + } 640 640 + 641 641 + #[tokio::test] 642 642 + async fn order_nokeys() { 643 643 + let location = location!(get_test_path!()); 644 644 + let node = Node { 645 645 + keys: vec![Key::default().into(), Key::default().into()], 646 646 + build_remotely: true, 647 647 + ..Default::default() 648 648 + }; 649 649 + let name = &Name(function_name!().into()); 650 650 + let should_quit = Arc::new(AtomicBool::new(false)); 651 651 + let target = SharedTarget(Arc::new(RwLock::new(node.target.clone()))); 652 652 + let plan = plan_for_node( 653 653 + &node.clone(), 654 654 + name.clone(), 655 655 + &Goal::Apply(ApplyGoalArgs { 656 656 + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), 657 657 + should_apply_locally: false, 658 658 + no_keys: true, 659 659 + substitute_on_destination: true, 660 660 + reboot: false, 661 661 + host_platform: "x86_64-linux".into(), 662 662 + handle_unreachable: HandleUnreachable::default(), 663 663 + }), 664 664 + location.clone().into(), 665 665 + &SubCommandModifiers::default(), 666 666 + should_quit.clone(), 667 667 + ); 668 668 + 669 669 + assert_eq!( 670 670 + plan.steps, 671 671 + vec![ 672 672 + Ping { 673 673 + target: target.clone() 674 674 + } 675 675 + .into(), 676 676 + Evaluate.into(), 677 677 + PushEvaluatedOutput { 678 678 + substitute_on_destination: true, 679 679 + target: target.clone() 680 680 + } 681 681 + .into(), 682 682 + Build { 683 683 + target: Some(target.clone()) 684 684 + } 685 685 + .into(), 686 686 + SwitchToConfiguration { 687 687 + goal: SwitchToConfigurationGoal::Switch, 688 688 + reboot: false, 689 689 + target: Some(target.clone()), 690 690 + privilege_escalation_command: node.privilege_escalation_command, 691 691 + } 692 692 + .into(), 693 693 + ] 694 694 + ); 695 695 + } 696 696 + 697 697 + #[tokio::test] 698 698 + async fn order_should_apply_locally() { 699 699 + let location = location!(get_test_path!()); 700 700 + let node = Node::default(); 701 701 + let name = &Name(function_name!().into()); 702 702 + let should_quit = Arc::new(AtomicBool::new(false)); 703 703 + let plan = plan_for_node( 704 704 + &node.clone(), 705 705 + name.clone(), 706 706 + &Goal::Apply(ApplyGoalArgs { 707 707 + goal: ApplyGoal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch), 708 708 + should_apply_locally: true, 709 709 + no_keys: true, 710 710 + substitute_on_destination: true, 711 711 + reboot: false, 712 712 + host_platform: "x86_64-linux".into(), 713 713 + handle_unreachable: HandleUnreachable::default(), 714 714 + }), 715 715 + location.clone().into(), 716 716 + &SubCommandModifiers::default(), 717 717 + should_quit.clone(), 718 718 + ); 719 719 + 720 720 + assert_eq!( 721 721 + plan.steps, 722 722 + vec![ 723 723 + Evaluate.into(), 724 724 + Build { target: None }.into(), 725 725 + SwitchToConfiguration { 726 726 + goal: SwitchToConfigurationGoal::Switch, 727 727 + reboot: false, 728 728 + target: None, 729 729 + privilege_escalation_command: node.privilege_escalation_command, 730 730 + } 731 731 + .into(), 732 732 + ] 733 733 + ); 734 734 + } 735 735 + }

+79 -85

crates/core/src/hive/steps/activate.rs

reviewed

··· 1 1 // SPDX-License-Identifier: AGPL-3.0-or-later 2 2 // Copyright 2024-2025 wire Contributors 3 3 4 4 - use std::fmt::Display; 4 4 + use std::{fmt::Display, sync::Arc}; 5 5 6 6 use tracing::{error, info, instrument, warn}; 7 7 ··· 9 9 HiveLibError, 10 10 commands::{CommandArguments, WireCommandChip, builder::CommandStringBuilder, run_command}, 11 11 errors::{ActivationError, NetworkError}, 12 12 - hive::node::{Context, ExecuteStep, Goal, Objective, SwitchToConfigurationGoal}, 12 12 + hive::node::{Context, ExecuteStep, SharedTarget, SwitchToConfigurationGoal}, 13 13 }; 14 14 15 15 - #[derive(Debug, PartialEq)] 16 16 - pub struct SwitchToConfiguration; 15 15 + #[derive(Debug)] 16 16 + #[cfg_attr(test, derive(PartialEq))] 17 17 + pub struct SwitchToConfiguration { 18 18 + pub goal: SwitchToConfigurationGoal, 19 19 + pub reboot: bool, 20 20 + pub target: Option<SharedTarget>, 21 21 + pub privilege_escalation_command: Arc<Vec<Arc<str>>>, 22 22 + } 17 23 18 24 impl Display for SwitchToConfiguration { 19 25 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { ··· 21 27 } 22 28 } 23 29 24 24 - async fn wait_for_ping(ctx: &Context<'_>) -> Result<(), HiveLibError> { 25 25 - let host = ctx.node.target.get_preferred_host()?; 26 26 - let mut result = ctx.node.ping(ctx.modifiers).await; 30 30 + async fn wait_for_ping(target: &SharedTarget, ctx: &Context) -> Result<(), HiveLibError> { 31 31 + let target = target.0.read().await; 32 32 + let host = target.get_preferred_host()?; 27 33 28 28 - for num in 0..2 { 34 34 + for num in 0..3 { 29 35 warn!("Trying to ping {host} (attempt {}/3)", num + 1); 30 36 31 31 - result = ctx.node.ping(ctx.modifiers).await; 37 37 + let result = target.ping(ctx.modifiers).await; 32 38 33 39 if result.is_ok() { 34 40 info!("Regained connection to {} via {host}", ctx.name); 35 41 36 36 - break; 42 42 + return Ok(()); 37 43 } 38 44 } 39 45 40 40 - result 46 46 + Err(HiveLibError::NetworkError(NetworkError::HostsExhausted)) 41 47 } 42 48 43 43 - async fn set_profile( 44 44 - goal: SwitchToConfigurationGoal, 45 45 - built_path: &String, 46 46 - ctx: &Context<'_>, 47 47 - ) -> Result<(), HiveLibError> { 48 48 - info!("Setting profiles in anticipation for switch-to-configuration {goal}"); 49 49 - 50 50 - let mut command_string = CommandStringBuilder::new("nix-env"); 51 51 - command_string.args(&["-p", "/nix/var/nix/profiles/system", "--set"]); 52 52 - command_string.arg(built_path); 49 49 + impl SwitchToConfiguration { 50 50 + async fn set_profile(&self, built_path: &String, ctx: &Context) -> Result<(), HiveLibError> { 51 51 + info!( 52 52 + "Setting profiles in anticipation for switch-to-configuration {}", 53 53 + self.goal 54 54 + ); 53 55 54 54 - let Objective::Apply(apply_objective) = ctx.objective else { 55 55 - unreachable!() 56 56 - }; 56 56 + let mut command_string = CommandStringBuilder::new("nix-env"); 57 57 + command_string.args(&["-p", "/nix/var/nix/profiles/system", "--set"]); 58 58 + command_string.arg(built_path); 57 59 58 58 - let child = run_command( 59 59 - &CommandArguments::new(command_string, ctx.modifiers) 60 60 - .mode(crate::commands::ChildOutputMode::Nix) 61 61 - .execute_on_remote(if apply_objective.should_apply_locally { 62 62 - None 63 63 - } else { 64 64 - Some(&ctx.node.target) 65 65 - }) 66 66 - .elevated(ctx.node), 67 67 - ) 68 68 - .await?; 60 60 + let child = run_command( 61 61 + &CommandArguments::new(command_string, ctx.modifiers) 62 62 + .mode(crate::commands::ChildOutputMode::Nix) 63 63 + .execute_on_remote(self.target.clone()) 64 64 + .privileged(&self.privilege_escalation_command), 65 65 + ) 66 66 + .await?; 69 67 70 70 - let _ = child 71 71 - .wait_till_success() 72 72 - .await 73 73 - .map_err(HiveLibError::CommandError)?; 68 68 + let _ = child 69 69 + .wait_till_success() 70 70 + .await 71 71 + .map_err(HiveLibError::CommandError)?; 74 72 75 75 - info!("Set system profile"); 73 73 + info!("Set system profile"); 76 74 77 77 - Ok(()) 75 75 + Ok(()) 76 76 + } 78 77 } 79 78 80 79 impl ExecuteStep for SwitchToConfiguration { 81 81 - fn should_execute(&self, ctx: &Context) -> bool { 82 82 - let Objective::Apply(apply_objective) = ctx.objective else { 83 83 - return false; 84 84 - }; 85 85 - 86 86 - matches!(apply_objective.goal, Goal::SwitchToConfiguration(..)) 87 87 - } 88 88 - 89 80 #[allow(clippy::too_many_lines)] 90 81 #[instrument(skip_all, name = "activate")] 91 91 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 82 82 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 92 83 let built_path = ctx.state.build.as_ref().unwrap(); 93 84 94 94 - let Objective::Apply(apply_objective) = ctx.objective else { 95 95 - unreachable!() 96 96 - }; 97 97 - 98 98 - let Goal::SwitchToConfiguration(goal) = &apply_objective.goal else { 99 99 - unreachable!("Cannot reach as guarded by should_execute") 100 100 - }; 101 101 - 102 85 if matches!( 103 103 - goal, 86 86 + self.goal, 104 87 // switch profile if switch or boot 105 88 // https://github.com/NixOS/nixpkgs/blob/a2c92aa34735a04010671e3378e2aa2d109b2a72/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/services.py#L224 106 89 SwitchToConfigurationGoal::Switch | SwitchToConfigurationGoal::Boot 107 90 ) { 108 108 - set_profile(*goal, built_path, ctx).await?; 91 91 + self.set_profile(built_path, ctx).await?; 109 92 } 110 93 111 111 - info!("Running switch-to-configuration {goal}"); 94 94 + info!("Running switch-to-configuration {}", self.goal); 112 95 113 96 let mut command_string = 114 97 CommandStringBuilder::new(format!("{built_path}/bin/switch-to-configuration")); 115 115 - command_string.arg(match goal { 98 98 + command_string.arg(match self.goal { 116 99 SwitchToConfigurationGoal::Switch => "switch", 117 100 SwitchToConfigurationGoal::Boot => "boot", 118 101 SwitchToConfigurationGoal::Test => "test", ··· 121 104 122 105 let child = run_command( 123 106 &CommandArguments::new(command_string, ctx.modifiers) 124 124 - .execute_on_remote(if apply_objective.should_apply_locally { 125 125 - None 126 126 - } else { 127 127 - Some(&ctx.node.target) 128 128 - }) 129 129 - .elevated(ctx.node) 107 107 + .execute_on_remote(self.target.clone()) 108 108 + .privileged(&self.privilege_escalation_command) 130 109 .log_stdout(), 131 110 ) 132 111 .await?; ··· 135 114 136 115 match result { 137 116 Ok(_) => { 138 138 - if !apply_objective.reboot { 117 117 + if !self.reboot { 139 118 return Ok(()); 140 119 } 141 120 142 142 - if apply_objective.should_apply_locally { 121 121 + let Some(ref target) = self.target else { 143 122 error!("Refusing to reboot local machine!"); 144 123 145 124 return Ok(()); 146 146 - } 125 125 + }; 147 126 148 127 warn!("Rebooting {name}!", name = ctx.name); 149 128 150 129 let reboot = run_command( 151 130 &CommandArguments::new("reboot now", ctx.modifiers) 152 131 .log_stdout() 153 153 - .execute_on_remote(Some(&ctx.node.target)) 154 154 - .elevated(ctx.node), 132 132 + .execute_on_remote(Some(target.clone())) 133 133 + .privileged(&self.privilege_escalation_command), 155 134 ) 156 135 .await?; 157 136 ··· 164 143 165 144 info!("Rebooted {name}, waiting to reconnect...", name = ctx.name); 166 145 167 167 - if wait_for_ping(ctx).await.is_ok() { 146 146 + if wait_for_ping(target, ctx).await.is_ok() { 168 147 return Ok(()); 169 148 } 149 149 + 150 150 + let target = target.0.read().await; 170 151 171 152 error!( 172 153 "Failed to get regain connection to {name} via {host} after reboot.", 173 154 name = ctx.name, 174 174 - host = ctx.node.target.get_preferred_host()? 155 155 + host = target.get_preferred_host()? 175 156 ); 176 157 177 158 return Err(HiveLibError::NetworkError( 178 159 NetworkError::HostUnreachableAfterReboot( 179 179 - ctx.node.target.get_preferred_host()?.to_string(), 160 160 + target.get_preferred_host()?.to_string(), 180 161 ), 181 162 )); 182 163 } ··· 188 169 189 170 // Bail if the command couldn't of broken the system 190 171 // and don't try to regain connection to localhost 191 191 - if matches!(goal, SwitchToConfigurationGoal::DryActivate) 192 192 - || apply_objective.should_apply_locally 193 193 - { 172 172 + let Some(target) = self 173 173 + .target 174 174 + .as_ref() 175 175 + .filter(|_| !matches!(self.goal, SwitchToConfigurationGoal::DryActivate)) 176 176 + else { 194 177 return Err(HiveLibError::ActivationError( 195 195 - ActivationError::SwitchToConfigurationError(*goal, ctx.name.clone(), error), 178 178 + ActivationError::SwitchToConfigurationError( 179 179 + self.goal, 180 180 + ctx.name.clone(), 181 181 + error, 182 182 + ), 196 183 )); 197 197 - } 184 184 + }; 198 185 199 199 - if wait_for_ping(ctx).await.is_ok() { 186 186 + if wait_for_ping(target, ctx).await.is_ok() { 200 187 return Err(HiveLibError::ActivationError( 201 201 - ActivationError::SwitchToConfigurationError(*goal, ctx.name.clone(), error), 188 188 + ActivationError::SwitchToConfigurationError( 189 189 + self.goal, 190 190 + ctx.name.clone(), 191 191 + error, 192 192 + ), 202 193 )); 203 194 } 195 195 + 196 196 + let target = target.0.read().await; 204 197 205 198 error!( 206 199 "Failed to get regain connection to {name} via {host} after {goal} activation.", 207 200 name = ctx.name, 208 208 - host = ctx.node.target.get_preferred_host()? 201 201 + host = target.get_preferred_host()?, 202 202 + goal = self.goal 209 203 ); 210 204 211 205 return Err(HiveLibError::NetworkError( 212 206 NetworkError::HostUnreachableAfterReboot( 213 213 - ctx.node.target.get_preferred_host()?.to_string(), 207 207 + target.get_preferred_host()?.to_string(), 214 208 ), 215 209 )); 216 210 }

+8 -23

crates/core/src/hive/steps/build.rs

reviewed

··· 11 11 CommandArguments, Either, WireCommandChip, builder::CommandStringBuilder, 12 12 run_command_with_env, 13 13 }, 14 14 - hive::node::{Context, ExecuteStep, Goal, Objective}, 14 14 + hive::node::{Context, ExecuteStep, SharedTarget}, 15 15 }; 16 16 17 17 - #[derive(Debug, PartialEq)] 18 18 - pub struct Build; 17 17 + #[derive(Debug)] 18 18 + #[cfg_attr(test, derive(PartialEq))] 19 19 + pub struct Build { 20 20 + pub(crate) target: Option<SharedTarget>, 21 21 + } 19 22 20 23 impl Display for Build { 21 24 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { ··· 24 27 } 25 28 26 29 impl ExecuteStep for Build { 27 27 - fn should_execute(&self, ctx: &Context) -> bool { 28 28 - match ctx.objective { 29 29 - Objective::Apply(apply_objective) => { 30 30 - !matches!(apply_objective.goal, Goal::Keys | Goal::Push) 31 31 - } 32 32 - Objective::BuildLocally => true, 33 33 - } 34 34 - } 35 35 - 36 30 #[instrument(skip_all, name = "build")] 37 37 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 31 31 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 38 32 let top_level = ctx.state.evaluation.as_ref().unwrap(); 39 33 40 34 let mut command_string = CommandStringBuilder::nix(); ··· 51 45 let status = run_command_with_env( 52 46 &CommandArguments::new(command_string, ctx.modifiers) 53 47 // build remotely if asked for AND we arent applying locally 54 54 - .execute_on_remote( 55 55 - if ctx.node.build_remotely 56 56 - && let Objective::Apply(apply_objective) = ctx.objective 57 57 - && !apply_objective.should_apply_locally 58 58 - { 59 59 - Some(&ctx.node.target) 60 60 - } else { 61 61 - None 62 62 - }, 63 63 - ) 48 48 + .execute_on_remote(self.target.clone()) 64 49 .mode(crate::commands::ChildOutputMode::Nix) 65 50 .log_stdout(), 66 51 std::collections::HashMap::new(),

-28

crates/core/src/hive/steps/cleanup.rs

reviewed

··· 1 1 - // SPDX-License-Identifier: AGPL-3.0-or-later 2 2 - // Copyright 2024-2025 wire Contributors 3 3 - 4 4 - use std::fmt::Display; 5 5 - 6 6 - use crate::{ 7 7 - errors::HiveLibError, 8 8 - hive::node::{Context, ExecuteStep}, 9 9 - }; 10 10 - 11 11 - #[derive(PartialEq, Debug)] 12 12 - pub(crate) struct CleanUp; 13 13 - 14 14 - impl Display for CleanUp { 15 15 - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 16 16 - write!(f, "Clean up") 17 17 - } 18 18 - } 19 19 - 20 20 - impl ExecuteStep for CleanUp { 21 21 - fn should_execute(&self, _ctx: &Context) -> bool { 22 22 - false 23 23 - } 24 24 - 25 25 - async fn execute(&self, _ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 26 26 - Ok(()) 27 27 - } 28 28 - }

+2 -9

crates/core/src/hive/steps/evaluate.rs

reviewed

··· 7 7 8 8 use crate::{ 9 9 HiveLibError, 10 10 - hive::node::{Context, ExecuteStep, Goal, Objective}, 10 10 + hive::node::{Context, ExecuteStep}, 11 11 }; 12 12 13 13 #[derive(Debug, PartialEq)] ··· 20 20 } 21 21 22 22 impl ExecuteStep for Evaluate { 23 23 - fn should_execute(&self, ctx: &Context) -> bool { 24 24 - match ctx.objective { 25 25 - Objective::Apply(apply_objective) => !matches!(apply_objective.goal, Goal::Keys), 26 26 - Objective::BuildLocally => true, 27 27 - } 28 28 - } 29 29 - 30 23 #[instrument(skip_all, name = "eval")] 31 31 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 24 24 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 32 25 let rx = ctx.state.evaluation_rx.take().unwrap(); 33 26 34 27 ctx.state.evaluation = Some(rx.await.unwrap()?);

+36 -139

crates/core/src/hive/steps/keys.rs

reviewed

··· 4 4 use base64::Engine; 5 5 use base64::prelude::BASE64_STANDARD; 6 6 use futures::future::join_all; 7 7 - use im::Vector; 8 7 use itertools::{Itertools, Position}; 9 8 use owo_colors::OwoColorize; 10 9 use prost::Message; ··· 19 18 use std::pin::Pin; 20 19 use std::process::Stdio; 21 20 use std::str::from_utf8; 21 21 + use std::sync::Arc; 22 22 use std::vec::IntoIter; 23 23 use tokio::io::AsyncReadExt as _; 24 24 use tokio::process::Command; ··· 31 31 use crate::commands::common::push; 32 32 use crate::commands::{CommandArguments, WireCommandChip, run_command}; 33 33 use crate::errors::KeyError; 34 34 - use crate::hive::node::{Context, ExecuteStep, Goal, Objective, Push, SwitchToConfigurationGoal}; 34 34 + use crate::hive::node::{Context, ExecuteStep, Push, SharedTarget}; 35 35 36 36 #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Hash)] 37 37 #[serde(tag = "t", content = "c")] ··· 179 179 )) 180 180 } 181 181 182 182 - #[derive(Debug, PartialEq)] 182 182 + #[derive(Debug)] 183 183 + #[cfg_attr(test, derive(PartialEq))] 183 184 pub struct Keys { 184 184 - pub filter: UploadKeyAt, 185 185 + pub keys: Vec<Arc<Key>>, 186 186 + pub target: Option<SharedTarget>, 187 187 + pub privilege_escalation_command: Arc<Vec<Arc<str>>>, 188 188 + } 189 189 + 190 190 + #[derive(Debug)] 191 191 + #[cfg_attr(test, derive(PartialEq))] 192 192 + pub struct PushKeyAgent { 193 193 + pub substitute_on_destination: bool, 194 194 + pub host_platform: Arc<str>, 195 195 + pub target: Option<SharedTarget>, 185 196 } 186 186 - #[derive(Debug, PartialEq)] 187 187 - pub struct PushKeyAgent; 188 197 189 198 impl Display for Keys { 190 199 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 191 191 - write!(f, "Upload key @ {:?}", self.filter) 200 200 + write!(f, "Upload {} key(s)", self.keys.len()) 192 201 } 193 202 } 194 203 ··· 225 234 } 226 235 227 236 impl ExecuteStep for Keys { 228 228 - fn should_execute(&self, ctx: &Context) -> bool { 229 229 - let Objective::Apply(apply_objective) = ctx.objective else { 230 230 - return false; 231 231 - }; 232 232 - 233 233 - if apply_objective.no_keys { 234 234 - return false; 235 235 - } 236 236 - 237 237 - // should execute if no filter, and the goal is keys. 238 238 - // otherwise, only execute if the goal is switch and non-nofilter 239 239 - matches!( 240 240 - (&self.filter, &apply_objective.goal), 241 241 - (UploadKeyAt::NoFilter, Goal::Keys) 242 242 - | ( 243 243 - UploadKeyAt::PreActivation | UploadKeyAt::PostActivation, 244 244 - Goal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) 245 245 - ) 246 246 - ) 247 247 - } 248 248 - 249 237 #[instrument(skip_all, name = "keys")] 250 250 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 238 238 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 251 239 let agent_directory = ctx.state.key_agent_directory.as_ref().unwrap(); 252 240 253 253 - let mut keys = self.select_keys(&ctx.node.keys).await?; 241 241 + let mut keys = self.select_keys(&self.keys).await?; 254 242 255 243 if keys.peek().is_none() { 256 244 debug!("Had no keys to push, ending KeyStep early."); ··· 259 247 260 248 let command_string = 261 249 CommandStringBuilder::new(format!("{agent_directory}/bin/wire-key-agent")); 262 262 - 263 263 - let Objective::Apply(apply_objective) = ctx.objective else { 264 264 - unreachable!() 265 265 - }; 266 250 267 251 let mut child = run_command( 268 252 &CommandArguments::new(command_string, ctx.modifiers) 269 269 - .execute_on_remote(if apply_objective.should_apply_locally { 270 270 - None 271 271 - } else { 272 272 - Some(&ctx.node.target) 273 273 - }) 274 274 - .elevated(ctx.node) 253 253 + .execute_on_remote(self.target.clone()) 254 254 + .privileged(&self.privilege_escalation_command) 275 255 .keep_stdin_open() 276 256 .log_stdout(), 277 257 ) ··· 306 286 impl Keys { 307 287 async fn select_keys( 308 288 &self, 309 309 - keys: &Vector<Key>, 289 289 + keys: &[Arc<Key>], 310 290 ) -> Result<Peekable<IntoIter<(wire_key_agent::keys::KeySpec, std::vec::Vec<u8>)>>, HiveLibError> 311 291 { 312 312 - let futures = keys 313 313 - .iter() 314 314 - .filter(|key| self.filter == UploadKeyAt::NoFilter || (key.upload_at == self.filter)) 315 315 - .map(|key| async move { 316 316 - process_key(key) 317 317 - .await 318 318 - .map_err(|err| HiveLibError::KeyError(key.name.clone(), err)) 319 319 - }); 292 292 + let futures = keys.iter().map(|key| async move { 293 293 + process_key(key) 294 294 + .await 295 295 + .map_err(|err| HiveLibError::KeyError(key.name.clone(), err)) 296 296 + }); 320 297 321 298 Ok(join_all(futures) 322 299 .await ··· 328 305 } 329 306 330 307 impl ExecuteStep for PushKeyAgent { 331 331 - fn should_execute(&self, ctx: &Context) -> bool { 332 332 - let Objective::Apply(apply_objective) = ctx.objective else { 333 333 - return false; 334 334 - }; 335 335 - 336 336 - if apply_objective.no_keys { 337 337 - return false; 338 338 - } 339 339 - 340 340 - matches!( 341 341 - &apply_objective.goal, 342 342 - Goal::Keys | Goal::SwitchToConfiguration(SwitchToConfigurationGoal::Switch) 343 343 - ) 344 344 - } 345 345 - 346 308 #[instrument(skip_all, name = "push_agent")] 347 347 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 309 309 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 348 310 let arg_name = format!( 349 311 "WIRE_KEY_AGENT_{platform}", 350 350 - platform = ctx.node.host_platform.replace('-', "_") 312 312 + platform = self.host_platform.replace('-', "_") 351 313 ); 352 314 353 315 let agent_directory = match env::var_os(&arg_name) { ··· 359 321 ), 360 322 }; 361 323 362 362 - let Objective::Apply(apply_objective) = ctx.objective else { 363 363 - unreachable!() 364 364 - }; 365 365 - 366 366 - if !apply_objective.should_apply_locally { 367 367 - push(ctx, Push::Path(&agent_directory)).await?; 324 324 + if let Some(ref target) = self.target { 325 325 + push( 326 326 + ctx, 327 327 + target, 328 328 + Push::Path(&agent_directory), 329 329 + self.substitute_on_destination, 330 330 + ) 331 331 + .await?; 368 332 } 369 333 370 334 ctx.state.key_agent_directory = Some(agent_directory); ··· 372 336 Ok(()) 373 337 } 374 338 } 375 375 - 376 376 - #[cfg(test)] 377 377 - mod tests { 378 378 - use im::Vector; 379 379 - 380 380 - use crate::hive::steps::keys::{Key, Keys, UploadKeyAt, process_key}; 381 381 - 382 382 - fn new_key(upload_at: &UploadKeyAt) -> Key { 383 383 - Key { 384 384 - upload_at: upload_at.clone(), 385 385 - source: super::Source::String(match upload_at { 386 386 - UploadKeyAt::PreActivation => "pre".into(), 387 387 - UploadKeyAt::PostActivation => "post".into(), 388 388 - UploadKeyAt::NoFilter => "none".into(), 389 389 - }), 390 390 - ..Default::default() 391 391 - } 392 392 - } 393 393 - 394 394 - #[tokio::test] 395 395 - async fn key_filtering() { 396 396 - let keys = Vector::from(vec![ 397 397 - new_key(&UploadKeyAt::PreActivation), 398 398 - new_key(&UploadKeyAt::PostActivation), 399 399 - new_key(&UploadKeyAt::PreActivation), 400 400 - new_key(&UploadKeyAt::PostActivation), 401 401 - ]); 402 402 - 403 403 - for (_, buf) in (Keys { 404 404 - filter: crate::hive::steps::keys::UploadKeyAt::PreActivation, 405 405 - }) 406 406 - .select_keys(&keys) 407 407 - .await 408 408 - .unwrap() 409 409 - { 410 410 - assert_eq!(String::from_utf8_lossy(&buf), "pre"); 411 411 - } 412 412 - 413 413 - for (_, buf) in (Keys { 414 414 - filter: crate::hive::steps::keys::UploadKeyAt::PostActivation, 415 415 - }) 416 416 - .select_keys(&keys) 417 417 - .await 418 418 - .unwrap() 419 419 - { 420 420 - assert_eq!(String::from_utf8_lossy(&buf), "post"); 421 421 - } 422 422 - 423 423 - // test that NoFilter processes all keys. 424 424 - let processed_all = 425 425 - futures::future::join_all(keys.iter().map(async |x| process_key(x).await)) 426 426 - .await 427 427 - .iter() 428 428 - .flatten() 429 429 - .cloned() 430 430 - .collect::<Vec<_>>(); 431 431 - let no_filter = (Keys { 432 432 - filter: crate::hive::steps::keys::UploadKeyAt::NoFilter, 433 433 - }) 434 434 - .select_keys(&keys) 435 435 - .await 436 436 - .unwrap() 437 437 - .collect::<Vec<_>>(); 438 438 - 439 439 - assert_eq!(processed_all, no_filter); 440 440 - } 441 441 - }

-1

crates/core/src/hive/steps/mod.rs

reviewed

··· 3 3 4 4 pub mod activate; 5 5 pub mod build; 6 6 - pub mod cleanup; 7 6 pub mod evaluate; 8 7 pub mod keys; 9 8 pub mod ping;

+17 -17

crates/core/src/hive/steps/ping.rs

reviewed

··· 7 7 8 8 use crate::{ 9 9 HiveLibError, 10 10 - hive::node::{Context, ExecuteStep, Objective}, 10 10 + hive::node::{Context, ExecuteStep, SharedTarget}, 11 11 }; 12 12 13 13 - #[derive(Debug, PartialEq)] 14 14 - pub struct Ping; 13 13 + #[derive(Debug)] 14 14 + #[cfg_attr(test, derive(PartialEq))] 15 15 + pub struct Ping { 16 16 + pub target: SharedTarget, 17 17 + } 15 18 16 19 impl Display for Ping { 17 20 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { ··· 20 23 } 21 24 22 25 impl ExecuteStep for Ping { 23 23 - fn should_execute(&self, ctx: &Context) -> bool { 24 24 - let Objective::Apply(apply_objective) = ctx.objective else { 25 25 - return false; 26 26 - }; 27 27 - 28 28 - !apply_objective.should_apply_locally 29 29 - } 30 30 - 31 26 #[instrument(skip_all, name = "ping")] 32 32 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 27 27 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 33 28 loop { 29 29 + let target = self.target.0.read().await; 30 30 + 34 31 event!( 35 32 Level::INFO, 36 33 status = "attempting", 37 37 - host = ctx.node.target.get_preferred_host()?.to_string() 34 34 + host = target.get_preferred_host()?.to_string() 38 35 ); 39 36 40 40 - if ctx.node.ping(ctx.modifiers).await.is_ok() { 37 37 + if target.ping(ctx.modifiers).await.is_ok() { 41 38 event!( 42 39 Level::INFO, 43 40 status = "success", 44 44 - host = ctx.node.target.get_preferred_host()?.to_string() 41 41 + host = target.get_preferred_host()?.to_string() 45 42 ); 46 43 return Ok(()); 47 44 } ··· 50 47 event!( 51 48 Level::WARN, 52 49 status = "failed to ping", 53 53 - host = ctx.node.target.get_preferred_host()?.to_string() 50 50 + host = target.get_preferred_host()?.to_string() 54 51 ); 55 55 - ctx.node.target.host_failed(); 52 52 + 53 53 + drop(target); 54 54 + 55 55 + self.target.0.write().await.host_failed(); 56 56 } 57 57 } 58 58 }

+30 -42

crates/core/src/hive/steps/push.rs

reviewed

··· 8 8 use crate::{ 9 9 HiveLibError, 10 10 commands::common::push, 11 11 - hive::node::{Context, ExecuteStep, Goal, Objective}, 11 11 + hive::node::{Context, ExecuteStep, SharedTarget}, 12 12 }; 13 13 14 14 - #[derive(Debug, PartialEq)] 15 15 - pub struct PushEvaluatedOutput; 16 16 - #[derive(Debug, PartialEq)] 17 17 - pub struct PushBuildOutput; 14 14 + #[derive(Debug)] 15 15 + #[cfg_attr(test, derive(PartialEq))] 16 16 + pub struct PushEvaluatedOutput { 17 17 + pub substitute_on_destination: bool, 18 18 + pub target: SharedTarget, 19 19 + } 20 20 + 21 21 + #[derive(Debug)] 22 22 + #[cfg_attr(test, derive(PartialEq))] 23 23 + pub struct PushBuildOutput { 24 24 + pub substitute_on_destination: bool, 25 25 + pub target: SharedTarget, 26 26 + } 18 27 19 28 impl Display for PushEvaluatedOutput { 20 29 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { ··· 29 38 } 30 39 31 40 impl ExecuteStep for PushEvaluatedOutput { 32 32 - fn should_execute(&self, ctx: &Context) -> bool { 33 33 - let Objective::Apply(apply_objective) = ctx.objective else { 34 34 - return false; 35 35 - }; 36 36 - 37 37 - !matches!(apply_objective.goal, Goal::Keys) 38 38 - && !apply_objective.should_apply_locally 39 39 - && (ctx.node.build_remotely | matches!(apply_objective.goal, Goal::Push)) 40 40 - } 41 41 - 42 41 #[instrument(skip_all, name = "push_eval")] 43 43 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 42 42 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 44 43 let top_level = ctx.state.evaluation.as_ref().unwrap(); 45 44 46 46 - push(ctx, crate::hive::node::Push::Derivation(top_level)).await?; 45 45 + push( 46 46 + ctx, 47 47 + &self.target, 48 48 + crate::hive::node::Push::Derivation(top_level), 49 49 + self.substitute_on_destination, 50 50 + ) 51 51 + .await?; 47 52 48 53 Ok(()) 49 54 } 50 55 } 51 56 52 57 impl ExecuteStep for PushBuildOutput { 53 53 - fn should_execute(&self, ctx: &Context) -> bool { 54 54 - let Objective::Apply(apply_objective) = ctx.objective else { 55 55 - return false; 56 56 - }; 57 57 - 58 58 - if matches!(apply_objective.goal, Goal::Keys | Goal::Push) { 59 59 - // skip if we are not building 60 60 - return false; 61 61 - } 62 62 - 63 63 - if ctx.node.build_remotely { 64 64 - // skip if we are building remotely 65 65 - return false; 66 66 - } 67 67 - 68 68 - if apply_objective.should_apply_locally { 69 69 - // skip step if we are applying locally 70 70 - return false; 71 71 - } 72 72 - 73 73 - true 74 74 - } 75 75 - 76 58 #[instrument(skip_all, name = "push_build")] 77 77 - async fn execute(&self, ctx: &mut Context<'_>) -> Result<(), HiveLibError> { 59 59 + async fn execute(&self, ctx: &mut Context) -> Result<(), HiveLibError> { 78 60 let built_path = ctx.state.build.as_ref().unwrap(); 79 61 80 80 - push(ctx, crate::hive::node::Push::Path(built_path)).await?; 62 62 + push( 63 63 + ctx, 64 64 + &self.target, 65 65 + crate::hive::node::Push::Path(built_path), 66 66 + self.substitute_on_destination, 67 67 + ) 68 68 + .await?; 81 69 82 70 Ok(()) 83 71 }