smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(brush-core-vendored): added process cancellation support with CancellationToken and platform-specific termination

Browse Source 
- Added process cancellation support via optional CancellationToken parameter to wait() methods across process execution layers.
- Implemented platform-specific process termination using SIGKILL on Unix and TerminateProcess on Windows when cancellation is triggered.
- Added ProcessWaitResult::Cancelled variant to represent cancelled process execution with exit code 130.
- Added windows-sys dependency with Win32_Foundation and Win32_System_Threading features for Windows process termination support.
- Removed process_id_tracker and process_group_id_tracker fields from ExecutionParameters struct to simplify process tracking.
This commit is contained in:
can1357
2026-02-01 19:23:57 +01:00
parent 00e743cb29
commit 654ee503f7
8 changed files with 121 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/brush-core-vendored/Cargo.toml
+4 -8
View File
@@ -41,14 +41,6 @@ name = "brush_core"
path = "src/lib.rs"
bench = false
[[example]]
name = "call-func"
path = "examples/call-func.rs"
[[example]]
name = "custom-builtin"
path = "examples/custom-builtin.rs"
[dependencies.async-recursion]
version = "1.1.1"
@@ -166,6 +158,10 @@ version = "=0.3.4"
[target."cfg(windows)".dependencies.whoami]
version = "1.6.1"
[target."cfg(windows)".dependencies.windows-sys]
version = "0.59"
features = ["Win32_Foundation", "Win32_System_Threading"]
[target.wasm32-unknown-unknown.dependencies.getrandom]
version = "0.3.4"
features = ["wasm_js"]
crates/brush-core-vendored/src/interp.rs
+5 -7
View File
@@ -24,6 +24,7 @@ impl From<processes::ProcessWaitResult> for results::ExecutionResult {
match wait_result {
processes::ProcessWaitResult::Completed(output) => output.into(),
processes::ProcessWaitResult::Stopped => Self::stopped(),
processes::ProcessWaitResult::Cancelled => Self::new(130),
}
}
}
@@ -66,15 +67,9 @@ pub struct ExecutionParameters {
pub process_group_policy: ProcessGroupPolicy,
/// Optional cancellation token shared with callers.
cancel_token: Option<CancellationToken>,
/// Optional tracker for the first spawned process ID.
process_id_tracker: Option<std::sync::Arc<std::sync::atomic::AtomicI32>>,
/// Optional tracker for the process group ID used by execution.
process_group_id_tracker: Option<std::sync::Arc<std::sync::atomic::AtomicI32>>,
}
impl ExecutionParameters {
const UNSET_TRACKER_VALUE: i32 = 0;
/// Assigns a cancellation token for this execution.
pub fn set_cancel_token(&mut self, token: CancellationToken) {
self.cancel_token = Some(token);
@@ -509,7 +504,10 @@ async fn wait_for_pipeline_processes_and_update_status(
while let Some(child) = process_spawn_results.pop_front() {
ensure_not_cancelled(params)?;
match child.wait(!stopped_children.is_empty()).await? {
match child
.wait(!stopped_children.is_empty(), params.cancel_token())
.await?
{
ExecutionWaitResult::Completed(current_result) => {
result = current_result;
*shell.last_exit_status_mut() = result.exit_code.into();
crates/brush-core-vendored/src/jobs.rs
+6 -1
View File
@@ -43,12 +43,17 @@ impl JobTask {
pub async fn wait(&mut self) -> Result<JobTaskWaitResult, error::Error> {
match self {
Self::External(process) => {
let wait_result = process.wait().await?;
// Background jobs don't receive cancellation tokens
let wait_result = process.wait(None).await?;
match wait_result {
processes::ProcessWaitResult::Completed(output) => {
Ok(JobTaskWaitResult::Completed(output.into()))
}
processes::ProcessWaitResult::Stopped => Ok(JobTaskWaitResult::Stopped),
processes::ProcessWaitResult::Cancelled => {
// Should never happen since we pass None, but handle gracefully
Ok(JobTaskWaitResult::Completed(ExecutionResult::new(130)))
}
}
}
Self::Internal(handle) => Ok(JobTaskWaitResult::Completed(handle.await??)),
crates/brush-core-vendored/src/processes.rs
+47 -1
View File
@@ -1,6 +1,7 @@
//! Process management
use futures::FutureExt;
use tokio_util::sync::CancellationToken;
use crate::{error, sys};
@@ -32,18 +33,35 @@ impl ChildProcess {
}
/// Waits for the process to exit.
pub async fn wait(&mut self) -> Result<ProcessWaitResult, error::Error> {
///
/// If a cancellation token is provided and triggered, the process will be killed.
pub async fn wait(
&mut self,
cancel_token: Option<CancellationToken>,
) -> Result<ProcessWaitResult, error::Error> {
#[allow(unused_mut, reason = "only mutated on some platforms")]
let mut sigtstp = sys::signal::tstp_signal_listener()?;
#[allow(unused_mut, reason = "only mutated on some platforms")]
let mut sigchld = sys::signal::chld_signal_listener()?;
let cancelled = async {
match &cancel_token {
Some(token) => token.cancelled().await,
None => std::future::pending().await,
}
};
tokio::pin!(cancelled);
#[allow(clippy::ignored_unit_patterns)]
loop {
tokio::select! {
output = &mut self.exec_future => {
break Ok(ProcessWaitResult::Completed(output?))
},
_ = &mut cancelled => {
self.kill();
break Ok(ProcessWaitResult::Cancelled)
},
_ = sigtstp.recv() => {
break Ok(ProcessWaitResult::Stopped)
},
@@ -61,6 +79,32 @@ impl ChildProcess {
}
}
/// Terminates the process if we have a PID.
fn kill(&self) {
let Some(pid) = self.pid else { return };
#[cfg(unix)]
{
let _ = nix::sys::signal::kill(nix::unistd::Pid::from_raw(pid), nix::sys::signal::Signal::SIGKILL);
}
#[cfg(windows)]
{
use windows_sys::Win32::Foundation::CloseHandle;
use windows_sys::Win32::System::Threading::{OpenProcess, TerminateProcess, PROCESS_TERMINATE};
// SAFETY: Windows API calls with proper handle management
unsafe {
#[expect(clippy::cast_sign_loss)]
let handle = OpenProcess(PROCESS_TERMINATE, 0, pid as u32);
if handle != 0 {
let _ = TerminateProcess(handle, 1);
CloseHandle(handle);
}
}
}
}
pub(crate) fn poll(&mut self) -> Option<Result<std::process::Output, error::Error>> {
let checkable_future = &mut self.exec_future;
checkable_future
@@ -75,4 +119,6 @@ pub enum ProcessWaitResult {
Completed(std::process::Output),
/// The process stopped and has not yet completed.
Stopped,
/// The process was killed due to cancellation.
Cancelled,
}
crates/brush-core-vendored/src/results.rs
+13 -2
View File
@@ -220,19 +220,26 @@ impl From<ExecutionResult> for ExecutionSpawnResult {
}
}
use tokio_util::sync::CancellationToken;
impl ExecutionSpawnResult {
/// Waits for the command to complete.
///
/// # Arguments
///
/// * `no_wait` - If true, do not wait for the command to complete; return immediately.
pub async fn wait(self, no_wait: bool) -> Result<ExecutionWaitResult, error::Error> {
/// * `cancel_token` - Optional cancellation token; if triggered, kills the process.
pub async fn wait(
self,
no_wait: bool,
cancel_token: Option<CancellationToken>,
) -> Result<ExecutionWaitResult, error::Error> {
match self {
Self::StartedProcess(mut child) => {
let process_wait_result = if !no_wait {
// Wait for the process to exit or for a relevant signal, whichever happens
// first.
child.wait().await?
child.wait(cancel_token).await?
} else {
processes::ProcessWaitResult::Stopped
};
@@ -242,6 +249,10 @@ impl ExecutionSpawnResult {
ExecutionWaitResult::Completed(ExecutionResult::from(output))
}
processes::ProcessWaitResult::Stopped => ExecutionWaitResult::Stopped(child),
processes::ProcessWaitResult::Cancelled => {
// 130 = 128 + SIGINT (2), standard shell interrupted exit code
ExecutionWaitResult::Completed(ExecutionResult::new(130))
}
};
Ok(wait_result)
crates/pi-natives/src/shell.rs
+39 -44
View File
@@ -26,9 +26,6 @@ use std::{
#[cfg(windows)]
mod windows;
#[cfg(windows)]
use windows::configure_windows_path;
use brush_builtins::{BuiltinSet, default_builtins};
use brush_core::{
CreateOptions, ExecutionContext, ExecutionControlFlow, ExecutionExitCode, ExecutionResult,
@@ -45,6 +42,8 @@ use napi::{
use napi_derive::napi;
use parking_lot::Mutex;
use tokio_util::sync::CancellationToken;
#[cfg(windows)]
use windows::configure_windows_path;
use crate::work::launch_task;
@@ -56,8 +55,6 @@ struct ExecutionControl {
session_key: String,
}
struct ExecutionGuard {
execution_id: String,
}
@@ -263,8 +260,7 @@ async fn execute_shell_with_options(
let run_result = {
let mut session = session.lock().await;
let run_future =
run_shell_command(&mut session, &options, on_chunk, cancel_token.clone());
let run_future = run_shell_command(&mut session, &options, on_chunk, cancel_token.clone());
tokio::pin!(run_future);
let run_result = if let Some(ms) = timeout_ms {
@@ -588,7 +584,6 @@ fn remove_session(session_key: &str) {
sessions.remove(session_key);
}
fn read_output(mut reader: std::fs::File, on_chunk: Option<ThreadsafeFunction<String>>) {
let mut buf = [0u8; 8192];
let mut pending = Vec::new();
@@ -691,9 +686,9 @@ impl builtins::Command for SleepCommand {
) -> impl Future<Output = std::result::Result<ExecutionResult, brush_core::Error>> + Send {
let durations = self.durations.clone();
async move {
if context.is_cancelled() {
return Ok(ExecutionExitCode::Interrupted.into());
}
if context.is_cancelled() {
return Ok(ExecutionExitCode::Interrupted.into());
}
let mut total = Duration::from_millis(0);
for duration in &durations {
let Some(parsed) = parse_duration(duration) else {
@@ -701,20 +696,20 @@ impl builtins::Command for SleepCommand {
return Ok(ExecutionResult::new(1));
};
total += parsed;
}
let sleep = time::sleep(total);
tokio::pin!(sleep);
if let Some(cancel_token) = context.cancel_token() {
tokio::select! {
() = &mut sleep => Ok(ExecutionResult::success()),
_ = cancel_token.cancelled() => Ok(ExecutionExitCode::Interrupted.into()),
}
} else {
sleep.await;
Ok(ExecutionResult::success())
let sleep = time::sleep(total);
tokio::pin!(sleep);
if let Some(cancel_token) = context.cancel_token() {
tokio::select! {
() = &mut sleep => Ok(ExecutionResult::success()),
() = cancel_token.cancelled() => Ok(ExecutionExitCode::Interrupted.into()),
}
} else {
sleep.await;
Ok(ExecutionResult::success())
}
}
}
}
}
#[derive(Parser)]
@@ -736,9 +731,9 @@ impl builtins::Command for TimeoutCommand {
let duration = self.duration.clone();
let command = self.command.clone();
async move {
if context.is_cancelled() {
return Ok(ExecutionExitCode::Interrupted.into());
}
if context.is_cancelled() {
return Ok(ExecutionExitCode::Interrupted.into());
}
let Some(timeout) = parse_duration(&duration) else {
let _ = writeln!(context.stderr(), "timeout: invalid time interval '{duration}'");
return Ok(ExecutionResult::new(125));
@@ -755,29 +750,29 @@ impl builtins::Command for TimeoutCommand {
for (idx, arg) in command.iter().enumerate() {
if idx > 0 {
command_line.push(' ');
}
}
command_line.push_str(&quote_arg(arg));
}
let cancel_token = context.cancel_token();
}
let cancel_token = context.cancel_token();
let run_future = context.shell.run_string(command_line, &params);
tokio::pin!(run_future);
let result = if let Some(cancel_token) = cancel_token {
tokio::select! {
result = &mut run_future => result,
() = time::sleep(timeout) => Ok(ExecutionResult::new(124)),
_ = cancel_token.cancelled() => Ok(ExecutionExitCode::Interrupted.into()),
}
} else {
tokio::select! {
result = &mut run_future => result,
() = time::sleep(timeout) => Ok(ExecutionResult::new(124)),
}
};
Ok(result?)
}
}
}
if let Some(cancel_token) = cancel_token {
tokio::select! {
result = &mut run_future => result,
() = time::sleep(timeout) => Ok(ExecutionResult::new(124)),
() = cancel_token.cancelled() => Ok(ExecutionExitCode::Interrupted.into()),
}
} else {
tokio::select! {
result = &mut run_future => result,
() = time::sleep(timeout) => Ok(ExecutionResult::new(124)),
}
}
}
}
}
fn parse_duration(input: &str) -> Option<Duration> {
let trimmed = input.trim();
if trimmed.is_empty() {
crates/pi-natives/src/shell/windows.rs
+6 -18
View File
@@ -1,14 +1,8 @@
use std::{
path::Path,
process::Command,
};
use std::{path::Path, process::Command};
use brush_core::{Shell as BrushShell, ShellValue, ShellVariable};
use napi::{Error, Result};
use winreg::{
enums::HKEY_LOCAL_MACHINE,
RegKey,
};
use winreg::{RegKey, enums::HKEY_LOCAL_MACHINE};
pub fn configure_windows_path(shell: &mut BrushShell) -> Result<()> {
let Some(git_usr_bin) = find_git_usr_bin() else {
@@ -55,12 +49,9 @@ fn path_contains_entry(path_value: &str, entry: &str) -> bool {
}
fn find_git_usr_bin() -> Option<String> {
for install_path in [
query_git_install_path_from_registry(),
query_git_install_path_from_where(),
]
.into_iter()
.flatten()
for install_path in [query_git_install_path_from_registry(), query_git_install_path_from_where()]
.into_iter()
.flatten()
{
if let Some(path) = git_usr_bin_with_ls(&install_path) {
return Some(path);
@@ -72,10 +63,7 @@ fn find_git_usr_bin() -> Option<String> {
fn query_git_install_path_from_registry() -> Option<String> {
let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
let key_paths = [
"SOFTWARE\\GitForWindows",
"SOFTWARE\\WOW6432Node\\GitForWindows",
];
let key_paths = ["SOFTWARE\\GitForWindows", "SOFTWARE\\WOW6432Node\\GitForWindows"];
for key_path in key_paths {
if let Ok(key) = hklm.open_subkey(key_path) {
rustfmt.toml
+1
View File
@@ -71,3 +71,4 @@ blank_lines_upper_bound = 1
force_explicit_abi = true
hex_literal_case = "Lower"
newline_style = "Unix"
ignore = ["crates/brush-core-vendored/**"]