Skip to content

Commit

Permalink
Handle platform differences and propagate sigterm
Browse files Browse the repository at this point in the history
This commit now only handles sigterm explicitly, using tokio sigint
handling, and implements the different sigterm propagation based on
windows vs unix systems
  • Loading branch information
shaneikennedy committed Nov 8, 2024
1 parent 091ca23 commit 7dd4436
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 48 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
60 changes: 37 additions & 23 deletions crates/uv/src/commands/project/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand All @@ -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>,
Expand Down
60 changes: 38 additions & 22 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 @@ -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) {
Expand All @@ -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,
Expand Down

0 comments on commit 7dd4436

Please sign in to comment.