Auto rollback if deployment is not confirmed

3 files changed, 226 insertions, 75 deletions

diff --git a/src/activate.rs b/src/activate.rs
index 0b7d28d..ce6b286 100644
--- a/src/activate.rs
+++ b/src/activate.rs
@@ -4,11 +4,20 @@
 
 use clap::Clap;
 
+use futures_util::FutureExt;
 use std::process::Stdio;
+use tokio::fs;
 use tokio::process::Command;
+use tokio::time::timeout;
+
+use std::time::Duration;
+
+use futures_util::StreamExt;
 
 use std::path::Path;
 
+use inotify::Inotify;
+
 extern crate pretty_env_logger;
 #[macro_use]
 extern crate log;
@@ -25,6 +34,8 @@ mod utils;
 struct Opts {
     profile_path: String,
     closure: String,
+    temp_path: String,
+    max_time: u16,
 
     /// Command for bootstrapping
     #[clap(long)]
@@ -35,11 +46,162 @@ struct Opts {
     auto_rollback: bool,
 }
 
+pub async fn deactivate(profile_path: &str) -> Result<(), Box<dyn std::error::Error>> {
+    error!("De-activating due to error");
+
+    let nix_env_rollback_exit_status = Command::new("nix-env")
+        .arg("-p")
+        .arg(&profile_path)
+        .arg("--rollback")
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .status()
+        .await?;
+
+    if !nix_env_rollback_exit_status.success() {
+        good_panic!("`nix-env --rollback` failed");
+    }
+
+    debug!("Listing generations");
+
+    let nix_env_list_generations_out = Command::new("nix-env")
+        .arg("-p")
+        .arg(&profile_path)
+        .arg("--list-generations")
+        .output()
+        .await?;
+
+    if !nix_env_list_generations_out.status.success() {
+        good_panic!("Listing `nix-env` generations failed");
+    }
+
+    let generations_list = String::from_utf8(nix_env_list_generations_out.stdout)?;
+
+    let last_generation_line = generations_list
+        .lines()
+        .last()
+        .expect("Expected to find a generation in list");
+
+    let last_generation_id = last_generation_line
+        .split_whitespace()
+        .next()
+        .expect("Expected to get ID from generation entry");
+
+    debug!("Removing generation entry {}", last_generation_line);
+    warn!("Removing generation by ID {}", last_generation_id);
+
+    let nix_env_delete_generation_exit_status = Command::new("nix-env")
+        .arg("-p")
+        .arg(&profile_path)
+        .arg("--delete-generations")
+        .arg(last_generation_id)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .status()
+        .await?;
+
+    if !nix_env_delete_generation_exit_status.success() {
+        good_panic!("Failed to delete failed generation");
+    }
+
+    info!("Attempting re-activate last generation");
+
+    let re_activate_exit_status = Command::new(format!("{}/deploy-rs-activate", profile_path))
+        .env("PROFILE", &profile_path)
+        .current_dir(&profile_path)
+        .status()
+        .await?;
+
+    if !re_activate_exit_status.success() {
+        good_panic!("Failed to re-activate the last generation");
+    }
+
+    Ok(())
+}
+
+async fn deactivate_on_err<A, B: core::fmt::Debug>(profile_path: &str, r: Result<A, B>) -> A {
+    match r {
+        Ok(x) => x,
+        Err(err) => {
+            error!("Deactivating due to error: {:?}", err);
+            match deactivate(profile_path).await {
+                Ok(_) => (),
+                Err(err) => {
+                    error!("Error de-activating, uh-oh: {:?}", err);
+                }
+            };
+
+            std::process::exit(1);
+        }
+    }
+}
+
+pub async fn activation_confirmation(
+    profile_path: String,
+    temp_path: String,
+    max_time: u16,
+    closure: String,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let lock_hash = &closure[11 /* /nix/store/ */ ..];
+    let lock_path = format!("{}/activating-{}", temp_path, lock_hash);
+
+    if let Some(parent) = Path::new(&lock_path).parent() {
+        fs::create_dir_all(parent).await?;
+    }
+
+    fs::File::create(&lock_path).await?;
+
+    let mut inotify = Inotify::init()?;
+    inotify.add_watch(lock_path, inotify::WatchMask::DELETE)?;
+
+    match fork::daemon(false, false).map_err(|x| x.to_string())? {
+        fork::Fork::Child => {
+            std::thread::spawn(move || {
+                let mut rt = tokio::runtime::Runtime::new().unwrap();
+
+                rt.block_on(async move {
+                    info!("Waiting for confirmation event...");
+
+                    let mut buffer = [0; 32];
+                    let mut stream =
+                        deactivate_on_err(&profile_path, inotify.event_stream(&mut buffer)).await;
+
+                    deactivate_on_err(
+                        &profile_path,
+                        deactivate_on_err(
+                            &profile_path,
+                            deactivate_on_err(
+                                &profile_path,
+                                timeout(Duration::from_secs(max_time as u64), stream.next()).await,
+                            )
+                            .await
+                            .ok_or("Watcher ended prematurely"),
+                        )
+                        .await,
+                    )
+                    .await;
+                });
+            })
+            .join()
+            .unwrap();
+
+            info!("Confirmation successful!");
+
+            std::process::exit(0);
+        }
+        fork::Fork::Parent(_) => {
+            std::process::exit(0);
+        }
+    }
+}
+
 pub async fn activate(
     profile_path: String,
     closure: String,
     bootstrap_cmd: Option<String>,
     auto_rollback: bool,
+    temp_path: String,
+    max_time: u16,
 ) -> Result<(), Box<dyn std::error::Error>> {
     info!("Activating profile");
 
@@ -83,80 +245,17 @@ pub async fn activate(
 
     match activate_status {
         Ok(s) if s.success() => (),
-        _ if auto_rollback => {
-            error!("Failed to execute activation command");
-
-            let nix_env_rollback_exit_status = Command::new("nix-env")
-                .arg("-p")
-                .arg(&profile_path)
-                .arg("--rollback")
-                .stdout(Stdio::null())
-                .stderr(Stdio::null())
-                .status()
-                .await?;
-
-            if !nix_env_rollback_exit_status.success() {
-                good_panic!("`nix-env --rollback` failed");
-            }
-
-            debug!("Listing generations");
-
-            let nix_env_list_generations_out = Command::new("nix-env")
-                .arg("-p")
-                .arg(&profile_path)
-                .arg("--list-generations")
-                .output()
-                .await?;
-
-            if !nix_env_list_generations_out.status.success() {
-                good_panic!("Listing `nix-env` generations failed");
-            }
-
-            let generations_list = String::from_utf8(nix_env_list_generations_out.stdout)?;
-
-            let last_generation_line = generations_list
-                .lines()
-                .last()
-                .expect("Expected to find a generation in list");
-
-            let last_generation_id = last_generation_line
-                .split_whitespace()
-                .next()
-                .expect("Expected to get ID from generation entry");
-
-            debug!("Removing generation entry {}", last_generation_line);
-            warn!("Removing generation by ID {}", last_generation_id);
-
-            let nix_env_delete_generation_exit_status = Command::new("nix-env")
-                .arg("-p")
-                .arg(&profile_path)
-                .arg("--delete-generations")
-                .arg(last_generation_id)
-                .stdout(Stdio::null())
-                .stderr(Stdio::null())
-                .status()
-                .await?;
-
-            if !nix_env_delete_generation_exit_status.success() {
-                good_panic!("Failed to delete failed generation");
-            }
-
-            info!("Attempting re-activate last generation");
-
-            let re_activate_exit_status = Command::new(format!("{}/deploy-rs-activate", profile_path))
-                .env("PROFILE", &profile_path)
-                .current_dir(&profile_path)
-                .status()
-                .await?;
+        _ if auto_rollback => return Ok(deactivate(&profile_path).await?),
+        _ => (),
+    }
 
-            if !re_activate_exit_status.success() {
-                good_panic!("Failed to re-activate the last generation");
-            }
+    info!("Activation succeeded, now performing post-activation checks");
 
-            std::process::exit(1);
-        }
-        _ => {}
-    }
+    deactivate_on_err(
+        &profile_path,
+        activation_confirmation(profile_path.clone(), temp_path, max_time, closure).await,
+    )
+    .await;
 
     Ok(())
 }
@@ -176,6 +275,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         opts.closure,
         opts.bootstrap_cmd,
         opts.auto_rollback,
+        opts.temp_path,
+        opts.max_time,
     )
     .await?;
 
diff --git a/src/utils/data.rs b/src/utils/data.rs
index f72f9a7..351b9ae 100644
--- a/src/utils/data.rs
+++ b/src/utils/data.rs
@@ -44,6 +44,10 @@ pub struct ProfileSettings {
     pub bootstrap: Option<String>,
     #[serde(rename(deserialize = "profilePath"))]
     pub profile_path: Option<String>,
+    #[serde(rename(deserialize = "maxTime"))]
+    pub max_time: Option<u16>,
+    #[serde(rename(deserialize = "tempPath"))]
+    pub temp_path: Option<String>,
 }
 
 #[derive(Deserialize, Debug, Clone)]
diff --git a/src/utils/deploy.rs b/src/utils/deploy.rs
index 7301967..e3493ba 100644
--- a/src/utils/deploy.rs
+++ b/src/utils/deploy.rs
@@ -2,6 +2,7 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
+use std::borrow::Cow;
 use tokio::process::Command;
 
 fn build_activate_command(
@@ -11,9 +12,13 @@ fn build_activate_command(
     closure: &str,
     bootstrap_cmd: &Option<String>,
     auto_rollback: bool,
+    temp_path: &Cow<str>,
+    max_time: u16,
 ) -> String {
-    let mut self_activate_command =
-        format!("{} '{}' '{}'", activate_path_str, profile_path, closure);
+    let mut self_activate_command = format!(
+        "{} '{}' '{}' {} {}",
+        activate_path_str, profile_path, closure, temp_path, max_time
+    );
 
     if let Some(sudo_cmd) = &sudo {
         self_activate_command = format!("{} {}", sudo_cmd, self_activate_command);
@@ -41,6 +46,8 @@ fn test_activation_command_builder() {
     let closure = "/blah/etc";
     let bootstrap_cmd = None;
     let auto_rollback = true;
+    let temp_path = &"/tmp/deploy-rs".into();
+    let max_time = 30;
 
     assert_eq!(
         build_activate_command(
@@ -50,8 +57,10 @@ fn test_activation_command_builder() {
             closure,
             &bootstrap_cmd,
             auto_rollback,
+            temp_path,
+            max_time
         ),
-        "sudo -u test /blah/bin/activate '/blah/profiles/test' '/blah/etc' --auto-rollback"
+        "sudo -u test /blah/bin/activate '/blah/profiles/test' '/blah/etc' /tmp/deploy-rs 30 --auto-rollback"
             .to_string(),
     );
 }
@@ -67,6 +76,13 @@ pub async fn deploy_profile(
 
     let activate_path_str = super::deploy_path_to_activate_path_str(&deploy_defs.current_exe)?;
 
+    let temp_path: Cow<str> = match &deploy_data.profile.profile_settings.temp_path {
+        Some(x) => x.into(),
+        None => "/tmp/deploy-rs".into(),
+    };
+
+    let max_time = deploy_data.profile.profile_settings.max_time.unwrap_or(30);
+
     let self_activate_command = build_activate_command(
         activate_path_str,
         &deploy_defs.sudo,
@@ -74,6 +90,8 @@ pub async fn deploy_profile(
         &deploy_data.profile.profile_settings.path,
         &deploy_data.profile.profile_settings.bootstrap,
         deploy_data.merged_settings.auto_rollback,
+        &temp_path,
+        max_time,
     );
 
     let hostname = match deploy_data.cmd_overrides.hostname {
@@ -94,5 +112,33 @@ pub async fn deploy_profile(
         good_panic!("Activation over SSH failed");
     }
 
+    info!("Success, attempting to connect to the node to confirm deployment");
+
+    let mut c = Command::new("ssh");
+    let mut ssh_confirm_command = c.arg(format!("ssh://{}@{}", deploy_defs.ssh_user, hostname));
+
+    for ssh_opt in &deploy_data.merged_settings.ssh_opts {
+        ssh_confirm_command = ssh_confirm_command.arg(ssh_opt);
+    }
+
+    let lock_hash = &deploy_data.profile.profile_settings.path[11 /* /nix/store/ */ ..];
+    let lock_path = format!("{}/activating-{}", temp_path, lock_hash);
+
+    let mut confirm_command = format!("rm {}", lock_path);
+    if let Some(sudo_cmd) = &deploy_defs.sudo {
+        confirm_command = format!("{} {}", sudo_cmd, confirm_command);
+    }
+
+    let ssh_exit_status = ssh_confirm_command.arg(confirm_command).status().await?;
+
+    if !ssh_exit_status.success() {
+        good_panic!(
+            "Failed to confirm deployment, the node will roll back in <{} seconds",
+            max_time
+        );
+    }
+
+    info!("Deployment confirmed.");
+
     Ok(())
 }