Fix terminating watch_pid task

Cargo.toml

@@ -11,7 +11,7 @@ members = [
 
 [package]
 name = "post-rs"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 [lib]

certifier/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "certifier"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 [dependencies]

ffi/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "post-cbindings"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 

initializer/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "initializer"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

profiler/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "profiler"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 [dependencies]

scrypt-ocl/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "scrypt-ocl"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 [dependencies]

service/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "service"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 
 [lib]

service/src/main.rs

@@ -3,6 +3,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 tokio::sync::oneshot::{self, error::TryRecvError, Receiver};
 use tonic::transport::{Certificate, Identity};
 
 use post::pow::randomx::RandomXFlag;
@@ -226,60 +227,71 @@ async fn main() -> eyre::Result<()> {
     let client = client::ServiceClient::new(args.address, tls, service)?;
     let client_handle = tokio::spawn(client.run(args.max_retries, args.reconnect_interval_s));
 
+    // A channel to communicate when the blocking task should quit.
+    let (term_tx, term_rx) = oneshot::channel();
+
     tokio::select! {
-        Some(err) = watch_pid_if_needed(args.watch_pid) => {
+        Some(err) = watch_pid_if_needed(args.watch_pid.map(|p| (p, term_rx))) => {
             log::info!("PID watcher exited: {err:?}");
             return Ok(())
         }
         err = client_handle => {
+            drop(term_tx);
             return err.unwrap();
         }
     }
 }
 
 async fn watch_pid_if_needed(
-    pid: Option<Pid>,
+    watch: Option<(Pid, Receiver<()>)>,
 ) -> Option<std::result::Result<(), tokio::task::JoinError>> {
-    match pid {
-        Some(pid) => {
-            Some(tokio::task::spawn_blocking(move || watch_pid(pid, Duration::from_secs(1))).await)
-        }
+    match watch {
+        Some((pid, term)) => Some(
+            tokio::task::spawn_blocking(move || watch_pid(pid, Duration::from_secs(1), term)).await,
+        ),
         None => None,
     }
 }
 
 // watch given PID and return when it dies
-fn watch_pid(pid: Pid, interval: Duration) {
+fn watch_pid(pid: Pid, interval: Duration, mut term: Receiver<()>) {
     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;
+                    log::info!("PID {pid} died");
+                    return;
                 }
                 _ => {}
             }
         }
-        std::thread::sleep(interval);
+        match term.try_recv() {
+            Ok(_) | Err(TryRecvError::Closed) => {
+                log::debug!("PID watcher received termination signal");
+                return;
+            }
+            _ => std::thread::sleep(interval),
+        }
     }
-
-    log::info!("PID {pid} died");
 }
 
 #[cfg(test)]
 mod tests {
     use std::process::Command;
 
     use sysinfo::{Pid, PidExt};
+    use tokio::sync::oneshot;
 
     #[tokio::test]
     async fn watch_pid_if_needed() {
         // Don't watch
         assert!(super::watch_pid_if_needed(None).await.is_none());
         // Watch
-        super::watch_pid_if_needed(Some(Pid::from(0)))
+        let (_term_tx, term_rx) = oneshot::channel();
+        super::watch_pid_if_needed(Some((Pid::from(0), term_rx)))
             .await
             .expect("should be some")
             .expect("should be OK");
@@ -289,11 +301,12 @@ mod tests {
     async fn watching_pid_zombie() {
         let mut proc = Command::new("sleep").arg("99999").spawn().unwrap();
         let pid = proc.id();
-
+        let (_term_tx, term_rx) = oneshot::channel();
         let handle = tokio::task::spawn_blocking(move || {
             super::watch_pid(
                 sysinfo::Pid::from_u32(pid),
                 std::time::Duration::from_millis(10),
+                term_rx,
             )
         });
         // just kill leaves a zombie process
@@ -305,11 +318,13 @@ mod tests {
     async fn watching_pid_reaped() {
         let mut proc = Command::new("sleep").arg("99999").spawn().unwrap();
         let pid = proc.id();
+        let (_term_tx, term_rx) = oneshot::channel();
 
         let handle = tokio::task::spawn_blocking(move || {
             super::watch_pid(
                 sysinfo::Pid::from_u32(pid),
                 std::time::Duration::from_millis(10),
+                term_rx,
             )
         });
 
@@ -318,4 +333,36 @@ mod tests {
         proc.wait().unwrap();
         handle.await.unwrap();
     }
+
+    #[tokio::test]
+    async fn terminate_watching_pid() {
+        let mut proc = Command::new("sleep").arg("99999").spawn().unwrap();
+        let pid = proc.id();
+        let (term_tx, term_rx) = oneshot::channel();
+        let handle = tokio::task::spawn_blocking(move || {
+            super::watch_pid(
+                sysinfo::Pid::from_u32(pid),
+                std::time::Duration::from_millis(10),
+                term_rx,
+            )
+        });
+        // Terminate by closing the channel
+        drop(term_tx);
+        handle.await.unwrap();
+
+        let (term_tx, term_rx) = oneshot::channel();
+        let handle = tokio::task::spawn_blocking(move || {
+            super::watch_pid(
+                sysinfo::Pid::from_u32(pid),
+                std::time::Duration::from_millis(10),
+                term_rx,
+            )
+        });
+        // Terminate by sending a signal
+        term_tx.send(()).unwrap();
+        handle.await.unwrap();
+
+        proc.kill().unwrap();
+        proc.wait().unwrap();
+    }
 }