feat: use process groups to improve zombie job handling (#6157)

* useProcessGroups

* nit

backend/Cargo.lock

@@ -8126,6 +8126,18 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "nix"
+version = "0.30.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6"
+dependencies = [
+ "bitflags 2.9.1",
+ "cfg-if",
+ "cfg_aliases 0.2.1",
+ "libc",
+]
+
 [[package]]
 name = "nkeys"
 version = "0.4.5"
@@ -9593,6 +9605,20 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "process-wrap"
+version = "8.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3ef4f2f0422f23a82ec9f628ea2acd12871c81a9362b02c43c1aa86acfc3ba1"
+dependencies = [
+ "futures",
+ "indexmap 2.10.0",
+ "nix 0.30.1",
+ "tokio",
+ "tracing",
+ "windows 0.61.3",
+]
+
 [[package]]
 name = "procfs"
 version = "0.17.0"
@@ -15271,6 +15297,7 @@ dependencies = [
  "pem 3.0.5",
  "pep440_rs",
  "postgres-native-tls 0.5.1",
+ "process-wrap",
  "prometheus",
  "rand 0.9.0",
  "regex",

backend/Cargo.toml

@@ -347,6 +347,8 @@ async-nats = "0.38.0"
 nkeys = "0.4.4"
 nu-parser = { version = "0.101.0", default-features = false }
 
+process-wrap = { version = "8.2.1", features = ["tokio1"] }
+
 datafusion = "47.0.0"
 object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "36752c975d4f29e20b57c91f81a10872dcd48ae7", features = ["aws", "azure", "gcp"] } 
 openidconnect = { version = "4.0.0-rc.1" }
@@ -402,3 +404,4 @@ oracle = { version = "0.6.3", features = ["chrono"] }
 rumqttc = { version = "0.24.0", features = ["use-native-tls"]}
 strum = { version = "0.27", features = ["derive"] }
 strum_macros = "^0"
+

backend/windmill-worker/Cargo.toml

@@ -120,6 +120,7 @@ yaml-rust.workspace = true
 backon.workspace = true
 winapi = { workspace = true, optional = true }
 pep440_rs.workspace = true
+process-wrap.workspace = true
 
 opentelemetry = { workspace = true, optional = true }
 bollard = { workspace = true, optional = true }

backend/windmill-worker/src/bash_executor.rs

@@ -618,11 +618,12 @@ pub async fn handle_powershell_job(
                     .collect::<Vec<_>>()
                     .join(", "),
             );
-        let child = Command::new(POWERSHELL_PATH.as_str())
-            .args(&["-Command", &install_string])
+        let mut cmd = Command::new(POWERSHELL_PATH.as_str());
+        cmd.args(&["-Command", &install_string])
             .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .spawn()?;
+            .stderr(Stdio::piped());
+
+        let child = start_child_process(cmd, POWERSHELL_PATH.as_str()).await?;
 
         handle_child(
             &job.id,
@@ -750,8 +751,8 @@ $env:PSModulePath = \"{};$PSModulePathBackup\"",
             "wrapper.sh",
         ];
         cmd_args.extend(pwsh_args.iter().map(|x| x.as_str()));
-        Command::new(NSJAIL_PATH.as_str())
-            .current_dir(job_dir)
+        let mut cmd = Command::new(NSJAIL_PATH.as_str());
+        cmd.current_dir(job_dir)
             .env_clear()
             .envs(PROXY_ENVS.clone())
             .envs(reserved_variables)
@@ -760,8 +761,9 @@ $env:PSModulePath = \"{};$PSModulePathBackup\"",
             .env("BASE_INTERNAL_URL", base_internal_url)
             .args(cmd_args)
             .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .spawn()?
+            .stderr(Stdio::piped());
+
+        start_child_process(cmd, NSJAIL_PATH.as_str()).await?
     } else {
         let mut cmd;
         let mut cmd_args;
@@ -833,7 +835,7 @@ $env:PSModulePath = \"{};$PSModulePathBackup\"",
                 .env("USERPROFILE", crate::USERPROFILE_ENV.as_str());
         }
 
-        cmd.spawn()?
+        start_child_process(cmd, POWERSHELL_PATH.as_str()).await?
     };
 
     handle_child(

backend/windmill-worker/src/bun_executor.rs

@@ -171,7 +171,7 @@ pub async fn gen_bun_lockfile(
             )
             .await?;
         } else {
-            child_process.wait().await?;
+            Box::into_pin(child_process.wait()).await?;
         }
 
         let new_package_json = read_file_content(&format!("{job_dir}/package.json")).await?;
@@ -366,7 +366,7 @@ pub async fn install_bun_lockfile(
         )
         .await?
     } else {
-        child_process.wait().await?;
+        Box::into_pin(child_process.wait()).await?;
     }
 
     if has_file {
@@ -589,7 +589,7 @@ pub async fn generate_bun_bundle(
         )
         .await?;
     } else {
-        child_process.wait().await?;
+        Box::into_pin(child_process.wait()).await?;
     }
     Ok(())
 }

backend/windmill-worker/src/common.rs

@@ -2,6 +2,7 @@ use async_recursion::async_recursion;
 
 use itertools::Itertools;
 use lazy_static::lazy_static;
+use process_wrap::tokio::TokioChildWrapper;
 use regex::Regex;
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
@@ -41,7 +42,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 use uuid::Uuid;
 use windmill_common::{variables, DB};
 
-use tokio::{io::AsyncWriteExt, process::Child, time::Instant};
+use tokio::{io::AsyncWriteExt, time::Instant};
 
 use crate::agent_workers::UPDATE_PING_URL;
 use crate::{DISABLE_NSJAIL, JOB_DEFAULT_TIMEOUT, MAX_RESULT_SIZE, MAX_TIMEOUT_DURATION, PATH_ENV};
@@ -611,7 +612,22 @@ impl OccupancyMetrics {
     }
 }
 
-pub async fn start_child_process(mut cmd: Command, executable: &str) -> Result<Child, Error> {
+pub async fn start_child_process(
+    cmd: Command,
+    executable: &str,
+) -> Result<Box<dyn TokioChildWrapper>, Error> {
+    use process_wrap::tokio::*;
+    let mut cmd = TokioCommandWrap::from(cmd);
+    #[cfg(unix)]
+    {
+        use process_wrap::tokio::ProcessGroup;
+
+        cmd.wrap(ProcessGroup::leader());
+    }
+    #[cfg(windows)]
+    {
+        cmd.wrap(JobObject);
+    }
     return cmd
         .spawn()
         .map_err(|err| tentatively_improve_error(err.into(), executable));

backend/windmill-worker/src/dedicated_worker.rs

@@ -102,12 +102,12 @@ pub async fn handle_dedicated_process(
     };
 
     let stdout = child
-        .stdout
+        .stdout()
         .take()
         .expect("child did not have a handle to stdout");
 
     let stderr = child
-        .stderr
+        .stderr()
         .take()
         .expect("child did not have a handle to stderr");
 
@@ -116,15 +116,14 @@ pub async fn handle_dedicated_process(
     let mut err_reader = BufReader::new(stderr).lines();
 
     let mut stdin = child
-        .stdin
+        .stdin()
         .take()
         .expect("child did not have a handle to stdin");
 
     // Ensure the child process is spawned in the runtime so it can
     // make progress on its own while we await for any output.
     let child = tokio::spawn(async move {
-        let status = child
-            .wait()
+        let status = Box::into_pin(child.wait())
             .await
             .expect("child process encountered an error");
         if let Err(e) = process_status(&cmd_name, status) {

backend/windmill-worker/src/deno_executor.rs

@@ -164,7 +164,7 @@ pub async fn generate_deno_lock(
         )
         .await?;
     } else {
-        child_process.wait().await?;
+        Box::into_pin(child_process.wait()).await?;
     }
 
     let path_lock = format!("{job_dir}/lock.json");

backend/windmill-worker/src/handle_child.rs

@@ -4,6 +4,7 @@ use futures::Future;
 use nix::sys::signal::{self, Signal};
 #[cfg(any(target_os = "linux", target_os = "macos"))]
 use nix::unistd::Pid;
+use process_wrap::tokio::TokioChildWrapper;
 use windmill_common::agent_workers::PingJobStatusResponse;
 use windmill_common::jobs::LARGE_LOG_THRESHOLD_SIZE;
 
@@ -41,7 +42,6 @@ use windmill_common::job_metrics;
 use tokio::io::AsyncWriteExt;
 use tokio::{
     io::{AsyncBufReadExt, BufReader},
-    process::Child,
     sync::{broadcast, watch},
     time::{interval, sleep, Instant, MissedTickBehavior},
 };
@@ -99,7 +99,7 @@ pub async fn handle_child(
     conn: &Connection,
     mem_peak: &mut i32,
     canceled_by_ref: &mut Option<CanceledBy>,
-    mut child: Child,
+    mut child: Box<dyn TokioChildWrapper>,
     nsjail: bool,
     worker: &str,
     w_id: &str,
@@ -197,7 +197,7 @@ pub async fn handle_child(
     let wait_on_child = async {
         let kill_reason = tokio::select! {
             biased;
-            result = child.wait() => return result.map(Ok),
+            result = Box::into_pin(child.wait()) => return result.map(Ok),
             Ok(()) = too_many_logs.changed() => KillReason::TooManyLogs,
             _ = sleep(timeout_duration) => KillReason::Timeout { is_job_specific },
             ex = update_job, if job_id != Uuid::nil() => match ex {
@@ -291,7 +291,7 @@ pub async fn handle_child(
         #[cfg(unix)]
         {
             /* send SIGKILL and reap child process */
-            let (_, kill) = future::join(set_reason, child.kill()).await;
+            let (_, kill) = future::join(set_reason, Box::into_pin(child.kill())).await;
             kill.map(|()| Err(kill_reason))
         }
     };
@@ -728,17 +728,17 @@ where
 ///
 /// builds a stream joining both stdout and stderr each read line by line
 fn child_joined_output_stream(
-    child: &mut Child,
+    child: &mut Box<dyn TokioChildWrapper>,
     job_id: Uuid,
     w_id: String,
 ) -> impl stream::FusedStream<Item = io::Result<String>> {
     let stderr = child
-        .stderr
+        .stderr()
         .take()
-        .expect("child did not have a handle to stdout");
+        .expect("child did not have a handle to stderr");
 
     let stdout = child
-        .stdout
+        .stdout()
         .take()
         .expect("child did not have a handle to stdout");
 

backend/windmill-worker/src/python_executor.rs

@@ -38,6 +38,8 @@ use windmill_common::variables::get_secret_value_as_admin;
 use std::env::var;
 use windmill_queue::{append_logs, CanceledBy, PrecomputedAgentInfo};
 
+use process_wrap::tokio::TokioChildWrapper;
+
 lazy_static::lazy_static! {
     pub(crate) static ref PYTHON_PATH: Option<String> = var("PYTHON_PATH").ok().map(|v| {
         tracing::warn!("PYTHON_PATH is set to {} and thus python will not be managed by uv and stay static regardless of annotation and instance settings. NOT RECOMMENDED", v);
@@ -88,7 +90,8 @@ async fn handle_piptar_uploads(mut rx: tokio::sync::mpsc::UnboundedReceiver<Pipt
 
     while let Some(task) = rx.recv().await {
         if let Some(os) = get_object_store().await {
-            match build_tar_and_push(os, task.venv_path.clone(), task.cache_dir, None, false).await {
+            match build_tar_and_push(os, task.venv_path.clone(), task.cache_dir, None, false).await
+            {
                 Ok(()) => {
                     tracing::info!("Successfully uploaded piptar for {}", task.venv_path);
                 }
@@ -97,7 +100,10 @@ async fn handle_piptar_uploads(mut rx: tokio::sync::mpsc::UnboundedReceiver<Pipt
                 }
             }
         } else {
-            tracing::warn!("S3 object store not available for piptar upload: {}", task.venv_path);
+            tracing::warn!(
+                "S3 object store not available for piptar upload: {}",
+                task.venv_path
+            );
         }
     }
 }
@@ -336,26 +342,30 @@ pub async fn uv_pip_compile(
                 )
                 .env(
                     "APPDATA",
-                    std::env::var("APPDATA")
-                        .unwrap_or_else(|_| format!("{}\\AppData\\Roaming", crate::USERPROFILE_ENV.as_str())),
+                    std::env::var("APPDATA").unwrap_or_else(|_| {
+                        format!("{}\\AppData\\Roaming", crate::USERPROFILE_ENV.as_str())
+                    }),
                 )
                 .env(
                     "ComSpec",
-                    std::env::var("ComSpec").unwrap_or_else(|_| String::from("C:\\Windows\\System32\\cmd.exe")),
+                    std::env::var("ComSpec")
+                        .unwrap_or_else(|_| String::from("C:\\Windows\\System32\\cmd.exe")),
                 )
                 .env(
                     "PATHEXT",
-                    std::env::var("PATHEXT").unwrap_or_else(|_|
+                    std::env::var("PATHEXT").unwrap_or_else(|_| {
                         String::from(".COM;.EXE;.BAT;.CMD;.VBS;.VBE;.JS;.JSE;.WSF;.WSH;.MSC;.CPL")
-                    ),
+                    }),
                 )
                 .env(
                     "ProgramData",
-                    std::env::var("ProgramData").unwrap_or_else(|_| String::from("C:\\ProgramData")),
+                    std::env::var("ProgramData")
+                        .unwrap_or_else(|_| String::from("C:\\ProgramData")),
                 )
                 .env(
                     "ProgramFiles",
-                    std::env::var("ProgramFiles").unwrap_or_else(|_| String::from("C:\\Program Files")),
+                    std::env::var("ProgramFiles")
+                        .unwrap_or_else(|_| String::from("C:\\Program Files")),
                 );
         }
 
@@ -1290,7 +1300,7 @@ async fn spawn_uv_install(
     // If none, it is system python
     py_path: Option<String>,
     worker_dir: &str,
-) -> Result<tokio::process::Child, Error> {
+) -> Result<Box<dyn TokioChildWrapper>, Error> {
     if !*DISABLE_NSJAIL {
         tracing::info!(
             workspace_id = %w_id,
@@ -1453,26 +1463,30 @@ async fn spawn_uv_install(
                 )
                 .env(
                     "APPDATA",
-                    std::env::var("APPDATA")
-                        .unwrap_or_else(|_| format!("{}\\AppData\\Roaming", crate::USERPROFILE_ENV.as_str())),
+                    std::env::var("APPDATA").unwrap_or_else(|_| {
+                        format!("{}\\AppData\\Roaming", crate::USERPROFILE_ENV.as_str())
+                    }),
                 )
                 .env(
                     "ComSpec",
-                    std::env::var("ComSpec").unwrap_or_else(|_| String::from("C:\\Windows\\System32\\cmd.exe")),
+                    std::env::var("ComSpec")
+                        .unwrap_or_else(|_| String::from("C:\\Windows\\System32\\cmd.exe")),
                 )
                 .env(
                     "PATHEXT",
-                    std::env::var("PATHEXT").unwrap_or_else(|_|
+                    std::env::var("PATHEXT").unwrap_or_else(|_| {
                         String::from(".COM;.EXE;.BAT;.CMD;.VBS;.VBE;.JS;.JSE;.WSF;.WSH;.MSC;.CPL")
-                    ),
+                    }),
                 )
                 .env(
                     "ProgramData",
-                    std::env::var("ProgramData").unwrap_or_else(|_| String::from("C:\\ProgramData")),
+                    std::env::var("ProgramData")
+                        .unwrap_or_else(|_| String::from("C:\\ProgramData")),
                 )
                 .env(
                     "ProgramFiles",
-                    std::env::var("ProgramFiles").unwrap_or_else(|_| String::from("C:\\Program Files")),
+                    std::env::var("ProgramFiles")
+                        .unwrap_or_else(|_| String::from("C:\\Program Files")),
                 )
                 .args(&command_args[1..])
                 .stdout(Stdio::piped())
@@ -1901,7 +1915,7 @@ pub async fn handle_python_reqs(
 
             let mut stderr_buf = String::new();
             let mut stderr_pipe = uv_install_proccess
-                .stderr
+                .stderr()
                 .take()
                 .ok_or(anyhow!("Cannot take stderr from uv_install_proccess"))?;
             let stderr_future = stderr_pipe.read_to_string(&mut stderr_buf);
@@ -1918,14 +1932,14 @@ pub async fn handle_python_reqs(
             tokio::select! {
                 // Canceled
                 _ = kill_rx.recv() => {
-                    uv_install_proccess.kill().await?;
+                    Box::into_pin(uv_install_proccess.kill()).await?;
                     pids.lock().await.get_mut(i).and_then(|e| e.take());
                     return Err(anyhow::anyhow!("uv pip install was canceled"));
                 },
                 (_, exitstatus) = async {
                     // See tokio::process::Child::wait_with_output() for more context
                     // Sometimes uv_install_proccess.wait() is not exiting if stderr is not awaited before it :/
-                    (stderr_future.await, uv_install_proccess.wait().await)
+                    (stderr_future.await, Box::into_pin(uv_install_proccess.wait()).await)
                 } => match exitstatus {
                     Ok(status) => if !status.success() {
                         tracing::warn!(
@@ -2294,4 +2308,3 @@ for line in sys.stdin:
     )
     .await
 }
-