Skip to content

Commit

Permalink
Windows shutdown (#68)
Browse files Browse the repository at this point in the history
  • Loading branch information
magick93 authored Dec 10, 2024
1 parent 04176ae commit b226486
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 0 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,6 @@ perf.data*

# VSCode
/.vscode

# cross
/zcross
23 changes: 23 additions & 0 deletions anchor/src/environment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@ use {
tokio::signal::unix::{signal, Signal, SignalKind},
};

#[cfg(target_family = "windows")]
#[path = "environment_windows.rs"]
mod environment_windows;

/// The maximum time in seconds the client will wait for all internal tasks to shutdown.
const MAXIMUM_SHUTDOWN_TIME: u64 = 15;

Expand Down Expand Up @@ -141,6 +145,25 @@ impl Environment {
}
}

#[cfg(target_family = "windows")]
pub fn block_until_shutdown_requested(&mut self) -> Result<ShutdownReason, String> {
let signal_rx = self
.signal_rx
.take()
.ok_or("Inner shutdown already received")?;

match self
.runtime()
.block_on(environment_windows::handle_shutdown_signals(signal_rx))
{
Ok(reason) => {
info!(reason = reason.message(), "Internal shutdown received");
Ok(reason)
}
Err(e) => Err(e),
}
}

/// Shutdown the `tokio` runtime when all tasks are idle.
pub fn shutdown_on_idle(self) {
match Arc::try_unwrap(self.runtime) {
Expand Down
64 changes: 64 additions & 0 deletions anchor/src/environment_windows.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
use futures::channel::mpsc::Receiver;
use futures::{future, Future, StreamExt};
use std::{pin::Pin, task::Context, task::Poll};
use task_executor::ShutdownReason;
use tokio::signal::windows::{ctrl_c, CtrlC};
use tracing::error;

pub(crate) async fn handle_shutdown_signals(
mut signal_rx: Receiver<ShutdownReason>,
) -> Result<ShutdownReason, String> {
let inner_shutdown = async move {
signal_rx
.next()
.await
.ok_or("Internal shutdown channel exhausted")
};
futures::pin_mut!(inner_shutdown);

let register_handlers = async {
let mut handles = vec![];

// Setup for handling Ctrl+C
match ctrl_c() {
Ok(ctrl_c) => {
let ctrl_c = SignalFuture::new(ctrl_c, "Received Ctrl+C");
handles.push(ctrl_c);
}
Err(e) => error!(error = ?e, "Could not register Ctrl+C handler"),
}

future::select(inner_shutdown, future::select_all(handles.into_iter())).await
};

match register_handlers.await {
future::Either::Left((Ok(reason), _)) => Ok(reason),
future::Either::Left((Err(e), _)) => Err(e.into()),
future::Either::Right(((res, _, _), _)) => {
res.ok_or_else(|| "Handler channel closed".to_string())
}
}
}

struct SignalFuture {
signal: CtrlC,
message: &'static str,
}

impl SignalFuture {
pub fn new(signal: CtrlC, message: &'static str) -> SignalFuture {
SignalFuture { signal, message }
}
}

impl Future for SignalFuture {
type Output = Option<ShutdownReason>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.signal.poll_recv(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Some(_)) => Poll::Ready(Some(ShutdownReason::Success(self.message))),
Poll::Ready(None) => Poll::Ready(None),
}
}
}

0 comments on commit b226486

Please sign in to comment.