Cleanup

include/ligo/detail/lifecycle.hpp

@@ -9,6 +9,8 @@
 #include <vector>
 
 #include <boost/asio/as_tuple.hpp>
+#include <boost/asio/co_spawn.hpp>
+#include <boost/asio/strand.hpp>
 #include <boost/asio/awaitable.hpp>
 #include <boost/asio/error.hpp>
 #include <boost/asio/post.hpp>
@@ -40,6 +42,7 @@ struct session_registry {
   auto cancel_all() -> void;
   auto empty() const -> std::bool_t;
   auto wait_empty(std::chrono::milliseconds timeout) -> boost::asio::awaitable<std::expected<void, modulus::error>>;
+  auto wait_registered(std::shared_ptr<boost::asio::steady_timer> timer, std::chrono::steady_clock::time_point deadline) -> boost::asio::awaitable<std::expected<void, modulus::error>>;
   auto notify_waiters() -> void;
   auto live_waiters_locked() -> std::vector<std::shared_ptr<boost::asio::steady_timer>>;
   mutable std::mutex mutex{};
@@ -74,7 +77,8 @@ auto ligo::detail::cancel_session_stream(boost::asio::ssl::stream<NEXT_LAYER> &s
 template <typename S>
 auto ligo::detail::session_registry<S>::cancel_waiters(std::vector<std::shared_ptr<boost::asio::steady_timer>> const &live) -> void {
   for (auto &waiter : live) {
-    waiter->cancel();
+    // Timer operations belong to the waiter strand, including registration.
+    boost::asio::post(waiter->get_executor(), [waiter] { waiter->cancel(); });
   }
 }
 
@@ -125,30 +129,28 @@ auto ligo::detail::session_registry<S>::wait_empty(std::chrono::milliseconds tim
   if (empty()) {
     co_return std::expected<void, modulus::error>{};
   }
-  auto notify_timer = std::make_shared<boost::asio::steady_timer>(executor);
-  auto deadline_timer = std::make_shared<boost::asio::steady_timer>(executor);
-  auto timed_out = std::make_shared<std::atomic<std::bool_t>>(false);
+  auto timer = std::make_shared<boost::asio::steady_timer>(boost::asio::make_strand(executor));
+  auto deadline = std::chrono::steady_clock::now() + timeout;
   {
     auto lock = std::scoped_lock{mutex};
-    waiters.push_back(notify_timer);
+    waiters.push_back(timer);
   }
-  deadline_timer->expires_after(timeout);
-  deadline_timer->async_wait([notify_timer, timed_out](boost::system::error_code const &error) {
-    if (!error) {
-      timed_out->store(true, std::memory_order_relaxed);
-      notify_timer->cancel();
-    }
-  });
+  co_return co_await boost::asio::co_spawn(timer->get_executor(), wait_registered(timer, deadline), boost::asio::use_awaitable);
+}
+
+template <typename S>
+auto ligo::detail::session_registry<S>::wait_registered(std::shared_ptr<boost::asio::steady_timer> timer, std::chrono::steady_clock::time_point deadline) -> boost::asio::awaitable<std::expected<void, modulus::error>> {
+  // This entire coroutine runs on the timer strand. Notifications cannot run
+  // between checking session state and initiating the asynchronous wait.
   for (;;) {
     if (empty()) {
-      deadline_timer->cancel();
       co_return std::expected<void, modulus::error>{};
     }
-    if (timed_out->load(std::memory_order_relaxed)) {
+    if (std::chrono::steady_clock::now() >= deadline) {
       co_return ligo::detail::fail_void(ligo::detail::error_timeout, "timed out waiting for sessions to drain");
     }
-    notify_timer->expires_at(std::chrono::steady_clock::time_point::max());
-    auto [error] = co_await notify_timer->async_wait(boost::asio::as_tuple(boost::asio::use_awaitable));
+    timer->expires_at(deadline);
+    auto [error] = co_await timer->async_wait(boost::asio::as_tuple(boost::asio::use_awaitable));
     if (error && error != boost::asio::error::operation_aborted) {
       co_return ligo::detail::fail_void(ligo::detail::network_failure(error));
     }