diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 624748e4b0..15caa2ab05 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -101,6 +101,8 @@ class MasterClient { ha_probe_client_accessor_( GetStoreRpcClientIoContextPool(), detail::MakeMasterRpcClientPoolConfig(/*ha_enabled=*/true)), + heartbeat_client_accessor_(GetStoreRpcClientIoContextPool(), + detail::MakeMasterRpcClientPoolConfig()), client_id_(client_id), tenant_id_(std::move(tenant_id)), metrics_(metrics) {} @@ -702,6 +704,9 @@ class MasterClient { [[nodiscard]] tl::expected invoke_rpc( Args&&... args); + void ConnectHeartbeatServer(const std::string& master_addr) + REQUIRES(connect_mutex_); + template [[nodiscard]] tl::expected invoke_rpc_with_pool( RpcClientPool& client_accessor, Args&&... args); @@ -728,6 +733,9 @@ class MasterClient { RpcClientPool client_accessor_; RpcClientPool ha_control_client_accessor_; RpcClientPool ha_probe_client_accessor_; + // Set on Connect when the master serves Ping on a separate server. + RpcClientPool heartbeat_client_accessor_; + std::atomic use_heartbeat_client_{false}; RpcDrainGuard rpc_drain_; // The client identification. diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index 18d11cccfe..caced3612c 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -101,6 +101,7 @@ struct MasterConfig { std::string rpc_interface; int32_t rpc_conn_timeout_seconds; bool rpc_enable_tcp_no_delay; + uint32_t heartbeat_rpc_port = 0; uint64_t default_kv_lease_ttl; uint64_t default_kv_soft_pin_ttl; diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index 322b8a2095..062779b3bf 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -190,6 +190,10 @@ class WrappedMasterService { tl::expected ServiceReady(); + // 0 when Ping is served only on the main RPC server. + tl::expected GetHeartbeatRpcPort(); + void SetHeartbeatRpcPort(uint32_t port) { heartbeat_rpc_port_ = port; } + [[nodiscard]] TieredStorageUsageSnapshot GetStorageUsageSnapshot() const; tl::expected, ErrorCode> @@ -346,9 +350,16 @@ class WrappedMasterService { private: MasterService master_service_; + std::atomic heartbeat_rpc_port_{0}; }; void RegisterRpcService(coro_rpc::coro_rpc_server& server, mooncake::WrappedMasterService& wrapped_master_service); +// For a server whose io threads serve nothing but heartbeats, so a long +// handler on the main server cannot keep a Ping from being read. +void RegisterHeartbeatRpcService( + coro_rpc::coro_rpc_server& server, + mooncake::WrappedMasterService& wrapped_master_service); + } // namespace mooncake diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index 7a081ba03d..166c487037 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -183,6 +183,10 @@ DEFINE_int32(rpc_conn_timeout_seconds, 0, "Connection timeout in seconds (0 = no timeout)"); DEFINE_bool(rpc_enable_tcp_no_delay, true, "Enable TCP_NODELAY for RPC connections"); +DEFINE_uint32(heartbeat_rpc_port, 0, + "Port of a separate RPC server that also serves Ping, so a long " + "handler on the main server cannot delay heartbeats; clients " + "discover it on Connect (0 = disabled, non-HA mode only)"); DEFINE_validator(eviction_ratio, [](const char* flagname, double value) { if (value < 0.0 || value > 1.0) { LOG(FATAL) << "Mem eviction ratio must be between 0.0 and 1.0"; @@ -512,6 +516,9 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, default_config.GetBool("rpc_enable_tcp_no_delay", &master_config.rpc_enable_tcp_no_delay, FLAGS_rpc_enable_tcp_no_delay); + default_config.GetUInt32("heartbeat_rpc_port", + &master_config.heartbeat_rpc_port, + FLAGS_heartbeat_rpc_port); default_config.GetDurationMs("default_kv_lease_ttl", &master_config.default_kv_lease_ttl, mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL); @@ -858,6 +865,11 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, !conf_set) { master_config.rpc_enable_tcp_no_delay = FLAGS_rpc_enable_tcp_no_delay; } + if ((google::GetCommandLineFlagInfo("heartbeat_rpc_port", &info) && + !info.is_default) || + !conf_set) { + master_config.heartbeat_rpc_port = FLAGS_heartbeat_rpc_port; + } if ((google::GetCommandLineFlagInfo("enable_metric_reporting", &info) && !info.is_default) || !conf_set) { @@ -1683,6 +1695,7 @@ int main(int argc, char* argv[]) { << ", rpc_conn_timeout_seconds=" << master_config.rpc_conn_timeout_seconds << ", rpc_enable_tcp_no_delay=" << master_config.rpc_enable_tcp_no_delay + << ", heartbeat_rpc_port=" << master_config.heartbeat_rpc_port << ", rpc protocol=" << protocol << ", cluster_id=" << master_config.cluster_id << ", root_fs_dir=" << master_config.root_fs_dir @@ -1751,6 +1764,8 @@ int main(int argc, char* argv[]) { } if (master_config.enable_ha) { + LOG_IF(WARNING, master_config.heartbeat_rpc_port != 0) + << "heartbeat_rpc_port is not supported in HA mode and is ignored"; mooncake::MasterServiceSupervisorConfig supervisor_config{ master_config}; supervisor_config.http_metadata_server = metadata_server_ptr; @@ -1786,6 +1801,27 @@ int main(int argc, char* argv[]) { mooncake::RegisterRpcService(server, *wrapped_master_service); + std::unique_ptr heartbeat_server; + if (master_config.heartbeat_rpc_port != 0) { + heartbeat_server = std::make_unique( + /*thread_num=*/1, master_config.heartbeat_rpc_port, + master_config.rpc_address, + std::chrono::seconds(master_config.rpc_conn_timeout_seconds), + master_config.rpc_enable_tcp_no_delay); + if (mooncake::RpcProtocolConfig::FromEnvironment().use_rdma) { + heartbeat_server->init_ibv(); + } + mooncake::RegisterHeartbeatRpcService(*heartbeat_server, + *wrapped_master_service); + if (auto ec = heartbeat_server->async_start(); ec.hasResult()) { + LOG(ERROR) << "Failed to start heartbeat RPC server on port " + << master_config.heartbeat_rpc_port; + return 1; + } + wrapped_master_service->SetHeartbeatRpcPort( + master_config.heartbeat_rpc_port); + } + static std::atomic shutdown_requested{false}; auto signal_handler = [](int /* signum */) { shutdown_requested.store(true); @@ -1806,6 +1842,9 @@ int main(int argc, char* argv[]) { server.stop(); server_thread.join(); + if (heartbeat_server) { + heartbeat_server->stop(); + } if (shutdown_requested.load()) { LOG(INFO) << "Shutdown signal received, exiting gracefully"; diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index f0c06f1381..f65a70e066 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -192,6 +192,11 @@ struct RpcNameTraits<&WrappedMasterService::Ping> { static constexpr const char* value = "Ping"; }; +template <> +struct RpcNameTraits<&WrappedMasterService::GetHeartbeatRpcPort> { + static constexpr const char* value = "GetHeartbeatRpcPort"; +}; + template <> struct RpcNameTraits<&WrappedMasterService::GetFsdir> { static constexpr const char* value = "GetFsdir"; @@ -477,6 +482,7 @@ ErrorCode MasterClient::Connect(const std::string& master_addr) { timer.LogRequest("master_addr=", master_addr); MutexLocker lock(&connect_mutex_); + use_heartbeat_client_.store(false); const bool use_ha_control_pool = ha_connection_policy_enabled_.load(std::memory_order_acquire); if (use_ha_control_pool) { @@ -511,11 +517,36 @@ ErrorCode MasterClient::Connect(const std::string& master_addr) { client_accessor_.GetOrCreateClientPool(master_addr); ha_control_client_accessor_.GetOrCreateClientPool(master_addr); } + ConnectHeartbeatServer(master_addr); client_addr_param_ = master_addr; timer.LogResponse("error_code=", ErrorCode::OK); return ErrorCode::OK; } +void MasterClient::ConnectHeartbeatServer(const std::string& master_addr) { + // An older master does not have this RPC; Ping then stays on the main + // server, which always serves it. + auto port = + invoke_rpc<&WrappedMasterService::GetHeartbeatRpcPort, uint32_t>(); + const auto host_end = master_addr.rfind(':'); + if (!port.has_value() || port.value() == 0 || + host_end == std::string::npos) { + return; + } + const std::string heartbeat_addr = + master_addr.substr(0, host_end + 1) + std::to_string(port.value()); + heartbeat_client_accessor_.GetOrCreateClientPool(heartbeat_addr); + auto ready = + invoke_rpc_with_pool<&WrappedMasterService::ServiceReady, std::string>( + heartbeat_client_accessor_); + if (!ready.has_value()) { + LOG(WARNING) << "Heartbeat RPC server " << heartbeat_addr + << " is unreachable, pinging the main server instead"; + return; + } + use_heartbeat_client_.store(true); +} + tl::expected MasterClient::ExistKey( const std::string& object_key) { ScopedVLogTimer timer(1, "MasterClient::ExistKey"); @@ -963,10 +994,14 @@ tl::expected MasterClient::Ping() { std::shared_ptr ping_pool; { MutexLocker lock(&connect_mutex_); - ping_pool = - ha_connection_policy_enabled_.load(std::memory_order_acquire) - ? ha_control_client_accessor_.GetClientPool() - : client_accessor_.GetClientPool(); + if (use_heartbeat_client_.load()) { + ping_pool = heartbeat_client_accessor_.GetClientPool(); + } else if (ha_connection_policy_enabled_.load( + std::memory_order_acquire)) { + ping_pool = ha_control_client_accessor_.GetClientPool(); + } else { + ping_pool = client_accessor_.GetClientPool(); + } } auto result = invoke_rpc_with_client_pool<&WrappedMasterService::Ping, PingResponse>( diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index b6e1cc1b89..3b7b3df511 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -1471,6 +1471,10 @@ tl::expected WrappedMasterService::ServiceReady() { return GetMooncakeStoreVersion(); } +tl::expected WrappedMasterService::GetHeartbeatRpcPort() { + return heartbeat_rpc_port_.load(); +} + TieredStorageUsageSnapshot WrappedMasterService::GetStorageUsageSnapshot() const { return master_service_.GetStorageUsageSnapshot(); @@ -1839,6 +1843,9 @@ void RegisterRpcService( &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::ServiceReady>( &wrapped_master_service); + server + .register_handler<&mooncake::WrappedMasterService::GetHeartbeatRpcPort>( + &wrapped_master_service); server.register_handler< &mooncake::WrappedMasterService::MountLocalDiskSegment>( &wrapped_master_service); @@ -1906,4 +1913,13 @@ void RegisterRpcService( &wrapped_master_service); } +void RegisterHeartbeatRpcService( + coro_rpc::coro_rpc_server& server, + mooncake::WrappedMasterService& wrapped_master_service) { + server.register_handler<&mooncake::WrappedMasterService::Ping>( + &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::ServiceReady>( + &wrapped_master_service); +} + } // namespace mooncake diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 92ea94b500..d7b93f0c81 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -283,6 +283,7 @@ add_ha_test(batch_oplog_snapshot_rebootstrap_test add_store_test(master_service_test_for_snapshot ha/snapshot/master_service_test_for_snapshot.cpp) add_store_test(non_ha_reconnect_test non_ha_reconnect_test.cpp) +add_store_test(heartbeat_rpc_server_test heartbeat_rpc_server_test.cpp) add_store_test(nvme_kv_connector_config_test nvme_kv/connector_config_test.cpp) add_store_test(ha_cluster_namespace_config_test ha_cluster_namespace_config_test.cpp) diff --git a/mooncake-store/tests/heartbeat_rpc_server_test.cpp b/mooncake-store/tests/heartbeat_rpc_server_test.cpp new file mode 100644 index 0000000000..d66c8e0ae3 --- /dev/null +++ b/mooncake-store/tests/heartbeat_rpc_server_test.cpp @@ -0,0 +1,187 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "common/network.h" +#include "master_client.h" +#include "rpc_service.h" +#include "types.h" + +namespace mooncake { +namespace { + +std::atomic main_io_thread_held{false}; +std::shared_future main_io_thread_release; + +// Stands in for any long handler: it keeps the io thread that read it. +void HoldMainIoThread() { + main_io_thread_held.store(true); + main_io_thread_release.wait(); +} + +class HeartbeatRpcServerTest : public ::testing::Test { + protected: + explicit HeartbeatRpcServerTest(size_t main_io_thread_num = 1) + : main_io_thread_num_(main_io_thread_num) {} + + void SetUp() override { + const auto ports = getFreeTcpPorts(2); + ASSERT_EQ(ports.size(), 2u); + main_port_ = ports[0]; + heartbeat_port_ = ports[1]; + + WrappedMasterServiceConfig config; + config.default_kv_lease_ttl = DEFAULT_DEFAULT_KV_LEASE_TTL; + config.enable_metric_reporting = false; + wrapped_ = std::make_unique(config); + main_server_ = std::make_unique( + main_io_thread_num_, main_port_, "127.0.0.1", + std::chrono::seconds(0), /*tcp_no_delay=*/true); + RegisterRpcService(*main_server_, *wrapped_); + main_server_->register_handler(); + ASSERT_FALSE(main_server_->async_start().hasResult()); + + heartbeat_server_ = std::make_unique( + /*thread_num=*/1, heartbeat_port_, "127.0.0.1", + std::chrono::seconds(0), /*tcp_no_delay=*/true); + RegisterHeartbeatRpcService(*heartbeat_server_, *wrapped_); + ASSERT_FALSE(heartbeat_server_->async_start().hasResult()); + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + + void TearDown() override { + ReleaseMainIoThread(); + if (hold_caller_.joinable()) { + hold_caller_.join(); + } + if (heartbeat_server_) { + heartbeat_server_->stop(); + } + if (main_server_) { + main_server_->stop(); + } + } + + std::string MainAddress() const { + return "127.0.0.1:" + std::to_string(main_port_); + } + + void HoldTheMainIoThread() { + main_io_thread_held.store(false); + release_ = std::promise(); + main_io_thread_release = release_.get_future().share(); + hold_caller_ = std::thread([port = main_port_] { + coro_rpc::coro_rpc_client client; + async_simple::coro::syncAwait( + client.connect("127.0.0.1", std::to_string(port))); + async_simple::coro::syncAwait(client.call()); + }); + while (!main_io_thread_held.load()) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + } + + void ReleaseMainIoThread() { + if (!released_) { + released_ = true; + release_.set_value(); + } + } + + const size_t main_io_thread_num_; + int main_port_ = 0; + int heartbeat_port_ = 0; + std::unique_ptr wrapped_; + std::unique_ptr main_server_; + std::unique_ptr heartbeat_server_; + std::promise release_; + bool released_ = false; + std::thread hold_caller_; +}; + +class HeartbeatRpcServerTwoMainIoThreadsTest : public HeartbeatRpcServerTest { + protected: + HeartbeatRpcServerTwoMainIoThreadsTest() + : HeartbeatRpcServerTest(/*main_io_thread_num=*/2) {} +}; + +TEST_F(HeartbeatRpcServerTest, PingIsServedWhileTheMainIoThreadIsBusy) { + wrapped_->SetHeartbeatRpcPort(heartbeat_port_); + MasterClient client(generate_uuid()); + ASSERT_EQ(client.Connect(MainAddress()), ErrorCode::OK); + + HoldTheMainIoThread(); + auto ping = std::async(std::launch::async, [&] { return client.Ping(); }); + const bool answered_while_busy = + ping.wait_for(std::chrono::seconds(2)) == std::future_status::ready; + ReleaseMainIoThread(); + + EXPECT_TRUE(answered_while_busy); + EXPECT_TRUE(ping.get().has_value()); +} + +TEST_F(HeartbeatRpcServerTest, PingWaitsForTheMainIoThreadWithoutTheServer) { + MasterClient client(generate_uuid()); + ASSERT_EQ(client.Connect(MainAddress()), ErrorCode::OK); + + HoldTheMainIoThread(); + auto ping = std::async(std::launch::async, [&] { return client.Ping(); }); + const bool answered_while_busy = + ping.wait_for(std::chrono::seconds(2)) == std::future_status::ready; + ReleaseMainIoThread(); + + EXPECT_FALSE(answered_while_busy); + EXPECT_TRUE(ping.get().has_value()); +} + +TEST_F(HeartbeatRpcServerTest, PingFallsBackToTheMainServerWhenUnreachable) { + wrapped_->SetHeartbeatRpcPort(getFreeTcpPorts(1).at(0)); + MasterClient client(generate_uuid()); + ASSERT_EQ(client.Connect(MainAddress()), ErrorCode::OK); + + EXPECT_TRUE(client.Ping().has_value()); +} + +TEST_F(HeartbeatRpcServerTwoMainIoThreadsTest, + PingWaitsForABusyIoThreadEvenWhenAnotherIsIdle) { + coro_rpc::coro_rpc_client first; + coro_rpc::coro_rpc_client second; + for (auto* client : {&first, &second}) { + ASSERT_FALSE(async_simple::coro::syncAwait( + client->connect("127.0.0.1", std::to_string(main_port_)))); + } + + // Connections take io threads round-robin, so the hold's connection, + // opened right after these two, shares its thread with exactly one. + HoldTheMainIoThread(); + auto ping_on = [](coro_rpc::coro_rpc_client& client) { + return std::async(std::launch::async, [&client] { + return async_simple::coro::syncAwait( + client.call<&WrappedMasterService::Ping>(generate_uuid())); + }); + }; + auto first_ping = ping_on(first); + auto second_ping = ping_on(second); + auto answered_while_busy = [](auto& ping) { + return ping.wait_for(std::chrono::seconds(2)) == + std::future_status::ready; + }; + const int pings_answered_while_busy = + answered_while_busy(first_ping) + answered_while_busy(second_ping); + ReleaseMainIoThread(); + + EXPECT_EQ(pings_answered_while_busy, 1); + EXPECT_TRUE(first_ping.get().has_value()); + EXPECT_TRUE(second_ping.get().has_value()); +} + +} // namespace +} // namespace mooncake