fix: windows process termination using Job Objects (NR-490832) (#2052)

* fix: Windows process termination using Job Objects

* nit

Author: gsanchezgavier

agent-control/Cargo.toml

@@ -85,6 +85,8 @@ nix = { workspace = true, features = ["signal", "user", "hostname"] }
 windows = { workspace = true, features = [
   "Win32_System_Console",
   "Win32_System_Threading",
+  "Win32_System_JobObjects",
+  "Win32_Security",
 ] }
 windows-service = "0.8.0"
 

agent-control/src/sub_agent/on_host/command.rs

@@ -1,5 +1,7 @@
 pub mod command_os;
 pub mod error;
 pub mod executable_data;
+#[cfg(target_family = "windows")]
+pub mod job_object;
 pub mod logging;
 pub mod restart_policy;

agent-control/src/sub_agent/on_host/command/command_os.rs

@@ -5,7 +5,6 @@ use crate::agent_control::defaults::{STDERR_LOG_PREFIX, STDOUT_LOG_PREFIX};
 use crate::sub_agent::on_host::command::executable_data::ExecutableData;
 use std::time::{Duration, Instant};
 use std::{
-    io,
     path::PathBuf,
     process::{Child, Command, ExitStatus, Stdio},
 };
@@ -18,6 +17,8 @@ use super::{
         logger::Logger,
     },
 };
+#[cfg(target_family = "windows")]
+use crate::sub_agent::on_host::command::job_object::JobObject;
 
 const POLL_INTERVAL: Duration = Duration::from_millis(100);
 
@@ -36,6 +37,9 @@ pub struct CommandOSStarted {
     process: Child,
     loggers: Option<FileSystemLoggers>,
     shutdown_timeout: Duration,
+
+    #[cfg(target_family = "windows")]
+    job_object: Option<JobObject>,
 }
 
 ////////////////////////////////////////////////////////////////////////////////////
@@ -54,9 +58,6 @@ impl CommandOSNotStarted {
             .stdout(Stdio::piped())
             .stderr(Stdio::piped());
 
-        #[cfg(target_family = "windows")]
-        Self::create_process_group(&mut cmd);
-
         Self {
             agent_id,
             cmd,
@@ -74,32 +75,31 @@ impl CommandOSNotStarted {
                 file_logger(&agent_id, self.logging_path, STDERR_LOG_PREFIX),
             )
         });
-        Ok(CommandOSStarted {
-            agent_id,
-            process: self.cmd.spawn()?,
-            loggers,
-            shutdown_timeout: self.shutdown_timeout,
-        })
-    }
+        let child = self.cmd.spawn()?;
 
-    #[cfg(target_family = "windows")]
-    /// Sets the process creation flags to create a new process group for Windows processes.
-    ///
-    /// This enables sending CTRL+BREAK events to it via the [`GenerateConsoleCtrlEvent`](windows::Win32::System::Console::GenerateConsoleCtrlEvent) function,
-    /// which is the mechanism we use to gracefully shut down the process. Otherwise, the Agent Control process needs to attach itself to the
-    /// console of the process to send a CTRL+C event which would need synchronization (many sub-agents making AC attach and reattach concurrently).
-    ///
-    /// For details, see the [task termination mechanism for GitLab runners](https://gitlab.com/gitlab-org/gitlab-runner/-/blob/397ba5dc2685e7b13feaccbfed4c242646955334/helpers/process/killer_windows.go#L75-108), which can use either mechanism dependent on a flag to use the legacy method (attach and reattach the parent process).
-    ///
-    /// Additional reading:
-    ///   - [`GenerateConsoleCtrlEvent` function](https://learn.microsoft.com/en-us/windows/console/generateconsolectrlevent), see second parameter `dwProcessGroupId`.
-    ///   - [Process Creation Flags](https://learn.microsoft.com/en-us/windows/win32/procthread/process-creation-flags)
-    fn create_process_group(cmd: &mut Command) {
-        use std::os::windows::process::CommandExt;
-        use windows::Win32::System::Threading::CREATE_NEW_PROCESS_GROUP;
-
-        // Create new process group so we can send CTRL+BREAK events to it
-        cmd.creation_flags(CREATE_NEW_PROCESS_GROUP.0);
+        #[cfg(target_family = "unix")]
+        {
+            Ok(CommandOSStarted {
+                agent_id,
+                process: child,
+                loggers,
+                shutdown_timeout: self.shutdown_timeout,
+            })
+        }
+        #[cfg(target_family = "windows")]
+        {
+            // Each started process gets its own JobObject. All sub-processes that the process spawns
+            // will be assigned to the same JobObject, allowing for a graceful shutdown of the entire process tree.
+            let job_object = JobObject::new()?;
+            job_object.assign_process(&child)?;
+            Ok(CommandOSStarted {
+                agent_id,
+                process: child,
+                job_object: Some(job_object),
+                loggers,
+                shutdown_timeout: self.shutdown_timeout,
+            })
+        }
     }
 }
 
@@ -164,12 +164,11 @@ impl CommandOSStarted {
     }
 
     pub fn shutdown(&mut self) -> Result<(), CommandError> {
-        let pid = self.get_pid();
         // Attempt a graceful shutdown (platform-dependent).
-        let graceful_shutdown_result = Self::graceful_shutdown(pid);
+        let graceful_shutdown_result = self.graceful_shutdown();
 
         if let Err(e) = &graceful_shutdown_result {
-            warn!(agent_id = %self.agent_id, "Graceful shutdown failed for process {pid}: {e}");
+            warn!(agent_id = %self.agent_id, "Graceful shutdown failed for process {}: {e}",self.get_pid());
         }
 
         if graceful_shutdown_result.is_err() || self.is_running_after_timeout(self.shutdown_timeout)
@@ -181,21 +180,25 @@ impl CommandOSStarted {
     }
 
     #[cfg(target_family = "unix")]
-    fn graceful_shutdown(pid: u32) -> Result<(), CommandError> {
+    fn graceful_shutdown(&self) -> Result<(), CommandError> {
         use nix::{sys::signal, unistd::Pid};
+        let pid = self.get_pid();
 
         signal::kill(Pid::from_raw(pid as i32), signal::SIGTERM)
-            .map_err(|e| CommandError::from(io::Error::from(e)))
+            .map_err(|e| CommandError::from(std::io::Error::from(e)))
     }
 
     #[cfg(target_family = "windows")]
-    fn graceful_shutdown(pid: u32) -> Result<(), CommandError> {
-        use windows::Win32::System::Console::{CTRL_BREAK_EVENT, GenerateConsoleCtrlEvent};
-        // Graceful shutdown for console applications
-        // <https://stackoverflow.com/a/12899284>
-        // <https://gitlab.com/gitlab-org/gitlab-runner/-/blob/397ba5dc2685e7b13feaccbfed4c242646955334/helpers/process/killer_windows.go#L75-108>
-        unsafe { GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, pid) }
-            .map_err(|e| CommandError::from(io::Error::from(e)))
+    /// On Windows there is no direct equivalent to sending SIGTERM. Applications that runs as
+    /// services handles stops signals via Service Control Manager (SCM), and console applications
+    /// can handle Ctrl-C or Ctrl-Break events via attached consoles.
+    /// The current implementation uses Job Objects to manage process groups, and there is no graceful
+    /// shutdown signal sent to the processes. The Job Object will terminate all associated processes.
+    fn graceful_shutdown(&mut self) -> Result<(), CommandError> {
+        if let Some(job_object) = self.job_object.take() {
+            job_object.kill()?;
+        }
+        Ok(())
     }
 }
 

agent-control/src/sub_agent/on_host/command/error.rs

@@ -9,6 +9,6 @@ pub enum CommandError {
     #[error("{0}")]
     IOError(#[from] std::io::Error),
 
-    #[error("Nix Error: {0}")]
-    NixError(String),
+    #[error("{0}")]
+    WinError(String),
 }

agent-control/src/sub_agent/on_host/command/job_object.rs

@@ -0,0 +1,68 @@
+use crate::sub_agent::on_host::command::error::CommandError;
+use std::os::windows::io::AsRawHandle;
+use std::process::Child;
+use tracing::error;
+use windows::Win32::Foundation::{CloseHandle, HANDLE};
+use windows::Win32::System::JobObjects::{
+    AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+    JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation,
+    SetInformationJobObject,
+};
+
+/// Represents a Windows Job Object used to manage and control a group of processes.
+/// When the Job Object is killed or dropped, all associated processes are terminated.
+pub struct JobObject {
+    handle: HANDLE,
+}
+impl JobObject {
+    /// Creates a new JobObject with the "kill on job close" configuration.
+    pub fn new() -> Result<Self, CommandError> {
+        unsafe {
+            let handle = CreateJobObjectW(None, None)
+                .map_err(|e| CommandError::WinError(format!("creating JobObject: {e}")))?;
+
+            // Set the JobObject to kill all associated processes when the JobObject is closed.
+            let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default();
+            limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+            SetInformationJobObject(
+                handle,
+                JobObjectExtendedLimitInformation,
+                &limits as *const _ as *const _,
+                std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
+            )
+            .map_err(|e| CommandError::WinError(format!("setting JobObject information: {e}")))?;
+
+            Ok(Self { handle })
+        }
+    }
+
+    /// Assigns the given process to this JobObject. The process will be terminated when the JobObject is closed.
+    pub fn assign_process(&self, process: &Child) -> Result<(), CommandError> {
+        unsafe {
+            let process_handle = HANDLE(process.as_raw_handle());
+            AssignProcessToJobObject(self.handle, process_handle).map_err(|e| {
+                CommandError::WinError(format!("assigning process to JobObject: {e}"))
+            })?;
+        }
+        Ok(())
+    }
+
+    /// Kills the JobObject, terminating all associated processes.
+    pub fn kill(self) -> Result<(), CommandError> {
+        unsafe {
+            CloseHandle(self.handle)
+                .map_err(|e| CommandError::WinError(format!("closing JobObject handle: {e}")))?;
+        }
+        Ok(())
+    }
+}
+
+/// Ensure the JobObject is killed when dropped.
+impl Drop for JobObject {
+    fn drop(&mut self) {
+        unsafe {
+            let _ =
+                CloseHandle(self.handle).inspect_err(|err| error!(%err,"Fail to kill a JobObject"));
+        }
+    }
+}

agent-control/tests/on_host/command/shutdown.rs

@@ -44,11 +44,11 @@ fn non_blocking_runner() {
 }
 
 #[test]
+#[cfg(target_family = "unix")]
 fn command_shutdown_when_sigterm_is_ignored() {
     let agent_id = "test".to_string().try_into().unwrap();
     let mut cmd = CommandOSNotStarted::new(
         agent_id,
-        #[cfg(target_family = "unix")]
         &ExecutableData {
             id: "test".to_string(),
             bin: "tests/on_host/data/ignore_sigterm.sh".to_string(),
@@ -57,21 +57,6 @@ fn command_shutdown_when_sigterm_is_ignored() {
             restart_policy: Default::default(),
             shutdown_timeout: Duration::from_millis(100),
         },
-        #[cfg(target_family = "windows")]
-        &ExecutableData {
-            id: "test".to_string(),
-            bin: "powershell".to_string(),
-            args: vec![
-                "-NoProfile".to_string(),
-                "-ExecutionPolicy".to_string(),
-                "Bypass".to_string(),
-                "-File".to_string(),
-                "tests\\on_host\\data\\ignore_sigbreak.ps1".to_string(),
-            ],
-            env: Default::default(),
-            restart_policy: Default::default(),
-            shutdown_timeout: Duration::from_millis(100),
-        },
         false,
         Default::default(),
     )
@@ -95,3 +80,56 @@ fn command_shutdown_when_sigterm_is_ignored() {
     assert!(!cmd.is_running());
     assert!(terminated.is_ok());
 }
+
+#[test]
+// This test ensure the Job Object is properly terminating orphan processes on Windows
+// On Unix, this is handled by systemd.
+#[cfg(target_family = "windows")]
+fn command_shutdown_kill_orphan_process() {
+    use crate::common::retry::retry;
+    use crate::on_host::tools::windows_process::{is_process_orphan, is_process_running};
+
+    let agent_id = "test".to_string().try_into().unwrap();
+    let id = chrono::Local::now().format("%Y%m%d_%H%M%S").to_string();
+    let mut cmd = CommandOSNotStarted::new(
+        agent_id,
+        &ExecutableData {
+            id: "test".to_string(),
+            bin: "powershell".to_string(),
+            args: vec![
+                "-NoProfile".to_string(),
+                "-ExecutionPolicy".to_string(),
+                "Bypass".to_string(),
+                "-File".to_string(),
+                "tests\\on_host\\data\\leak_sub_process.ps1".to_string(),
+                id.clone(),
+            ],
+            env: Default::default(),
+            restart_policy: Default::default(),
+            shutdown_timeout: Duration::from_millis(100),
+        },
+        false,
+        Default::default(),
+    )
+    .start()
+    .unwrap();
+
+    // Check that we have a leaked process
+    retry(30, Duration::from_secs(1), || {
+        if is_process_running(&id) && is_process_orphan(&id) {
+            Ok(())
+        } else {
+            Err("Process not running or not orphaned yet".into())
+        }
+    });
+
+    cmd.shutdown().unwrap();
+
+    retry(30, Duration::from_secs(1), || {
+        if !is_process_running(&id) {
+            Ok(())
+        } else {
+            Err("Orphan process leaked".into())
+        }
+    });
+}

agent-control/tests/on_host/data/ignore_sigbreak.ps1

@@ -0,0 +1,15 @@
+param(
+  # Just a unique identifier to make simple to find if running from tests
+  # with a command like: `gwmi Win32_Process | ? CommandLine -match "$Id" | select ProcessId, CommandLine`
+  [string]$Id = "default"
+
+)
+
+# Windows-only simple sub-process that sleeps
+$proc = Start-Process -FilePath 'powershell.exe' -ArgumentList @(
+  '-NoProfile'
+  '-NonInteractive'
+  '-Command'
+  "# test-id: $Id
+   Start-Sleep -Seconds 30"
+) -PassThru -WindowStyle Hidden

agent-control/tests/on_host/data/leak_sub_process.ps1

@@ -3,3 +3,5 @@ pub mod custom_agent_type;
 pub mod instance_id;
 pub mod oci_artifact;
 pub mod oci_package_manager;
+#[cfg(target_family = "windows")]
+pub mod windows_process;