From d2a73beb43c68b8cf43ddf7edc9ff17f05372a27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bartosz=20R=C3=B3=C5=BCa=C5=84ski?= <bartek.roza@gmail.com>
Date: Thu, 26 Oct 2023 12:22:51 +0200
Subject: [PATCH] Add mode for watching PID and exit if it dies

---
 Cargo.lock          | 25 +++++++++++++++
 service/Cargo.toml  |  1 +
 service/src/main.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 49dfcc3d..f14171b8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1535,6 +1535,15 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "61807f77802ff30975e01f4f071c8ba10c022052f98b3294119f3e615d13e5be"
 
+[[package]]
+name = "ntapi"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "num-complex"
 version = "0.4.4"
@@ -2503,6 +2512,7 @@ dependencies = [
  "prost",
  "rcgen",
  "rstest 0.18.2",
+ "sysinfo",
  "tempfile",
  "tokio",
  "tokio-stream",
@@ -2638,6 +2648,21 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160"
 
+[[package]]
+name = "sysinfo"
+version = "0.29.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a18d114d420ada3a891e6bc8e96a2023402203296a47cdd65083377dad18ba5"
+dependencies = [
+ "cfg-if",
+ "core-foundation-sys",
+ "libc",
+ "ntapi",
+ "once_cell",
+ "rayon",
+ "winapi",
+]
+
 [[package]]
 name = "tap"
 version = "1.0.1"
diff --git a/service/Cargo.toml b/service/Cargo.toml
index 6bbbd05d..b36b568a 100644
--- a/service/Cargo.toml
+++ b/service/Cargo.toml
@@ -25,6 +25,7 @@ env_logger = "0.10.0"
 clap = { version = "4.4.4", features = ["derive"] }
 hex = "0.4.3"
 mockall = "0.11.4"
+sysinfo = "0.29.10"
 
 [build-dependencies]
 tonic-build = "0.10.0"
diff --git a/service/src/main.rs b/service/src/main.rs
index 75cf63ac..b9ceca3a 100644
--- a/service/src/main.rs
+++ b/service/src/main.rs
@@ -2,6 +2,7 @@ use std::{fs::read_to_string, path::PathBuf, time::Duration};
 
 use clap::{Args, Parser, ValueEnum};
 use eyre::Context;
+use sysinfo::{Pid, ProcessExt, ProcessStatus, System, SystemExt};
 use tonic::transport::{Certificate, Identity};
 
 use post::pow::randomx::RandomXFlag;
@@ -29,6 +30,10 @@ struct Cli {
 
     #[command(flatten, next_help_heading = "TLS configuration")]
     tls: Option<Tls>,
+
+    /// watch PID and exit if it dies
+    #[arg(long)]
+    watch_pid: Option<sysinfo::Pid>,
 }
 
 #[derive(Args, Debug)]
@@ -200,6 +205,73 @@ async fn main() -> eyre::Result<()> {
     };
 
     let client = client::ServiceClient::new(args.address, args.reconnect_interval_s, tls, service)?;
+    let client_handle = tokio::spawn(client.run());
+
+    if let Some(pid) = args.watch_pid {
+        tokio::task::spawn_blocking(move || watch_pid(pid, Duration::from_secs(1))).await?;
+        Ok(())
+    } else {
+        client_handle.await?
+    }
+}
+
+// watch given PID and return when it dies
+fn watch_pid(pid: Pid, interval: Duration) {
+    log::info!("watching PID {pid}");
+
+    let mut sys = System::new();
+    while sys.refresh_process(pid) {
+        if let Some(p) = sys.process(pid) {
+            match p.status() {
+                ProcessStatus::Zombie | ProcessStatus::Dead => {
+                    break;
+                }
+                _ => {}
+            }
+        }
+        std::thread::sleep(interval);
+    }
 
-    client.run().await
+    log::info!("PID {pid} died");
+}
+
+#[cfg(test)]
+mod tests {
+    use std::process::Command;
+
+    use sysinfo::PidExt;
+
+    #[tokio::test]
+    async fn watching_pid_zombie() {
+        let mut proc = Command::new("sleep").arg("99999").spawn().unwrap();
+        let pid = proc.id();
+
+        let handle = tokio::task::spawn_blocking(move || {
+            super::watch_pid(
+                sysinfo::Pid::from_u32(pid),
+                std::time::Duration::from_millis(10),
+            )
+        });
+        // just kill leaves a zombie process
+        proc.kill().unwrap();
+        handle.await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn watching_pid_reaped() {
+        let mut proc = Command::new("sleep").arg("99999").spawn().unwrap();
+        let pid = proc.id();
+
+        let handle = tokio::task::spawn_blocking(move || {
+            super::watch_pid(
+                sysinfo::Pid::from_u32(pid),
+                std::time::Duration::from_millis(10),
+            )
+        });
+
+        // kill and wait
+        proc.kill().unwrap();
+        proc.wait().unwrap();
+        handle.await.unwrap();
+    }
 }