From 052b4e77a60847de0ccd275f4c7fc241378c9bc0 Mon Sep 17 00:00:00 2001 From: Shane Kennedy Date: Tue, 12 Nov 2024 03:48:21 +0100 Subject: [PATCH] Handle sigterm calls, fixes #6724 (#8933) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR builds off of https://github.com/astral-sh/uv/pull/6738 to fix #6724 (sorry for the new PR @charliermarsh I didn't want to push to your branch, not even sure if I could). The reason the original PR doesn't fix the issue described in #6724 is because the fastapi is ran in the project context (as I assume a lot of use cases are). This PR adds an extra commit to handle the signals in the project/run.rs file ~It also addresses the comment [here](https://github.com/astral-sh/uv/pull/6738/files#r1734757548) to not use the tokio ctrl-c method since we are now handling SIGINT ourselves~ update, tokio handles SIGINT in a platform agnostic way, intercepting this ouselves makes the logic more complicated with windows, decided to leave the tokio ctrl-c handler ~[This comment](https://github.com/astral-sh/uv/pull/6738/files#r1743510140) remains unaddressed, however, the Child process does not have any other methods besides kill() so I don't see how we can "preserve" the interrupt call :/ I tried looking around but no luck.~ updated, this PR is reduced to only handling SIGTERM propagation on unix machines, and the sigterm call to the child is preserved by making use of the nix package, instead of relying on tokio which only allowed for `kill()` on a child process ## Test Plan I tested this by building the docker container locally with these changes and tagging it "myuv", and then using that as the base image in uv-docker-example, (and ofc following the rest of the repro issues in #6724. In my tests I see that ctrl-c in the docker-compose up command exits the process almost immediately 👍 --------- Co-authored-by: Charlie Marsh --- Cargo.lock | 3 ++- Cargo.toml | 5 ++-- crates/uv/Cargo.toml | 3 +++ crates/uv/src/commands/project/run.rs | 34 +++++++++++++++++++++++++-- crates/uv/src/commands/tool/run.rs | 34 +++++++++++++++++++++++++-- 5 files changed, 72 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fdd98f432..6d886e80f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4376,6 +4376,7 @@ dependencies = [ "itertools 0.13.0", "jiff", "miette", + "nix", "owo-colors", "petgraph", "predicates", @@ -5798,7 +5799,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 3075f6793..f055ca98e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -120,15 +120,15 @@ md-5 = { version = "0.10.6" } memchr = { version = "2.7.4" } miette = { version = "7.2.0" } nanoid = { version = "0.4.0" } +nix = { version = "0.29.0" } owo-colors = { version = "4.1.0" } path-slash = { version = "0.2.1" } pathdiff = { version = "0.2.1" } petgraph = { version = "0.6.5" } platform-info = { version = "2.0.3" } -procfs = { version = "0.17.0", default-features = false, features = ["flate2"] } proc-macro2 = { version = "1.0.86" } +procfs = { version = "0.17.0", default-features = false, features = ["flate2"] } pubgrub = { git = "https://github.com/astral-sh/pubgrub", rev = "95e1390399cdddee986b658be19587eb1fdb2d79" } -version-ranges = { git = "https://github.com/astral-sh/pubgrub", rev = "95e1390399cdddee986b658be19587eb1fdb2d79" } quote = { version = "1.0.37" } rayon = { version = "1.10.0" } reflink-copy = { version = "0.1.19" } @@ -172,6 +172,7 @@ unicode-width = { version = "0.1.13" } unscanny = { version = "0.1.0" } url = { version = "2.5.2" } urlencoding = { version = "2.1.3" } +version-ranges = { git = "https://github.com/astral-sh/pubgrub", rev = "95e1390399cdddee986b658be19587eb1fdb2d79" } walkdir = { version = "2.5.0" } which = { version = "7.0.0", features = ["regex"] } windows-registry = { version = "0.3.0" } diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml index b2be62851..1d0af552e 100644 --- a/crates/uv/Cargo.toml +++ b/crates/uv/Cargo.toml @@ -115,6 +115,9 @@ similar = { version = "2.6.0" } tempfile = { workspace = true } zip = { workspace = true } +[target.'cfg(unix)'.dependencies] +nix = { workspace = true } + [package.metadata.cargo-shear] ignored = [ "flate2", diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index 76c759cb9..6ce8bcbcc 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -996,9 +996,30 @@ 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")?; + // Exit based on the result of the command. + #[cfg(unix)] + let status = { + use tokio::select; + use tokio::signal::unix::{signal, SignalKind}; + + let mut term_signal = signal(SignalKind::terminate())?; + loop { + select! { + result = handle.wait() => { + break result; + }, + + // `SIGTERM` + _ = term_signal.recv() => { + let _ = terminate_process(&mut handle); + } + }; + } + }?; + + #[cfg(not(unix))] + let status = handle.wait().await?; - // Exit based on the result of the command if let Some(code) = status.code() { debug!("Command exited with code: {code}"); if let Ok(code) = u8::try_from(code) { @@ -1017,6 +1038,15 @@ pub(crate) async fn run( } } +#[cfg(unix)] +fn terminate_process(child: &mut tokio::process::Child) -> anyhow::Result<()> { + use nix::sys::signal::{self, Signal}; + use nix::unistd::Pid; + + let pid = child.id().context("Failed to get child process ID")?; + signal::kill(Pid::from_raw(pid.try_into()?), Signal::SIGTERM).context("Failed to send SIGTERM") +} + /// Returns `true` if we can skip creating an additional ephemeral environment in `uv run`. fn can_skip_ephemeral( spec: Option<&RequirementsSpecification>, diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs index 43816a5d7..16062ccc9 100644 --- a/crates/uv/src/commands/tool/run.rs +++ b/crates/uv/src/commands/tool/run.rs @@ -236,9 +236,30 @@ 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")?; + // Exit based on the result of the command. + #[cfg(unix)] + let status = { + use tokio::select; + use tokio::signal::unix::{signal, SignalKind}; + + let mut term_signal = signal(SignalKind::terminate())?; + loop { + select! { + result = handle.wait() => { + break result; + }, + + // `SIGTERM` + _ = term_signal.recv() => { + let _ = terminate_process(&mut handle); + } + }; + } + }?; + + #[cfg(not(unix))] + let status = handle.wait().await?; - // Exit based on the result of the command if let Some(code) = status.code() { debug!("Command exited with code: {code}"); if let Ok(code) = u8::try_from(code) { @@ -257,6 +278,15 @@ pub(crate) async fn run( } } +#[cfg(unix)] +fn terminate_process(child: &mut tokio::process::Child) -> anyhow::Result<()> { + use nix::sys::signal::{self, Signal}; + use nix::unistd::Pid; + + let pid = child.id().context("Failed to get child process ID")?; + signal::kill(Pid::from_raw(pid.try_into()?), Signal::SIGTERM).context("Failed to send SIGTERM") +} + /// Return the entry points for the specified package. fn get_entrypoints( from: &PackageName,