Skip to content

Commit

Permalink
Handle platform differences and propagate sigterm
Browse files Browse the repository at this point in the history
This commit handles sigterm explicitly,and implements the different
sigterm propagation based on windows vs unix systems. This commit
does not provide a windows implementation for handling sigterm.

Affected paths are project/run and tool/run
  • Loading branch information
shaneikennedy committed Nov 8, 2024
1 parent 6da70d1 commit ee2e751
Show file tree
Hide file tree
Showing 4 changed files with 105 additions and 28 deletions.
30 changes: 27 additions & 3 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions crates/uv/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ url = { workspace = true }
walkdir = { workspace = true }
which = { workspace = true }
zip = { workspace = true }
cfg-if = "1.0"

[dev-dependencies]
assert_cmd = { version = "2.0.16" }
Expand All @@ -115,6 +116,9 @@ similar = { version = "2.6.0" }
tempfile = { workspace = true }
zip = { workspace = true }

[target.'cfg(unix)'.dependencies]
nix = "0.23"

[package.metadata.cargo-shear]
ignored = [
"flate2",
Expand Down
39 changes: 37 additions & 2 deletions crates/uv/src/commands/project/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@ use std::path::{Path, PathBuf};

use anstream::eprint;
use anyhow::{anyhow, bail, Context};
use cfg_if::cfg_if;
use futures::StreamExt;
use itertools::Itertools;
use owo_colors::OwoColorize;
use tokio::process::Command;
use tokio::process::{Child, Command};
use tracing::{debug, warn};
use url::Url;
use uv_cache::Cache;
Expand Down Expand Up @@ -996,9 +997,27 @@ pub(crate) async fn run(
// signal handlers after the command completes.
let _handler = tokio::spawn(async { while tokio::signal::ctrl_c().await.is_ok() {} });

let status = handle.wait().await.context("Child process disappeared")?;
#[cfg(unix)]
{
use tokio::select;
use tokio::signal::unix::{signal, SignalKind};
let mut term_signal = signal(SignalKind::terminate())?;
loop {
select! {
_ = handle.wait() => {
break
},

// `SIGTERM`
_ = term_signal.recv() => {
let _ = terminate_process(&mut handle);
}
};
}
}

// Exit based on the result of the command
let status = handle.wait().await?;
if let Some(code) = status.code() {
debug!("Command exited with code: {code}");
if let Ok(code) = u8::try_from(code) {
Expand All @@ -1017,6 +1036,22 @@ pub(crate) async fn run(
}
}

#[allow(clippy::items_after_statements, dead_code)]
fn terminate_process(_child: &mut Child) -> Result<(), anyhow::Error> {
cfg_if! {
if #[cfg(unix)] {
let pid = _child.id().context("Failed to get child process ID")?;
use nix::sys::signal::{self, Signal};
use nix::unistd::Pid;
signal::kill(Pid::from_raw(pid.try_into().unwrap()), Signal::SIGTERM)
.context("Failed to send SIGTERM")
} else if #[cfg(windows)] {
// On Windows, use winapi to terminate the process gracefully
todo!("Implement graceful termination on Windows");
}
}
}

/// Returns `true` if we can skip creating an additional ephemeral environment in `uv run`.
fn can_skip_ephemeral(
spec: Option<&RequirementsSpecification>,
Expand Down
60 changes: 37 additions & 23 deletions crates/uv/src/commands/tool/run.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use cfg_if::cfg_if;
use std::fmt::Display;
use std::fmt::Write;
use std::path::PathBuf;
Expand All @@ -7,9 +8,8 @@ use anstream::eprint;
use anyhow::{bail, Context};
use itertools::Itertools;
use owo_colors::OwoColorize;
use tokio::process::Child;
use tokio::process::Command;
use tokio::select;
use tokio::signal::unix::{signal, SignalKind};
use tracing::{debug, warn};

use uv_cache::{Cache, Refresh};
Expand Down Expand Up @@ -238,30 +238,27 @@ pub(crate) async fn run(
// signal handlers after the command completes.
let _handler = tokio::spawn(async { while tokio::signal::ctrl_c().await.is_ok() {} });

let mut term_signal = signal(SignalKind::terminate())?;
let mut int_signal = signal(SignalKind::interrupt())?;

let status = select! {
status = handle.wait() => status,

// `SIGTERM`
_ = term_signal.recv() => {
handle.kill().await?;
handle.wait().await.context("Child process disappeared")?;
return Ok(ExitStatus::Failure);
}

// `SIGINT`
_ = int_signal.recv() => {
handle.kill().await?;
handle.wait().await.context("Child process disappeared")?;
return Ok(ExitStatus::Failure);
#[cfg(unix)]
{
use tokio::select;
use tokio::signal::unix::{signal, SignalKind};
let mut term_signal = signal(SignalKind::terminate())?;
loop {
select! {
_ = handle.wait() => {
break
},

// `SIGTERM`
_ = term_signal.recv() => {
let _ = terminate_process(&mut handle);
}
};
}
};

let status = status.context("Child process disappeared")?;
}

// Exit based on the result of the command
let status = handle.wait().await?;
if let Some(code) = status.code() {
debug!("Command exited with code: {code}");
if let Ok(code) = u8::try_from(code) {
Expand All @@ -280,6 +277,23 @@ pub(crate) async fn run(
}
}

#[allow(clippy::items_after_statements, dead_code)]
fn terminate_process(_child: &mut Child) -> Result<(), anyhow::Error> {
cfg_if! {
if #[cfg(unix)] {
let pid = _child.id().context("Failed to get child process ID")?;
// On Unix, send SIGTERM for a graceful shutdown
use nix::sys::signal::{self, Signal};
use nix::unistd::Pid;
signal::kill(Pid::from_raw(pid.try_into().unwrap()), Signal::SIGTERM)
.context("Failed to send SIGTERM")
} else if #[cfg(windows)] {
// On Windows, use winapi to terminate the process gracefully
todo!("Implement graceful termination on Windows");
}
}
}

/// Return the entry points for the specified package.
fn get_entrypoints(
from: &PackageName,
Expand Down

0 comments on commit ee2e751

Please sign in to comment.