Skip to content

Commit

Permalink
wip: blocking
Browse files Browse the repository at this point in the history
  • Loading branch information
@oneofthezombies
oneofthezombies committed Feb 11, 2024
1 parent ab4ef6f commit 1ffb2da
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 124 deletions.
70 changes: 25 additions & 45 deletions crates/libs/kill_tree/src/unix.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,37 +38,6 @@ pub(crate) fn validate_process_id(process_id: ProcessId, available_max: ProcessI
}
}

fn parse_process_id_sign(process_id: ProcessId) -> Result<i32> {
let process_id_sign = i32::try_from(process_id).map_err(|e| Error::InvalidCast {
reason: "Failed to cast process id to i32".into(),
source: e,
})?;
Ok(process_id_sign)
}

fn parse_kill_result(result: nix::Result<()>, process_id: ProcessId) -> Result<KillOutput> {
match result {
Ok(()) => Ok(KillOutput::Killed { process_id }),
Err(e) => {
// ESRCH: No such process.
// This happens when the process has already terminated.
// This treat as success.
if e == nix::errno::Errno::ESRCH {
Ok(KillOutput::MaybeAlreadyTerminated {
process_id,
source: e.into(),
})
} else {
Err(e.into())
}
}
}
}

fn kill(process_id_sign: i32, signal: nix::sys::signal::Signal) -> nix::Result<()> {
nix::sys::signal::kill(nix::unistd::Pid::from_raw(process_id_sign), signal)
}

pub(crate) mod blocking {
use super::*;

Expand All @@ -92,16 +61,33 @@ pub(crate) mod blocking {
process_id: ProcessId,
signal: nix::sys::signal::Signal,
) -> Result<KillOutput> {
let process_id_sign = parse_process_id_sign(process_id)?;
let result = crate::unix::kill(process_id_sign, signal);
parse_kill_result(result, process_id)
let process_id_sign = i32::try_from(process_id).map_err(|e| Error::InvalidCast {
reason: "Failed to cast process id to i32".into(),
source: e,
})?;
let result = nix::sys::signal::kill(nix::unistd::Pid::from_raw(process_id_sign), signal);
match result {
Ok(()) => Ok(KillOutput::Killed { process_id }),
Err(e) => {
// ESRCH: No such process.
// This happens when the process has already terminated.
// This treat as success.
if e == nix::errno::Errno::ESRCH {
Ok(KillOutput::MaybeAlreadyTerminated {
process_id,
source: e.into(),
})
} else {
Err(e.into())
}
}
}
}
}

pub(crate) mod tokio {
use async_trait::async_trait;

use super::*;
use async_trait::async_trait;

#[derive(Clone)]
struct Killer {
Expand All @@ -121,15 +107,9 @@ pub(crate) mod tokio {
}

#[instrument]
pub(crate) async fn kill(
process_id: ProcessId,
signal: nix::sys::signal::Signal,
) -> Result<KillOutput> {
let process_id_sign = parse_process_id_sign(process_id)?;
let result =
::tokio::task::spawn_blocking(move || crate::unix::kill(process_id_sign, signal))
.await?;
parse_kill_result(result, process_id)
async fn kill(process_id: ProcessId, signal: nix::sys::signal::Signal) -> Result<KillOutput> {
::tokio::task::spawn_blocking(move || crate::unix::blocking::kill(process_id_sign, signal))
.await;
}
}

Expand Down
207 changes: 128 additions & 79 deletions crates/libs/kill_tree/src/windows.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,102 +43,151 @@ pub(crate) fn child_process_id_map_filter(process_info: &ProcessInfo) -> bool {
process_info.parent_process_id == process_info.process_id
}

#[instrument]
pub(crate) fn kill(process_id: ProcessId, config: &Config) -> Result<KillOutput> {
let result: Result<KillOutput>;
unsafe {
let open_result = OpenProcess(PROCESS_TERMINATE, false, process_id);
match open_result {
Ok(process_handle) => {
{
// do NOT return early from this block
result = TerminateProcess(process_handle, 1)
.and(Ok(KillOutput::Killed { process_id }))
.or_else(|e| {
if e.code() == E_ACCESSDENIED {
// Access is denied.
// This happens when the process is already terminated.
// This treat as success.
Ok(KillOutput::MaybeAlreadyTerminated {
process_id,
source: e.into(),
})
} else {
Err(e.into())
}
});
#[cfg(feature = "blocking")]
pub(crate) mod blocking {
use super::*;

struct Killer {}

impl crate::blocking::Killable for Killer {
fn kill(&self, process_id: ProcessId) -> Result<KillOutput> {
crate::windows::blocking::kill(process_id)
}
}

pub(crate) fn new_killer(_config: &Config) -> Result<impl crate::blocking::Killable> {
Ok(Killer {})
}

#[instrument]
pub(crate) fn kill(process_id: ProcessId) -> Result<KillOutput> {
let result: Result<KillOutput>;
unsafe {
let open_result = OpenProcess(PROCESS_TERMINATE, false, process_id);
match open_result {
Ok(process_handle) => {
{
// do NOT return early from this block
result = TerminateProcess(process_handle, 1)
.and(Ok(KillOutput::Killed { process_id }))
.or_else(|e| {
if e.code() == E_ACCESSDENIED {
// Access is denied.
// This happens when the process is already terminated.
// This treat as success.
Ok(KillOutput::MaybeAlreadyTerminated {
process_id,
source: e.into(),
})
} else {
Err(e.into())
}
});
}
CloseHandle(process_handle)?;
}
CloseHandle(process_handle)?;
}
Err(e) => {
if e.code() == E_INVALIDARG {
// The parameter is incorrect.
// This happens when the process is already terminated.
// This treat as success.
result = Ok(KillOutput::MaybeAlreadyTerminated {
process_id,
source: e.into(),
});
} else {
result = Err(e.into());
Err(e) => {
if e.code() == E_INVALIDARG {
// The parameter is incorrect.
// This happens when the process is already terminated.
// This treat as success.
result = Ok(KillOutput::MaybeAlreadyTerminated {
process_id,
source: e.into(),
});
} else {
result = Err(e.into());
}
}
}
}
result
}
result
}

#[instrument]
pub(crate) fn get_process_infos() -> Result<ProcessInfos> {
let mut process_infos = ProcessInfos::new();
let mut error: Option<Error> = None;
unsafe {
let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?;
{
// do NOT return early from this block
let mut process_entry = std::mem::zeroed::<PROCESSENTRY32>();
let process_entry_size = u32::try_from(std::mem::size_of::<PROCESSENTRY32>());
match process_entry_size {
Ok(process_entry_size) => {
process_entry.dwSize = process_entry_size;
match Process32First(snapshot_handle, &mut process_entry) {
Ok(()) => loop {
process_infos.push(ProcessInfo {
process_id: process_entry.th32ProcessID,
parent_process_id: process_entry.th32ParentProcessID,
name: ffi::CStr::from_ptr(process_entry.szExeFile.as_ptr().cast())
#[instrument]
pub(crate) fn get_process_infos() -> Result<ProcessInfos> {
let mut process_infos = ProcessInfos::new();
let mut error: Option<Error> = None;
unsafe {
let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?;
{
// do NOT return early from this block
let mut process_entry = std::mem::zeroed::<PROCESSENTRY32>();
let process_entry_size = u32::try_from(std::mem::size_of::<PROCESSENTRY32>());
match process_entry_size {
Ok(process_entry_size) => {
process_entry.dwSize = process_entry_size;
match Process32First(snapshot_handle, &mut process_entry) {
Ok(()) => loop {
process_infos.push(ProcessInfo {
process_id: process_entry.th32ProcessID,
parent_process_id: process_entry.th32ParentProcessID,
name: ffi::CStr::from_ptr(
process_entry.szExeFile.as_ptr().cast(),
)
.to_string_lossy()
.into_owned(),
});
match Process32Next(snapshot_handle, &mut process_entry) {
Ok(()) => {}
Err(e) => {
if e.code() != ERROR_NO_MORE_FILES.into() {
error = Some(e.into());
});
match Process32Next(snapshot_handle, &mut process_entry) {
Ok(()) => {}
Err(e) => {
if e.code() != ERROR_NO_MORE_FILES.into() {
error = Some(e.into());
}
break;
}
break;
}
},
Err(e) => {
error = Some(e.into());
}
},
Err(e) => {
error = Some(e.into());
}
}
}
Err(e) => {
error = Some(Error::InvalidCast {
source: e,
reason: "size of PROCESSENTRY32 to u32".into(),
});
Err(e) => {
error = Some(Error::InvalidCast {
source: e,
reason: "size of PROCESSENTRY32 to u32".into(),
});
}
}
}
CloseHandle(snapshot_handle)?;
}
if let Some(e) = error {
Err(e)
} else {
Ok(process_infos)
}
CloseHandle(snapshot_handle)?;
}
if let Some(e) = error {
Err(e)
} else {
Ok(process_infos)
}

#[cfg(feature = "tokio")]
pub(crate) mod tokio {
use super::*;
use async_trait::async_trait;

#[derive(Clone)]
struct Killer {}

#[async_trait]
impl crate::tokio::Killable for Killer {
async fn kill(&self, process_id: ProcessId) -> Result<KillOutput> {
crate::windows::tokio::kill(process_id).await
}
}

pub(crate) fn new_killer(_config: &Config) -> Result<impl crate::tokio::Killable> {
Ok(Killer {})
}

#[instrument]
async fn kill(process_id: ProcessId) -> Result<KillOutput> {
::tokio::task::spawn_blocking(move || crate::windows::blocking::kill(process_id)).await?
}

#[instrument]
pub(crate) async fn get_process_infos() -> Result<ProcessInfos> {
::tokio::task::spawn_blocking(move || blocking::get_process_infos()).await?
}
}

Expand Down

0 comments on commit 1ffb2da

Please sign in to comment.