From a44fd24523bce226b3ac6a72fcdb3e8f9282f7b6 Mon Sep 17 00:00:00 2001
From: Klemens Morgenstern <klemens.d.morgenstern@gmail.com>
Date: Fri, 25 Oct 2024 11:46:58 +0800
Subject: [PATCH] unified cancellation on process.async_wait()

---
 .../boost/process/v2/detail/process_handle_fd.hpp  |  1 +
 .../v2/detail/process_handle_fd_or_signal.hpp      |  9 ++++++++-
 .../process/v2/detail/process_handle_signal.hpp    |  1 +
 .../process/v2/detail/process_handle_windows.hpp   | 14 ++++++++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/include/boost/process/v2/detail/process_handle_fd.hpp b/include/boost/process/v2/detail/process_handle_fd.hpp
index ba12c7f7..7a94dd0d 100644
--- a/include/boost/process/v2/detail/process_handle_fd.hpp
+++ b/include/boost/process/v2/detail/process_handle_fd.hpp
@@ -299,6 +299,7 @@ struct basic_process_handle_fd
         template<typename Self>
         void operator()(Self &&self)
         {
+            self.reset_cancellation_state(asio::enable_total_cancellation());
             error_code ec;
             native_exit_code_type exit_code{};
             int wait_res = -1;
diff --git a/include/boost/process/v2/detail/process_handle_fd_or_signal.hpp b/include/boost/process/v2/detail/process_handle_fd_or_signal.hpp
index 9dc6d93c..97d9400a 100644
--- a/include/boost/process/v2/detail/process_handle_fd_or_signal.hpp
+++ b/include/boost/process/v2/detail/process_handle_fd_or_signal.hpp
@@ -329,8 +329,15 @@ struct basic_process_handle_fd_or_signal
         pid_type pid_;
         bool needs_post = true;
 
+        template<typename Self.
+        void operator()(Self && self)
+        {
+          self.reset_cancellation_state(asio::enable_total_cancellation());
+          (*this)(std::move(self), error_code{});
+        }
+
         template<typename Self>
-        void operator()(Self &&self, error_code ec = {}, int = 0)
+        void operator()(Self &&self, error_code ec, int = 0)
         {
             native_exit_code_type exit_code{};
             int wait_res = -1;
diff --git a/include/boost/process/v2/detail/process_handle_signal.hpp b/include/boost/process/v2/detail/process_handle_signal.hpp
index 0c27a547..ca486b34 100644
--- a/include/boost/process/v2/detail/process_handle_signal.hpp
+++ b/include/boost/process/v2/detail/process_handle_signal.hpp
@@ -292,6 +292,7 @@ struct basic_process_handle_signal
         template<typename Self>
         void operator()(Self &&self)
         {
+            self.reset_cancellation_state(asio::enable_total_cancellation());
             handle.async_wait(std::move(self));
             handle.cancel();
             // we cancel so we end up on the signal-sets executor 
diff --git a/include/boost/process/v2/detail/process_handle_windows.hpp b/include/boost/process/v2/detail/process_handle_windows.hpp
index 274caf7b..e92b8653 100644
--- a/include/boost/process/v2/detail/process_handle_windows.hpp
+++ b/include/boost/process/v2/detail/process_handle_windows.hpp
@@ -288,6 +288,17 @@ struct basic_process_handle_win
         template<typename Self>
         void operator()(Self &&self)
         {
+
+            self.reset_cancellation_state(asio::enable_total_cancellation());
+            auto sl = self.get_cancellation_state().slot();
+            auto & h = handle;
+            if (sl.is_connected())
+                sl.assign(
+                    [&h](asio::cancellation_type ct)
+                    {
+                      boost::system::error_code ec;
+                      h.cancel(ec);
+                    });
             handle.async_wait(std::move(self));
         }
 
@@ -295,6 +306,9 @@ struct basic_process_handle_win
         void operator()(Self &&self, error_code ec)
         {
             native_exit_code_type exit_code{};
+            if (ec == asio::error::operation_aborted && !self.get_cancellation_state().cancelled())
+              return handle.async_wait(std::move(self));
+
             if (!ec)
                 detail::get_exit_code_(handle.native_handle(), exit_code, ec);
             std::move(self).complete(ec, exit_code);