handle conn lifetime

include/rest_rpc/rest_rpc_server.hpp

@@ -11,7 +11,8 @@ public:
   rest_rpc_server(std::string address,
                   size_t num_thread = std::thread::hardware_concurrency())
       : io_context_pool_(num_thread),
-        acceptor_(io_context_pool_.get_io_context()) {
+        acceptor_(io_context_pool_.get_io_context()),
+        check_timer_(io_context_pool_.get_io_context()) {
     size_t pos = address.find(':');
     if (pos != std::string::npos) {
       host_ = address.substr(0, pos);
@@ -23,13 +24,27 @@ public:
                   size_t num_thread = std::thread::hardware_concurrency())
       : io_context_pool_(num_thread),
         acceptor_(io_context_pool_.get_io_context()), host_(std::move(host)),
-        port_(std::move(port)) {}
+        port_(std::move(port)),
+        check_timer_(io_context_pool_.get_io_context()) {}
   ~rest_rpc_server() { stop(); }
 
   std::error_code start() { return start_impl(false); }
 
   std::error_code async_start() { return start_impl(true); }
 
+  void set_conn_max_age(std::chrono::steady_clock::duration dur) {
+    if (dur > std::chrono::steady_clock::duration::zero()) {
+      need_check_ = true;
+      timeout_duration_ = dur;
+      asio::co_spawn(check_timer_.get_executor(), start_check_timer(),
+                     asio::detached);
+    }
+  }
+
+  void set_check_conn_interval(std::chrono::steady_clock::duration dur) {
+    check_duration_ = dur;
+  }
+
   void stop() {
     if (has_stop_.load(std::memory_order_acquire)) {
       return;
@@ -43,6 +58,16 @@ public:
         (void)acceptor_.close(ec);
       });
 
+      stop_timer_ = true;
+
+      {
+        std::scoped_lock lock(*conn_mtx_);
+        for (auto &conn : conns_) {
+          conn.second->close(false);
+        }
+        conns_.clear();
+      }
+
       io_context_pool_.stop();
     });
 
@@ -72,6 +97,11 @@ public:
 
   void enable_cross_ending(bool r) { cross_ending_ = r; }
 
+  size_t connection_count() {
+    std::scoped_lock lock(*conn_mtx_);
+    return conns_.size();
+  }
+
 private:
   std::error_code listen() {
     using asio::ip::tcp;
@@ -136,7 +166,24 @@ private:
       REST_LOG_INFO << "new connction comming...";
       auto conn = std::make_shared<rpc_connection>(std::move(socket), conn_id,
                                                    router_, cross_ending_);
-      conns_.emplace(conn_id++, conn);
+      if (need_check_) {
+        conn->set_check_timeout(true);
+      }
+      std::weak_ptr<std::mutex> weak(conn_mtx_);
+      conn->set_quit_callback([this, weak](const uint64_t &id) {
+        auto mtx = weak.lock();
+        if (mtx) {
+          std::scoped_lock lock(*mtx);
+          if (!conns_.empty())
+            conns_.erase(id);
+        }
+      });
+
+      {
+        std::scoped_lock lock(*conn_mtx_);
+        conns_.emplace(conn_id++, conn);
+      }
+
       co_spawn(socket.get_executor(), conn->start(), asio::detached);
     }
   }
@@ -167,6 +214,35 @@ private:
     return ec;
   }
 
+  asio::awaitable<void> check_timeout() {
+    auto cur_time = std::chrono::system_clock::now();
+
+    {
+      std::scoped_lock lock(*conn_mtx_);
+      for (auto it = conns_.begin(); it != conns_.end();) {
+        if (cur_time - co_await it->second->get_last_rwtime() >
+            timeout_duration_) {
+          it->second->close(false);
+          conns_.erase(it++);
+        } else {
+          ++it;
+        }
+      }
+    }
+  }
+
+  asio::awaitable<void> start_check_timer() {
+    while (true) {
+      check_timer_.expires_after(check_duration_);
+      auto [ec] =
+          co_await check_timer_.async_wait(asio::as_tuple(asio::use_awaitable));
+      if (ec || stop_timer_) {
+        co_return;
+      }
+      co_await check_timeout();
+    }
+  }
+
   io_context_pool io_context_pool_;
   std::thread thd_;
   asio::ip::tcp::acceptor acceptor_;
@@ -176,6 +252,16 @@ private:
   std::once_flag stop_flag_;
   std::atomic<bool> has_stop_ = false;
   std::unordered_map<uint64_t, std::shared_ptr<rpc_connection>> conns_;
+
+  std::shared_ptr<std::mutex> conn_mtx_ = std::make_shared<std::mutex>();
+  std::chrono::steady_clock::duration check_duration_ =
+      std::chrono::seconds(12);
+  std::chrono::steady_clock::duration timeout_duration_{
+      std::chrono::seconds(10)};
+  asio::steady_timer check_timer_;
+  std::atomic<bool> stop_timer_ = false;
+  bool need_check_ = false;
+
   rpc_router router_;
   bool tcp_no_delay_ = true;
   bool cross_ending_ = false;

include/rest_rpc/rpc_connection.hpp

@@ -46,12 +46,14 @@ public:
       : socket_(std::move(socket)), conn_id_(conn_id), router_(router),
         cross_ending_(cross_ending) {}
 
+  ~rpc_connection() { close(); }
   asio::awaitable<void> start() {
     rest_rpc_header header;
     auto self = this->shared_from_this();
     while (true) {
       std::error_code ec;
       size_t size;
+      set_last_time();
       std::tie(ec, size) = co_await asio::async_read(
           socket_, asio::buffer(&header, sizeof(rest_rpc_header)),
           asio::as_tuple(asio::use_awaitable));
@@ -72,6 +74,7 @@ public:
       detail::resize(body_, header.body_len);
 
       if (header.body_len > 0) {
+        set_last_time();
         std::tie(ec, size) = co_await asio::async_read(
             socket_, asio::buffer(body_), asio::as_tuple(asio::use_awaitable));
         if (ec) {
@@ -111,18 +114,63 @@ public:
     if (!result.empty())
       buffers.push_back(asio::buffer(result.data()));
 
+    set_last_time();
     auto [ec, size] = co_await asio::async_write(
         socket_, buffers, asio::as_tuple(asio::use_awaitable));
+    if (ec) {
+      REST_LOG_WARNING << "write error: " << ec.message();
+      close();
+    }
     co_return ec;
   }
 
   uint64_t id() const { return conn_id_; }
   auto get_executor() { return socket_.get_executor(); }
 
+  void
+  set_quit_callback(std::function<void(const uint64_t &conn_id)> callback) {
+    quit_cb_ = std::move(callback);
+  }
+
+  void close(bool need_cb = true) {
+    if (has_closed_) {
+      return;
+    }
+
+    asio::dispatch(socket_.get_executor(),
+                   [this, need_cb, self = shared_from_this()] {
+                     std::error_code ec;
+                     socket_.shutdown(asio::socket_base::shutdown_both, ec);
+                     socket_.close(ec);
+                     if (need_cb && quit_cb_) {
+                       quit_cb_(conn_id_);
+                     }
+                     has_closed_ = true;
+                   });
+  }
+
+  void set_last_time() {
+    if (checkout_timeout_) {
+      last_rwtime_ = std::chrono::system_clock::now();
+    }
+  }
+
+  asio::awaitable<std::chrono::system_clock::time_point> get_last_rwtime() {
+    co_await asio::this_coro::executor;
+    co_return last_rwtime_;
+  }
+
+  void set_check_timeout(bool r) { checkout_timeout_ = r; }
+
 private:
   tcp_socket socket_;
   uint64_t conn_id_;
   std::string body_;
+  std::function<void(const uint64_t &conn_id)> quit_cb_ = nullptr;
+  std::atomic<bool> has_closed_{false};
+  std::chrono::system_clock::time_point last_rwtime_ =
+      std::chrono::system_clock::now();
+  bool checkout_timeout_ = false;
   rpc_router &router_;
   bool cross_ending_;
 };