Skip to content

Commit

Permalink
windows cancellation
Browse files Browse the repository at this point in the history
  • Loading branch information
erhant committed Sep 6, 2024
1 parent 47e81ae commit fcc881d
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 15 deletions.
2 changes: 1 addition & 1 deletion Cargo.lock

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

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "dkn-compute"
version = "0.2.0"
version = "0.2.1"
edition = "2021"
license = "Apache-2.0"
readme = "README.md"
Expand Down
55 changes: 42 additions & 13 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let cancellation_token = token.clone();
// add cancellation check
tokio::spawn(async move {
// FIXME: weird feature-gating here bugs with IDE, fix this later
#[cfg(feature = "profiling")]
{
const PROFILE_DURATION_SECS: u64 = 120;
Expand Down Expand Up @@ -86,22 +87,50 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(())
}

/// Waits for SIGTERM or SIGINT, and cancels the given token when the signal is received.
// FIXME: remove this `unused` once we have a better way to handle this
/// Waits for various termination signals, and cancels the given token when the signal is received.
///
/// Handles Unix and Windows [target families](https://doc.rust-lang.org/reference/conditional-compilation.html#target_family).
#[allow(unused)]
async fn wait_for_termination(cancellation: CancellationToken) -> std::io::Result<()> {
use tokio::signal::unix::{signal, SignalKind};
#[cfg(unix)]
{
use tokio::signal::unix::{signal, SignalKind};
let mut sigterm = signal(SignalKind::terminate())?; // Docker sends SIGTERM
let mut sigint = signal(SignalKind::interrupt())?; // Ctrl+C sends SIGINT
tokio::select! {
_ = sigterm.recv() => log::warn!("Recieved SIGTERM"),
_ = sigint.recv() => log::warn!("Recieved SIGINT"),
_ = cancellation.cancelled() => {
// no need to wait if cancelled anyways
// although this is not likely to happen
return Ok(());
}
};
}

let mut sigterm = signal(SignalKind::terminate())?; // Docker sends SIGTERM
let mut sigint = signal(SignalKind::interrupt())?; // Ctrl+C sends SIGINT
tokio::select! {
_ = sigterm.recv() => log::warn!("Recieved SIGTERM"),
_ = sigint.recv() => log::warn!("Recieved SIGINT"),
_ = cancellation.cancelled() => {
// no need to wait if cancelled anyways
// although this is not likely to happen
return Ok(());
}
};
#[cfg(windows)]
{
use tokio::signal::windows;

// https://learn.microsoft.com/en-us/windows/console/handlerroutine
let mut signal_c = windows::ctrl_c()?;
let mut signal_break = windows::ctrl_break()?;
let mut signal_close = windows::ctrl_close()?;
let mut signal_shutdown = windows::ctrl_shutdown()?;

tokio::select! {
_ = signal_c.recv() => log::warn!("Received CTRL_C"),
_ = signal_break.recv() => log::warn!("Received CTRL_BREAK"),
_ = signal_close.recv() => log::warn!("Received CTRL_CLOSE"),
_ = signal_shutdown.recv() => log::warn!("Received CTRL_SHUTDOWN"),
_ = cancellation.cancelled() => {
// no need to wait if cancelled anyways
// although this is not likely to happen
return Ok(());
}
};
}

log::info!("Terminating the node...");
cancellation.cancel();
Expand Down

0 comments on commit fcc881d

Please sign in to comment.