diff --git a/include/rest_rpc/rpc_client.hpp b/include/rest_rpc/rpc_client.hpp index 582638e..5167177 100644 --- a/include/rest_rpc/rpc_client.hpp +++ b/include/rest_rpc/rpc_client.hpp @@ -20,22 +20,31 @@ template struct call_result { R value; }; -template <> struct call_result { rpc_errc ec; }; +template <> struct call_result { + rpc_errc ec; +}; class rpc_client { public: - rpc_client() : socket_(get_global_executor()) {} + rpc_client() : socket_(std::make_shared(get_global_executor())) {} + ~rpc_client() { close(); } - auto get_executor() { return socket_.get_executor(); } + auto get_executor() { return socket_->get_executor(); } asio::awaitable 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) || - resolver.async_resolve( - host, port, asio::as_tuple(asio::use_awaitable))); + asio::ip::tcp::resolver resolver(socket_->get_executor()); + + 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( - watchdog(duration) || - socket_.async_connect(endpoint, asio::as_tuple(asio::use_awaitable))); + auto conn_r = co_await (watchdog(duration) || + socket_->impl_.async_connect( + 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(); using R = function_return_type_t; - auto r = co_await(watchdog(duration) || - call_impl(header, std::forward(args)...)); + auto r = co_await (watchdog(duration) || + call_impl(header, std::forward(args)...)); if (r.index() == 0) { co_return call_result{rpc_errc::request_timeout}; } @@ -119,24 +130,24 @@ public: MD5::MD5Hash32(topic.data(), (uint32_t)topic.size()); // topic id bool b = false; call_result ret{}; - auto it = sub_ops_.find(topic_id); - if (it == sub_ops_.end()) { + auto it = socket_->sub_ops_.find(topic_id); + 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{rpc_errc::duplicate_topic}; } - std::tie(b, ret) = co_await( + std::tie(b, ret) = co_await ( asio::async_compose( std::ref(it->second), asio::use_awaitable) && call_impl(header)); } else { - std::tie(b, ret) = co_await( + std::tie(b, ret) = co_await ( asio::async_compose( std::ref(it->second), asio::use_awaitable) && wait_response()); @@ -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 @@ -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) { rpc_service::msgpack_codec codec; - result.value = codec.unpack( - std::string_view(body_.data() + 1, resp_header.body_len - 1)); + result.value = codec.unpack(std::string_view( + 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); - it != sub_ops_.end()) { + if (auto it = socket_->sub_ops_.find(resp_header.function_id); + it != socket_->sub_ops_.end()) { it->second.complete(true); } } @@ -227,9 +251,12 @@ private: } asio::awaitable 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 complete_handler_; }; - tcp_socket socket_; - std::string body_; + struct socket_t { + socket_t(auto executor) : impl_(executor) {} + asio::any_io_executor get_executor() { return impl_.get_executor(); } + asio::ip::tcp::socket impl_; + std::atomic has_closed_ = true; + std::string body_; + std::unordered_map 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_; bool tcp_no_delay_ = true; bool cross_ending_ = false; - std::unordered_map sub_ops_; + bool should_reset_ = false; }; -} // namespace rest_rpc \ No newline at end of file +} // namespace rest_rpc diff --git a/tests/test_rest_rpc.cpp b/tests/test_rest_rpc.cpp index 4460800..fd2dd9c 100644 --- a/tests/test_rest_rpc.cpp +++ b/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)