Skip to content

Commit

Permalink
Runtime process termination on Windows with timeout
Browse files Browse the repository at this point in the history
  • Loading branch information
pwalski committed Jan 10, 2024
1 parent df51447 commit b8d689e
Showing 1 changed file with 39 additions and 35 deletions.
74 changes: 39 additions & 35 deletions utils/process/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use anyhow::{anyhow, Result};
use derive_more::Display;
use futures::channel::oneshot::channel;
use futures::future::{AbortHandle, Abortable};
use shared_child::SharedChild;
use std::process::Command;
use std::sync::Arc;
Expand All @@ -11,17 +12,14 @@ use std::time::Duration;
pub mod lock;

#[cfg(unix)]
use {
futures::future::{AbortHandle, Abortable},
shared_child::unix::SharedChildExt,
};
use shared_child::unix::SharedChildExt;

#[cfg(windows)]
use {
winapi::um::handleapi::CloseHandle,
winapi::um::processthreadsapi::OpenProcess,
winapi::um::wincon::{GenerateConsoleCtrlEvent, CTRL_BREAK_EVENT},
winapi::um::winnt::{PROCESS_QUERY_INFORMATION, PROCESS_TERMINATE, SYNCHRONIZE},
use std::process::ExitStatus;
#[cfg(windows)]
use winapi::um::{
errhandlingapi,
wincon::{GenerateConsoleCtrlEvent, CTRL_BREAK_EVENT},
};

pub trait ProcessGroupExt<T> {
Expand Down Expand Up @@ -125,36 +123,28 @@ impl ProcessHandle {
}

#[cfg(windows)]
pub async fn terminate(&self, _timeout: Duration) -> Result<()> {
pub async fn terminate(&self, timeout: Duration) -> Result<()> {
use anyhow::bail;

let process = self.process.clone();
let process_pid = process.id();
let (abort_handle, abort_registration) = AbortHandle::new_pair();
let _process_pid = process.id();

unsafe {
let process_handle = OpenProcess(
PROCESS_QUERY_INFORMATION | SYNCHRONIZE | PROCESS_TERMINATE,
0,
process_pid,
);

if process_handle.is_null() {
return Err(anyhow!(
"Unable to open the process. Error code: {}",
winapi::um::errhandlingapi::GetLastError()
));
}

let event_result = GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, process_pid);

if event_result == 0 {
return Err(anyhow!(
"Unable to send CTRL+C event to the process. Error code: {}",
winapi::um::errhandlingapi::GetLastError()
));
};
CloseHandle(process_handle);
tokio::task::spawn_local(async move {
tokio::time::sleep(timeout).await;
abort_handle.abort();
});

if let Ok(Err(err)) = Abortable::new(
async { unsafe { ctrl_break_process(process) } },
abort_registration,
)
.await
{
bail!(err.context("Failed to CTRL-BREAK process"));
}

Ok(())
self.check_if_running()
}

pub fn check_if_running(&self) -> Result<()> {
Expand Down Expand Up @@ -200,3 +190,17 @@ impl ProcessHandle {
receiver.await.unwrap()
}
}

#[cfg(windows)]
unsafe fn ctrl_break_process(process: Arc<SharedChild>) -> Result<ExitStatus> {
let process_pid = process.id();
let event_result = GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, process_pid);

if event_result == 0 {
return Err(anyhow!(
"Unable to send CTRL-BREAK event to the process. Error code: {}",
errhandlingapi::GetLastError()
));
};
Ok(process.wait()?)
}

0 comments on commit b8d689e

Please sign in to comment.