reset

include/rest_rpc/rpc_client.hpp

@@ -20,22 +20,31 @@ template <typename R> struct call_result {
   R value;
 };
 
-template <> struct call_result<void> { rpc_errc ec; };
+template <> struct call_result<void> {
+  rpc_errc ec;
+};
 
 class rpc_client {
 public:
-  rpc_client() : socket_(get_global_executor()) {}
+  rpc_client() : socket_(std::make_shared<socket_t>(get_global_executor())) {}
+  ~rpc_client() { close(); }
 
-  auto get_executor() { return socket_.get_executor(); }
+  auto get_executor() { return socket_->get_executor(); }
 
   asio::awaitable<std::error_code> connect(
       std::string_view host, std::string_view port,
       std::chrono::steady_clock::duration duration = std::chrono::seconds(5)) {
-    asio::ip::tcp::resolver resolver(socket_.get_executor());
+    if (should_reset_) {
+      reset();
+    } else {
+      should_reset_ = true;
+    }
 
-    auto r = co_await(watchdog(duration) ||
+    asio::ip::tcp::resolver resolver(socket_->get_executor());
-                      resolver.async_resolve(
+
-                          host, port, asio::as_tuple(asio::use_awaitable)));
+    auto r = co_await (watchdog(duration) ||
+                       resolver.async_resolve(
+                           host, port, asio::as_tuple(asio::use_awaitable)));
     if (r.index() == 0) {
       REST_LOG_ERROR << "resolve timeout";
       co_return make_error_code(rpc_errc::resolve_timeout);
@@ -54,9 +63,9 @@ public:
     }
 
     auto endpoint = it->endpoint();
-    auto conn_r = co_await(
+    auto conn_r = co_await (watchdog(duration) ||
-        watchdog(duration) ||
+                            socket_->impl_.async_connect(
-        socket_.async_connect(endpoint, asio::as_tuple(asio::use_awaitable)));
+                                endpoint, asio::as_tuple(asio::use_awaitable)));
     if (conn_r.index() == 0) {
       REST_LOG_ERROR << "connect timeout";
       co_return make_error_code(rpc_errc::connection_timeout);
@@ -68,8 +77,10 @@ public:
       co_return conn_ec;
     }
 
+    socket_->has_closed_ = false;
+
     if (tcp_no_delay_) {
-      socket_.set_option(asio::ip::tcp::no_delay(true));
+      socket_->impl_.set_option(asio::ip::tcp::no_delay(true));
     }
 
     co_return std::error_code{};
@@ -105,8 +116,8 @@ public:
     rest_rpc_header header{};
     header.function_id = get_key<func>();
     using R = function_return_type_t<decltype(func)>;
-    auto r = co_await(watchdog(duration) ||
+    auto r = co_await (watchdog(duration) ||
-                      call_impl<R>(header, std::forward<Args>(args)...));
+                       call_impl<R>(header, std::forward<Args>(args)...));
     if (r.index() == 0) {
       co_return call_result<R>{rpc_errc::request_timeout};
     }
@@ -119,24 +130,24 @@ public:
         MD5::MD5Hash32(topic.data(), (uint32_t)topic.size()); // topic id
     bool b = false;
     call_result<R> ret{};
-    auto it = sub_ops_.find(topic_id);
+    auto it = socket_->sub_ops_.find(topic_id);
-    if (it == sub_ops_.end()) {
+    if (it == socket_->sub_ops_.end()) {
       rest_rpc_header header{};
       header.msg_type = 1; // pub/sub
 
       header.function_id = topic_id;
-      auto [it, r] = sub_ops_.emplace(topic_id, sub_operation{});
+      auto [it, r] = socket_->sub_ops_.emplace(topic_id, sub_operation{});
       if (!r) {
         REST_LOG_ERROR << "subscribe duplicate topic";
         co_return call_result<R>{rpc_errc::duplicate_topic};
       }
 
-      std::tie(b, ret) = co_await(
+      std::tie(b, ret) = co_await (
           asio::async_compose<decltype(asio::use_awaitable), void(bool)>(
               std::ref(it->second), asio::use_awaitable) &&
           call_impl<R>(header));
     } else {
-      std::tie(b, ret) = co_await(
+      std::tie(b, ret) = co_await (
           asio::async_compose<decltype(asio::use_awaitable), void(bool)>(
               std::ref(it->second), asio::use_awaitable) &&
           wait_response<R>());
@@ -148,6 +159,15 @@ public:
   void enable_tcp_no_delay(bool r) { tcp_no_delay_ = r; }
 
   void enable_cross_ending(bool r) { cross_ending_ = r; }
+  bool has_closed() const { return socket_->has_closed_; }
+
+  void close() {
+    if (socket_ == nullptr || socket_->has_closed_)
+      return;
+
+    asio::dispatch(socket_->get_executor(),
+                   [socket = socket_] { close_socket(*socket); });
+  }
 
 private:
   template <typename R, typename... Args>
@@ -171,9 +191,10 @@ private:
     std::error_code ec;
     size_t size;
     std::tie(ec, size) = co_await asio::async_write(
-        socket_, buffers, asio::as_tuple(asio::use_awaitable));
+        socket_->impl_, buffers, asio::as_tuple(asio::use_awaitable));
     if (ec) {
       result.ec = rpc_errc::write_error;
+      close_socket(*socket_);
       co_return result;
     }
 
@@ -186,10 +207,11 @@ private:
     size_t size;
     rest_rpc_header resp_header;
     std::tie(ec, size) = co_await asio::async_read(
-        socket_, asio::buffer(&resp_header, sizeof(rest_rpc_header)),
+        socket_->impl_, asio::buffer(&resp_header, sizeof(rest_rpc_header)),
         asio::as_tuple(asio::use_awaitable));
     if (ec) {
       result.ec = rpc_errc::write_error;
+      close_socket(*socket_);
       co_return result;
     }
     if (resp_header.magic != 39) {
@@ -201,25 +223,27 @@ private:
       parse_recieved(resp_header);
     }
 
-    detail::resize(body_, resp_header.body_len);
+    detail::resize(socket_->body_, resp_header.body_len);
     std::tie(ec, size) = co_await asio::async_read(
-        socket_, asio::buffer(body_.data(), body_.size()),
+        socket_->impl_,
+        asio::buffer(socket_->body_.data(), socket_->body_.size()),
         asio::as_tuple(asio::use_awaitable));
     if (ec) {
       REST_LOG_WARNING << "read body error: " << ec.message();
       result.ec = rpc_errc::read_error;
+      close_socket(*socket_);
       co_return result;
     }
-    result.ec = (rpc_errc)body_[0];
+    result.ec = (rpc_errc)socket_->body_[0];
     if constexpr (!std::is_void_v<R>) {
       rpc_service::msgpack_codec codec;
-      result.value = codec.unpack<R>(
+      result.value = codec.unpack<R>(std::string_view(
-          std::string_view(body_.data() + 1, resp_header.body_len - 1));
+          socket_->body_.data() + 1, resp_header.body_len - 1));
     }
 
     if (resp_header.msg_type == 1) { // pubsub
-      if (auto it = sub_ops_.find(resp_header.function_id);
+      if (auto it = socket_->sub_ops_.find(resp_header.function_id);
-          it != sub_ops_.end()) {
+          it != socket_->sub_ops_.end()) {
         it->second.complete(true);
       }
     }
@@ -227,9 +251,12 @@ private:
   }
 
   asio::awaitable<std::error_code> watchdog(auto duration) {
-    asio::steady_timer timer(socket_.get_executor());
+    asio::steady_timer timer(socket_->get_executor());
     timer.expires_after(duration);
     auto [ec] = co_await timer.async_wait(asio::as_tuple(asio::use_awaitable));
+    if (!ec) {
+      close_socket(*socket_);
+    }
     co_return ec;
   }
 
@@ -250,10 +277,43 @@ private:
     std::function<void(bool)> complete_handler_;
   };
 
-  tcp_socket socket_;
+  struct socket_t {
-  std::string body_;
+    socket_t(auto executor) : impl_(executor) {}
+    asio::any_io_executor get_executor() { return impl_.get_executor(); }
+    asio::ip::tcp::socket impl_;
+    std::atomic<bool> has_closed_ = true;
+    std::string body_;
+    std::unordered_map<uint32_t, sub_operation> sub_ops_;
+  };
+
+  inline static void close_socket(socket_t &socket) {
+    std::error_code ec;
+    socket.impl_.shutdown(asio::ip::tcp::socket::shutdown_both, ec);
+    socket.impl_.close(ec);
+    socket.has_closed_ = true;
+  }
+
+  void reset() {
+    auto executor = socket_->get_executor();
+    if (!has_closed()) {
+      close_socket(*socket_);
+    }
+
+    socket_->impl_ = asio::ip::tcp::socket{executor};
+    if (!socket_->impl_.is_open()) {
+      std::error_code ec;
+      socket_->impl_.open(asio::ip::tcp::v4(), ec);
+      if (ec) {
+        REST_LOG_WARNING << "client reset socket failed, reason: "
+                         << ec.message();
+        return;
+      }
+    }
+  }
+
+  std::shared_ptr<socket_t> socket_;
   bool tcp_no_delay_ = true;
   bool cross_ending_ = false;
-  std::unordered_map<uint32_t, sub_operation> sub_ops_;
+  bool should_reset_ = false;
 };
-} // namespace rest_rpc
+} // namespace rest_rpc

tests/test_rest_rpc.cpp

@@ -320,6 +320,19 @@ TEST_CASE("test pub sub") {
   promise.get_future().wait();
 }
 
+TEST_CASE("test reconnect") {
+  rpc_server server("127.0.0.1:9004");
+  server.async_start();
+  rpc_server server1("127.0.0.1:9005");
+  server1.async_start();
+
+  rpc_client client{};
+  auto ec = sync_wait(get_global_executor(), client.connect("127.0.0.1:9004"));
+  CHECK(!ec);
+  ec = sync_wait(get_global_executor(), client.connect("127.0.0.1:9005"));
+  CHECK(!ec);
+}
+
 // doctest comments
 // 'function' : must be 'attribute' - see issue #182
 DOCTEST_MSVC_SUPPRESS_WARNING_WITH_PUSH(4007)