From 16ccdbc7000872d124c38f1688422db4b21741a2 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Wed, 28 Apr 2021 15:09:40 -0700
Subject: [PATCH] Better handling of killed processes

---
 src/nix/host/key_uploader.rs |  4 ++--
 src/nix/mod.rs               | 26 +++++++++++++++-----------
 src/nix/profile.rs           |  2 +-
 src/util.rs                  |  4 ++--
 4 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/src/nix/host/key_uploader.rs b/src/nix/host/key_uploader.rs
index 58e9a90..b3ce0c5 100644
--- a/src/nix/host/key_uploader.rs
+++ b/src/nix/host/key_uploader.rs
@@ -12,7 +12,7 @@ use shell_escape::unix::escape;
 use tokio::io::{AsyncWriteExt, BufReader};
 use tokio::process::Child;
 
-use crate::nix::{Key, NixError, NixResult};
+use crate::nix::{Key, NixResult};
 use crate::progress::TaskProgress;
 use crate::util::capture_stream;
 
@@ -54,6 +54,6 @@ pub async fn feed_uploader(mut uploader: Child, key: &Key, progress: TaskProgres
     if exit.success() {
         Ok(())
     } else {
-        Err(NixError::NixFailure { exit_code: exit.code().unwrap() })
+        Err(exit.into())
     }
 }
diff --git a/src/nix/mod.rs b/src/nix/mod.rs
index 4a1f88a..740699f 100644
--- a/src/nix/mod.rs
+++ b/src/nix/mod.rs
@@ -1,7 +1,8 @@
 use std::collections::HashMap;
 use std::convert::TryFrom;
+use std::os::unix::process::ExitStatusExt;
 use std::path::Path;
-use std::process::Stdio;
+use std::process::{ExitStatus, Stdio};
 
 use async_trait::async_trait;
 use serde::de::DeserializeOwned;
@@ -50,8 +51,8 @@ pub enum NixError {
     #[snafu(display("Nix exited with error code: {}", exit_code))]
     NixFailure { exit_code: i32 },
 
-    #[snafu(display("Nix was interrupted"))]
-    NixKilled,
+    #[snafu(display("Nix was killed by signal {}", signal))]
+    NixKilled { signal: i32 },
 
     #[snafu(display("This operation is not supported"))]
     Unsupported,
@@ -93,6 +94,15 @@ impl From<ValidationErrors> for NixError {
     }
 }
 
+impl From<ExitStatus> for NixError {
+    fn from(status: ExitStatus) -> Self {
+        match status.code() {
+            Some(exit_code) => Self::NixFailure { exit_code },
+            None => Self::NixKilled { signal: status.signal().unwrap() },
+        }
+    }
+}
+
 #[derive(Debug, Clone, Validate, Deserialize)]
 pub struct NodeConfig {
     #[serde(rename = "targetHost")]
@@ -152,10 +162,7 @@ impl NixCommand for Command {
         if exit.success() {
             Ok(())
         } else {
-            Err(match exit.code() {
-                Some(exit_code) => NixError::NixFailure { exit_code },
-                None => NixError::NixKilled,
-            })
+            Err(exit.into())
         }
     }
 
@@ -173,10 +180,7 @@ impl NixCommand for Command {
             // FIXME: unwrap
             Ok(String::from_utf8(output.stdout).unwrap())
         } else {
-            Err(match output.status.code() {
-                Some(exit_code) => NixError::NixFailure { exit_code },
-                None => NixError::NixKilled,
-            })
+            Err(output.status.into())
         }
     }
 
diff --git a/src/nix/profile.rs b/src/nix/profile.rs
index 823acac..769f5e0 100644
--- a/src/nix/profile.rs
+++ b/src/nix/profile.rs
@@ -131,7 +131,7 @@ impl ProfileMap {
 
             let status = command.status().await?;
             if !status.success() {
-                return Err(NixError::NixFailure { exit_code: status.code().unwrap() });
+                return Err(status.into());
             }
         }
 
diff --git a/src/util.rs b/src/util.rs
index fec6451..3a13a31 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -9,7 +9,7 @@ use glob::Pattern as GlobPattern;
 use tokio::io::{AsyncRead, AsyncBufReadExt, BufReader};
 use tokio::process::Command;
 
-use super::nix::{NodeConfig, Hive, NixResult, NixError};
+use super::nix::{NodeConfig, Hive, NixResult};
 use super::progress::TaskProgress;
 
 enum NodeFilter {
@@ -74,7 +74,7 @@ impl CommandExecution {
         if exit.success() {
             Ok(())
         } else {
-            Err(NixError::NixFailure { exit_code: exit.code().unwrap() })
+            Err(exit.into())
         }
     }
 }