From b2264864c97a4f0510cae1e3caef00435237ca9c Mon Sep 17 00:00:00 2001
From: magick93 <kurrent93@gmail.com>
Date: Tue, 10 Dec 2024 15:00:15 +1300
Subject: [PATCH] Windows shutdown (#68)

---
 .gitignore                        |  3 ++
 anchor/src/environment.rs         | 23 +++++++++++
 anchor/src/environment_windows.rs | 64 +++++++++++++++++++++++++++++++
 3 files changed, 90 insertions(+)
 create mode 100644 anchor/src/environment_windows.rs

diff --git a/.gitignore b/.gitignore
index b627be9d..4240a362 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,6 @@ perf.data*
 
 # VSCode
 /.vscode
+
+# cross
+/zcross
\ No newline at end of file
diff --git a/anchor/src/environment.rs b/anchor/src/environment.rs
index 716b1319..32eedccf 100644
--- a/anchor/src/environment.rs
+++ b/anchor/src/environment.rs
@@ -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;
 
@@ -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) {
diff --git a/anchor/src/environment_windows.rs b/anchor/src/environment_windows.rs
new file mode 100644
index 00000000..25ddfbe2
--- /dev/null
+++ b/anchor/src/environment_windows.rs
@@ -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),
+        }
+    }
+}