ligo/include/ligo/detail/lifecycle.hpp

#pragma once
#include <expected>

#include <atomic>
#include <chrono>
#include <memory>
#include <mutex>
#include <set>
#include <vector>

#include <boost/asio/as_tuple.hpp>
#include <boost/asio/awaitable.hpp>
#include <boost/asio/error.hpp>
#include <boost/asio/post.hpp>
#include <boost/asio/steady_timer.hpp>
#include <boost/asio/this_coro.hpp>
#include <boost/asio/use_awaitable.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/ssl/stream.hpp>
#include <boost/beast/core.hpp>
#include <boost/system/error_code.hpp>

#include <modulus.hpp>

#include <ligo/detail/errors.hpp>
#include <ligo/detail/transport.hpp>

namespace ligo::detail {

auto cancel_session_stream(boost::beast::tcp_stream& stream) -> void;
template <typename NEXT_LAYER>
auto cancel_session_stream(boost::asio::ssl::stream<NEXT_LAYER>& stream) -> void;

template <typename S>
struct session_registry {
  static auto cancel_waiters(std::vector<std::shared_ptr<boost::asio::steady_timer>> const& live) -> void;
  auto add(std::shared_ptr<S> const& stream) -> void;
  auto remove(std::shared_ptr<S> const& stream) -> void;
  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 notify_waiters() -> void;
  auto live_waiters_locked() -> std::vector<std::shared_ptr<boost::asio::steady_timer>>;
  mutable std::mutex mutex{};
  std::set<std::shared_ptr<S>> sessions{};
  std::vector<std::weak_ptr<boost::asio::steady_timer>> waiters{};
};

template <typename S>
struct session_guard {
  ~session_guard();
  auto release() -> void;
  std::shared_ptr<ligo::detail::session_registry<S>> registry{};
  std::shared_ptr<S> stream{};
};

template <typename S>
auto stop_accepting_server(std::shared_ptr<S> state) -> void;
template <typename S>
auto stop_accepting_server_now(std::shared_ptr<S> const& state) -> void;
template <typename S>
auto async_stop_accepting_server(std::shared_ptr<S> state, std::chrono::milliseconds timeout) -> boost::asio::awaitable<std::expected<void, modulus::error>>;

}

template <typename NEXT_LAYER>
auto ligo::detail::cancel_session_stream(boost::asio::ssl::stream<NEXT_LAYER>& stream) -> void {
  auto ignored = boost::system::error_code{};
  stream.lowest_layer().cancel(ignored);
  stream.lowest_layer().close(ignored);
}

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();
  }
}

template <typename S>
auto ligo::detail::session_registry<S>::add(std::shared_ptr<S> const& stream) -> void {
  auto lock = std::scoped_lock{mutex};
  sessions.insert(stream);
}

template <typename S>
auto ligo::detail::session_registry<S>::remove(std::shared_ptr<S> const& stream) -> void {
  auto live = std::vector<std::shared_ptr<boost::asio::steady_timer>>{};
  {
    auto lock = std::scoped_lock{mutex};
    sessions.erase(stream);
    live = live_waiters_locked();
  }
  cancel_waiters(live);
}

template <typename S>
auto ligo::detail::session_registry<S>::cancel_all() -> void {
  auto copy = std::vector<std::shared_ptr<S>>{};
  {
    auto lock = std::scoped_lock{mutex};
    copy.assign(sessions.begin(), sessions.end());
  }
  for (auto& stream : copy) {
    boost::asio::post(stream->get_executor(), [stream] {
      ligo::detail::cancel_session_stream(*stream);
    });
  }
}

template <typename S>
auto ligo::detail::session_registry<S>::empty() const -> std::bool_t {
  auto lock = std::scoped_lock{mutex};
  return sessions.empty();
}

template <typename S>
auto ligo::detail::session_registry<S>::wait_empty(std::chrono::milliseconds timeout) -> boost::asio::awaitable<std::expected<void, modulus::error>> {
  auto valid_timeout = ligo::detail::validate_positive_timeout(timeout, "shutdown timeout");
  if (!valid_timeout) {
    co_return valid_timeout;
  }
  auto executor = co_await boost::asio::this_coro::executor;
  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 lock = std::scoped_lock{mutex};
    waiters.push_back(notify_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();
    }
  });
  for (;;) {
    if (empty()) {
      deadline_timer->cancel();
      co_return std::expected<void, modulus::error>{};
    }
    if (timed_out->load(std::memory_order_relaxed)) {
      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));
    if (error && error != boost::asio::error::operation_aborted) {
      co_return ligo::detail::fail_void(ligo::detail::network_failure(error));
    }
  }
}

template <typename S>
auto ligo::detail::session_registry<S>::notify_waiters() -> void {
  auto live = std::vector<std::shared_ptr<boost::asio::steady_timer>>{};
  {
    auto lock = std::scoped_lock{mutex};
    live = live_waiters_locked();
  }
  cancel_waiters(live);
}

template <typename S>
auto ligo::detail::session_registry<S>::live_waiters_locked() -> std::vector<std::shared_ptr<boost::asio::steady_timer>> {
  auto live = std::vector<std::shared_ptr<boost::asio::steady_timer>>{};
  auto retained = std::vector<std::weak_ptr<boost::asio::steady_timer>>{};
  for (auto& waiter : waiters) {
    if (auto locked = waiter.lock()) {
      live.push_back(locked);
      retained.push_back(locked);
    }
  }
  waiters = std::move(retained);
  return live;
}

template <typename S>
ligo::detail::session_guard<S>::~session_guard() {
  if (registry && stream) {
    registry->remove(stream);
  }
}

template <typename S>
auto ligo::detail::session_guard<S>::release() -> void {
  if (registry && stream) {
    registry->remove(stream);
    stream.reset();
    registry.reset();
  }
}

template <typename S>
auto ligo::detail::stop_accepting_server(std::shared_ptr<S> state) -> void {
  auto lock = std::scoped_lock{state->lifecycle_mutex};
  if (state->stopped.exchange(true)) {
    return;
  }
  state->sessions->cancel_all();
  auto server_executor = state->strand;
  boost::asio::post(server_executor, [state = std::move(state)] {
    auto ignored = boost::system::error_code{};
    state->acceptor->close(ignored);
  });
}

template <typename S>
auto ligo::detail::stop_accepting_server_now(std::shared_ptr<S> const& state) -> void {
  auto lock = std::scoped_lock{state->lifecycle_mutex};
  if (state->stopped.exchange(true)) {
    return;
  }
  state->sessions->cancel_all();
  auto ignored = boost::system::error_code{};
  state->acceptor->close(ignored);
}

template <typename S>
auto ligo::detail::async_stop_accepting_server(std::shared_ptr<S> state, std::chrono::milliseconds timeout) -> boost::asio::awaitable<std::expected<void, modulus::error>> {
  auto valid_timeout = ligo::detail::validate_positive_timeout(timeout, "shutdown timeout");
  if (!valid_timeout) {
    co_return valid_timeout;
  }
  auto should_stop = false;
  {
    auto lock = std::scoped_lock{state->lifecycle_mutex};
    should_stop = !state->stopped.exchange(true);
    if (should_stop) {
      state->sessions->cancel_all();
    }
  }
  if (should_stop) {
    co_await boost::asio::post(state->strand, boost::asio::use_awaitable);
    auto ignored = boost::system::error_code{};
    state->acceptor->close(ignored);
  }
  co_return co_await state->sessions->wait_empty(timeout);
}