[Store] Serve Ping on a separate RPC server that clients discover on Connect

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<ReturnType, ErrorCode> invoke_rpc(
         Args&&... args);
 
+    void ConnectHeartbeatServer(const std::string& master_addr)
+        REQUIRES(connect_mutex_);
+
     template <auto ServiceMethod, typename ReturnType, typename... Args>
     [[nodiscard]] tl::expected<ReturnType, ErrorCode> 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<bool> use_heartbeat_client_{false};
     RpcDrainGuard rpc_drain_;
 
     // The client identification.

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;

mooncake-store/include/rpc_service.h

@@ -190,6 +190,10 @@ class WrappedMasterService {
 
     tl::expected<std::string, ErrorCode> ServiceReady();
 
+    // 0 when Ping is served only on the main RPC server.
+    tl::expected<uint32_t, ErrorCode> GetHeartbeatRpcPort();
+    void SetHeartbeatRpcPort(uint32_t port) { heartbeat_rpc_port_ = port; }
+
     [[nodiscard]] TieredStorageUsageSnapshot GetStorageUsageSnapshot() const;
 
     tl::expected<std::vector<TenantQuotaSnapshot>, ErrorCode>
@@ -346,9 +350,16 @@ class WrappedMasterService {
 
    private:
     MasterService master_service_;
+    std::atomic<uint32_t> 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

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<coro_rpc::coro_rpc_server> heartbeat_server;
+        if (master_config.heartbeat_rpc_port != 0) {
+            heartbeat_server = std::make_unique<coro_rpc::coro_rpc_server>(
+                /*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<bool> 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";

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<bool, ErrorCode> MasterClient::ExistKey(
     const std::string& object_key) {
     ScopedVLogTimer timer(1, "MasterClient::ExistKey");
@@ -963,10 +994,14 @@ tl::expected<PingResponse, ErrorCode> MasterClient::Ping() {
     std::shared_ptr<RpcClientPool::ClientPool> 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>(

mooncake-store/src/rpc_service.cpp

@@ -1471,6 +1471,10 @@ tl::expected<std::string, ErrorCode> WrappedMasterService::ServiceReady() {
     return GetMooncakeStoreVersion();
 }
 
+tl::expected<uint32_t, ErrorCode> 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

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)