process: fix cleanup race (#818)

Summary:
Pull Request resolved: #818

remove `.kill_on_drop()` introduced by D79898223. fix the underlying race that necessitated it and kill child processes via a `Drop` implementation for `struct Child`. fix provided by mariusae

Reviewed By: mariusae

Differential Revision: D80022315

fbshipit-source-id: 1757e1eaf830cf50989173142e5c20c451fe7e70

Author: shayne-fletcher

hyperactor_mesh/src/alloc/process.rs

@@ -135,9 +135,10 @@ struct Child {
     channel: ChannelState,
     group: monitor::Group,
     exit_flag: Option<flag::Flag>,
-    stdout: LogTailer,
-    stderr: LogTailer,
+    stdout: Option<LogTailer>,
+    stderr: Option<LogTailer>,
     stop_reason: Arc<OnceLock<ProcStopReason>>,
+    process_pid: Arc<std::sync::Mutex<Option<i32>>>,
 }
 
 impl Child {
@@ -179,27 +180,23 @@ impl Child {
             stderr_tee,
         );
 
+        let process_pid = Arc::new(std::sync::Mutex::new(Some(process.id().unwrap() as i32)));
+
         let child = Self {
             local_rank,
             channel: ChannelState::NotConnected,
             group,
             exit_flag: Some(exit_flag),
-            stdout,
-            stderr,
+            stdout: Some(stdout),
+            stderr: Some(stderr),
             stop_reason: Arc::clone(&stop_reason),
+            process_pid: process_pid.clone(),
         };
 
         let monitor = async move {
             let reason = tokio::select! {
                 _ = handle => {
-                    let Some(id) = process.id() else {
-                        tracing::error!("could not get child process id");
-                        return ProcStopReason::Unknown;
-                    };
-                    if let Err(e) = signal::kill(Pid::from_raw(id as i32), signal::SIGTERM) {
-                        tracing::error!("failed to kill child process: {}", e);
-                        return ProcStopReason::Unknown;
-                    };
+                    Self::ensure_killed(process_pid);
                     Self::exit_status_to_reason(process.wait().await)
                 }
                 result = process.wait() => {
@@ -214,6 +211,25 @@ impl Child {
         (child, monitor)
     }
 
+    fn ensure_killed(pid: Arc<std::sync::Mutex<Option<i32>>>) {
+        match pid.lock().unwrap().take() {
+            Some(pid) => {
+                if let Err(e) = signal::kill(Pid::from_raw(pid), signal::SIGTERM) {
+                    match e {
+                        nix::errno::Errno::ESRCH => {
+                            // Process already gone.
+                            tracing::debug!("pid {} already exited", pid);
+                        }
+                        _ => {
+                            tracing::error!("failed to kill {}: {}", pid, e);
+                        }
+                    }
+                }
+            }
+            None => (),
+        }
+    }
+
     fn exit_status_to_reason(result: io::Result<ExitStatus>) -> ProcStopReason {
         match result {
             Ok(status) if status.success() => ProcStopReason::Stopped,
@@ -300,6 +316,12 @@ impl Child {
     }
 }
 
+impl Drop for Child {
+    fn drop(&mut self) {
+        Self::ensure_killed(self.process_pid.clone());
+    }
+}
+
 impl ProcessAlloc {
     // Also implement exit (for graceful exit)
 
@@ -371,7 +393,6 @@ impl ProcessAlloc {
         cmd.env(bootstrap::BOOTSTRAP_LOG_CHANNEL, log_channel.to_string());
         cmd.stdout(Stdio::piped());
         cmd.stderr(Stdio::piped());
-        cmd.kill_on_drop(true);
 
         let proc_id = ProcId::Ranked(WorldId(self.name.to_string()), index);
         tracing::debug!("Spawning process {:?}", cmd);
@@ -472,7 +493,9 @@ impl Alloc for ProcessAlloc {
                 },
 
                 Some(Ok((index, mut reason))) = self.children.join_next() => {
-                    let stderr_content = if let Some(Child { stdout, stderr, ..} ) = self.remove(index) {
+                    let stderr_content =  if let Some(mut child) = self.remove(index) {
+                        let stdout = child.stdout.take().unwrap();
+                        let stderr = child.stderr.take().unwrap();
                         stdout.abort();
                         stderr.abort();
                         let (_stdout, _) = stdout.join().await;