From 7dd4436a913bfd5b6c54583c3a76150d7f81d234 Mon Sep 17 00:00:00 2001 From: Shane Kennedy Date: Fri, 8 Nov 2024 14:34:04 +0100 Subject: [PATCH] Handle platform differences and propagate sigterm This commit now only handles sigterm explicitly, using tokio sigint handling, and implements the different sigterm propagation based on windows vs unix systems --- Cargo.lock | 30 ++++++++++++-- crates/uv/Cargo.toml | 4 ++ crates/uv/src/commands/project/run.rs | 60 +++++++++++++++++---------- crates/uv/src/commands/tool/run.rs | 60 +++++++++++++++++---------- 4 files changed, 106 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 846c8605d01a0..2cb3e9389e48c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -844,7 +844,7 @@ version = "3.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "90eeab0aa92f3f9b4e87f258c72b139c207d251f9cbc1080a0086b86a8870dd3" dependencies = [ - "nix", + "nix 0.29.0", "windows-sys 0.59.0", ] @@ -1419,7 +1419,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5bdbbd5bc8c5749697ccaa352fa45aff8730cf21c68029c0eef1ffed7c3d6ba2" dependencies = [ "cfg-if", - "nix", + "nix 0.29.0", "widestring", "windows 0.57.0", ] @@ -1939,6 +1939,15 @@ dependencies = [ "libc", ] +[[package]] +name = "memoffset" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" +dependencies = [ + "autocfg", +] + [[package]] name = "miette" version = "7.2.0" @@ -2053,6 +2062,19 @@ dependencies = [ "rand", ] +[[package]] +name = "nix" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f3790c00a0150112de0f4cd161e3d7fc4b2d8a5542ffc35f099a2562aecb35c" +dependencies = [ + "bitflags 1.3.2", + "cc", + "cfg-if", + "libc", + "memoffset", +] + [[package]] name = "nix" version = "0.29.0" @@ -4183,6 +4205,7 @@ dependencies = [ "axoupdater", "base64 0.22.1", "byteorder", + "cfg-if", "clap", "console", "ctrlc", @@ -4200,6 +4223,7 @@ dependencies = [ "itertools 0.13.0", "jiff", "miette", + "nix 0.23.2", "owo-colors", "petgraph", "predicates", @@ -5623,7 +5647,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/crates/uv/Cargo.toml b/crates/uv/Cargo.toml index b930238b32f0f..eea238b5712d3 100644 --- a/crates/uv/Cargo.toml +++ b/crates/uv/Cargo.toml @@ -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" } @@ -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", diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index b8310af272304..134818f2d0ca0 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -7,12 +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::select; -use tokio::signal::unix::{signal, SignalKind}; +use tokio::process::{Child, Command}; use tracing::{debug, warn}; use url::Url; use uv_cache::Cache; @@ -996,30 +995,29 @@ pub(crate) async fn run( // Ignore signals in the parent process, deferring them to the child. This is safe as long as // the command is the last thing that runs in this process; otherwise, we'd need to restore the // signal handlers after the command completes. - 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); - } + let _handler = tokio::spawn(async { while tokio::signal::ctrl_c().await.is_ok() {} }); + + #[cfg(unix)] + { + use tokio::select; + use tokio::signal::unix::{signal, SignalKind}; + let mut term_signal = signal(SignalKind::terminate())?; + loop { + select! { + _ = handle.wait() => { + break + }, - // `SIGINT` - _ = int_signal.recv() => { - handle.kill().await?; - handle.wait().await.context("Child process disappeared")?; - return Ok(ExitStatus::Failure); + // `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) { @@ -1038,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>, diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs index f06c3a5c9f719..85f578567d4c8 100644 --- a/crates/uv/src/commands/tool/run.rs +++ b/crates/uv/src/commands/tool/run.rs @@ -1,3 +1,4 @@ +use cfg_if::cfg_if; use std::fmt::Display; use std::fmt::Write; use std::path::PathBuf; @@ -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}; @@ -236,30 +236,29 @@ pub(crate) async fn run( // Ignore signals in the parent process, deferring them to the child. This is safe as long as // the command is the last thing that runs in this process; otherwise, we'd need to restore the // signal handlers after the command completes. - let mut term_signal = signal(SignalKind::terminate())?; - let mut int_signal = signal(SignalKind::interrupt())?; + let _handler = tokio::spawn(async { while tokio::signal::ctrl_c().await.is_ok() {} }); - 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) { @@ -278,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,