diff --git a/utils/process/src/lib.rs b/utils/process/src/lib.rs index 76f5e06030..21110d5d19 100644 --- a/utils/process/src/lib.rs +++ b/utils/process/src/lib.rs @@ -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; @@ -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 { @@ -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<()> { @@ -200,3 +190,17 @@ impl ProcessHandle { receiver.await.unwrap() } } + +#[cfg(windows)] +unsafe fn ctrl_break_process(process: Arc) -> Result { + 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()?) +}