diff --git a/src/client/clnt_flags.cc b/src/client/clnt_flags.cc index 25ea741..a6ebff4 100644 --- a/src/client/clnt_flags.cc +++ b/src/client/clnt_flags.cc @@ -10,3 +10,7 @@ DEFINE_uint32(clnt_cm_addr_check_interval_inSecs, DEFINE_int32(clnt_sync_req_timeout_ms, 1000, "simm client sync request timeout in milliseconds, default is 1s"); DEFINE_int32(clnt_async_req_timeout_ms, 3000, "simm client sync request timeout in milliseconds, default is 3s"); DEFINE_string(clnt_log_file, "/var/log/simm/simm_clnt.log", "simm client log file path & name"); +// client-side wait window before triggering ReInit on DS failure +DEFINE_uint32(clnt_deferred_reshard_wait_inSecs, 30, + "Client-side wait window before triggering ReInit on DS failure. " + "Should be less than CM's cm_deferred_reshard_window_inSecs (default 120s)"); diff --git a/src/client/clnt_messenger.cc b/src/client/clnt_messenger.cc index f912d73..77d4e29 100644 --- a/src/client/clnt_messenger.cc +++ b/src/client/clnt_messenger.cc @@ -43,6 +43,7 @@ DECLARE_uint32(clnt_syncreq_retry_count); DECLARE_bool(clnt_syncreq_enable_retry); DECLARE_bool(clnt_use_k8s); DECLARE_bool(simm_enable_trace); +DECLARE_uint32(clnt_deferred_reshard_wait_inSecs); DECLARE_LOG_MODULE("simm_client"); @@ -151,13 +152,47 @@ error_code_t ClientMessenger::Init() { } else { for (auto [addr, ds_ctx] : ds_conn_ctxs_) { if (!ds_ctx->active.load()) { - if (CommonErr::OK != build_connection(addr)) { - // FIXME(ytji): current behavior is rude to reconnect to all data servers when one or part of them - // have issues. The better action is sync with cm and get latest data servers address info, reconnect - // to servers which are new extended. + // track when DS was first seen as dead + { + std::lock_guard lg(ds_dead_since_mtx_); + if (!ds_dead_since_.count(addr)) { + ds_dead_since_[addr] = std::chrono::steady_clock::now(); + } + } + + // Try to reconnect — new DS may have come up with same port + if (CommonErr::OK == build_connection(addr)) { + std::lock_guard lg(ds_dead_since_mtx_); + ds_dead_since_.erase(addr); + continue; + } + + // Reconnect failed — check if deferred reshard wait window exceeded. + // NOTE: CM updates its routing table immediately upon DS handshake (IP update or + // replacement), but the client has no way to learn about it promptly because + // CM-to-client routing push (RPC_ROUTING_TABLE_UPDATE) is not yet implemented + // (see cm_service.cc TODO). Until push is available, the client can only discover + // the new IP by polling CM via ReInit() after this wait window expires. + // If the DS restarts with a different IP, IO to the affected shards will fail for + // up to clnt_deferred_reshard_wait_inSecs seconds. This is a known limitation; + // implementing CM→client push will eliminate the gap. + std::chrono::duration dur; + { + std::lock_guard lg(ds_dead_since_mtx_); + dur = std::chrono::steady_clock::now() - ds_dead_since_[addr]; + } + + if (dur > std::chrono::seconds(FLAGS_clnt_deferred_reshard_wait_inSecs)) { + // Window expired: CM may have done reshard or IP update, pull new routes + { + std::lock_guard lg(ds_dead_since_mtx_); + ds_dead_since_.erase(addr); + } should_reinit = true; - MLOG_ERROR("Failover thread failed to reconnect to {}, will trigger reinit", addr); + MLOG_ERROR("Failover thread: DS {} unreachable for {}s (> {}s window), triggering reinit", + addr, static_cast(dur.count()), FLAGS_clnt_deferred_reshard_wait_inSecs); } + // else: still within wait window, keep retrying next cycle } } } diff --git a/src/client/clnt_messenger.h b/src/client/clnt_messenger.h index 6ad1b79..b7beab0 100644 --- a/src/client/clnt_messenger.h +++ b/src/client/clnt_messenger.h @@ -217,6 +217,11 @@ class ClientMessenger { std::atomic failover_flag_{true}; std::mutex failover_mutex_; std::condition_variable failover_condv_; + + // track when each DS was first seen as dead, used for failover wait window + std::mutex ds_dead_since_mtx_; + std::unordered_map ds_dead_since_; + std::unique_ptr trace_server_{nullptr}; #if defined(SIMM_UNIT_TEST) diff --git a/src/cluster_manager/cm_flags.cc b/src/cluster_manager/cm_flags.cc index 2e73908..92e8919 100644 --- a/src/cluster_manager/cm_flags.cc +++ b/src/cluster_manager/cm_flags.cc @@ -26,6 +26,16 @@ DEFINE_uint32(cm_heartbeat_timeout_inSecs, "Timeout strategy : dataserver heartbeat timeout in seconds," "if no heartbeat received in this time, the dataserver will be considered to mark as dead"); +// Deferred Reshard related flags +DEFINE_bool(cm_deferred_reshard_enabled, + true, + "Enable Deferred Reshard: when a DS heartbeat times out, CM waits for a replacement DS " + "with the same logical_node_id before triggering reshard"); +DEFINE_uint32(cm_deferred_reshard_window_inSecs, + 120, + "Deferred Reshard window in seconds: if no replacement DS registers within this window, " + "CM falls back to standard reshard"); + // Node Manager related flags DEFINE_uint32(dataserver_resource_interval_inSecs, 60, diff --git a/src/cluster_manager/cm_hb_monitor.cc b/src/cluster_manager/cm_hb_monitor.cc index 3b6c954..ddb4fde 100644 --- a/src/cluster_manager/cm_hb_monitor.cc +++ b/src/cluster_manager/cm_hb_monitor.cc @@ -13,6 +13,9 @@ DECLARE_LOG_MODULE("cluster_manager"); DECLARE_uint32(cm_heartbeat_records_perserver); DECLARE_uint32(cm_heartbeat_bg_scan_interval_inSecs); DECLARE_uint32(cm_heartbeat_timeout_inSecs); +DECLARE_bool(cm_deferred_reshard_enabled); +DECLARE_uint32(cm_deferred_reshard_window_inSecs); +DECLARE_uint32(dataserver_min_num); namespace simm { namespace cm { @@ -48,6 +51,7 @@ error_code_t ClusterManagerHBMonitor::Stop() { return CommonErr::OK; } +// Legacy: heartbeat keyed by ip:port (backward compat) error_code_t ClusterManagerHBMonitor::OnRecvNodeHeartbeat(const std::string &node_addr_str) { simm::common::NodeHeartbeatTs hb_ts(std::chrono::steady_clock::now(), std::chrono::system_clock::now()); auto uomap_locked = ds_hb_records_.wlock(); @@ -65,6 +69,38 @@ error_code_t ClusterManagerHBMonitor::OnRecvNodeHeartbeat(const std::string &nod return CommonErr::OK; } +// heartbeat keyed by logical_node_id +error_code_t ClusterManagerHBMonitor::OnRecvNodeHeartbeat(const std::string &logical_node_id, + const std::string &ip_port) { + simm::common::NodeHeartbeatTs hb_ts(std::chrono::steady_clock::now(), std::chrono::system_clock::now()); + + // Store heartbeat keyed by logical_node_id (not ip:port) + auto uomap_locked = ds_hb_records_.wlock(); + auto &entry = (*uomap_locked)[logical_node_id]; + if (entry.size() >= FLAGS_cm_heartbeat_records_perserver) { + entry.pop_front(); + } + entry.push_back(hb_ts); + + // Sync NodeManager: update logical_id → ip mapping if needed + cm_node_manager_ptr_->OnHeartbeat(logical_node_id, ip_port); + + return CommonErr::OK; +} + +void ClusterManagerHBMonitor::OnDeferredReshardResolved(const std::string &logical_node_id) { + // Reset heartbeat records for this logical_id so the new DS starts fresh + auto uomap_locked = ds_hb_records_.wlock(); + auto it = uomap_locked->find(logical_node_id); + if (it != uomap_locked->end()) { + it->second.clear(); + // Push a fresh heartbeat so the new DS doesn't immediately time out + simm::common::NodeHeartbeatTs hb_ts(std::chrono::steady_clock::now(), std::chrono::system_clock::now()); + it->second.push_back(hb_ts); + } + MLOG_INFO("Node handshake resolved: logical_id={}, heartbeat records reset", logical_node_id); +} + void ClusterManagerHBMonitor::BgHBScanLoop() { while (!stop_flag_.load()) { std::vector dead_dataservers{}; @@ -74,28 +110,121 @@ void ClusterManagerHBMonitor::BgHBScanLoop() { for (auto it = uomap_locked->begin(); it != uomap_locked->end();) { auto &hb_records = it->second; if (hb_records.empty()) { - // it = uomap_locked->erase(it); MLOG_WARN("Dataserver node({}) has no heartbeat records, skip it.", it->first); ++it; continue; } - // TODO(ytji): add more mark dead strategies // Check the latest heartbeat record auto previous_hb_ts = hb_records.back(); if (now - previous_hb_ts.monotonic_tp_ > std::chrono::seconds(FLAGS_cm_heartbeat_timeout_inSecs)) { - MLOG_ERROR("Dataserver node({}) heartbeat timeout({} secs), mark it as DEAD", - it->first, - FLAGS_cm_heartbeat_timeout_inSecs); - if (cm_node_manager_ptr_->QueryNodeStatus(it->first) == NodeStatus::DEAD) { - ++it; - continue; + const std::string& hb_key = it->first; // logical_node_id or ip:port + + // Try to resolve logical_node_id for this key + // If the key is a logical_node_id, GetNodeEntry will find it directly. + // If the key is an ip:port (legacy), try reverse lookup. + auto entry_opt = cm_node_manager_ptr_->GetNodeEntry(hb_key); + std::string logical_id = hb_key; + if (!entry_opt) { + // Maybe hb_key is an ip:port — try reverse lookup + logical_id = cm_node_manager_ptr_->ResolveLogicalId(hb_key); + if (!logical_id.empty()) { + entry_opt = cm_node_manager_ptr_->GetNodeEntry(logical_id); + } + } + + if (entry_opt && FLAGS_cm_deferred_reshard_enabled) { + // logical_node_id available, deferred reshard enabled + auto& entry = *entry_opt; + + if (entry.status == NodeStatus::DEFERRED_RESHARD) { + // Already waiting — check if window expired + auto elapsed = now - entry.deferred_reshard_since; + if (elapsed > std::chrono::seconds(FLAGS_cm_deferred_reshard_window_inSecs)) { + // Window expired → fallback to standard reshard + // Pre-check: is reshard feasible (enough alive nodes after marking DEAD)? + // DEFERRED_RESHARD nodes are already excluded from GetAllNodeAddress(alive=true) + // since legacy map has them as non-RUNNING, so no need to subtract 1. + auto alive = cm_node_manager_ptr_->GetAllNodeAddress(true); + if (alive.size() >= FLAGS_dataserver_min_num) { + MLOG_WARN("Deferred window expired for logical_id={} ({}s), triggering reshard", + logical_id, FLAGS_cm_deferred_reshard_window_inSecs); + // CAS: only mark DEAD if still in DEFERRED_RESHARD (guards against concurrent + // ProcessHandshake that already moved it back to RUNNING) + auto ret = cm_node_manager_ptr_->SetNodeStatus( + logical_id, NodeStatus::DEAD, {}, NodeStatus::DEFERRED_RESHARD); + if (ret == CommonErr::OK) { + dead_dataservers.emplace_back(entry.current_ip_port); + } + } else { + // Not safe to reshard — keep waiting in DEFERRED_RESHARD + MLOG_WARN("Deferred window expired for logical_id={}, but alive={} < min_required={}. " + "Staying in DEFERRED_RESHARD to avoid orphaned shards.", + logical_id, alive.size(), FLAGS_dataserver_min_num); + } + } + // else: window not expired, keep waiting + + } else if (entry.status == NodeStatus::RUNNING) { + // First timeout → enter DEFERRED_RESHARD state + // CAS: only transition if still RUNNING (guards against concurrent ProcessHandshake + // that may have re-registered the node after we read the snapshot) + MLOG_WARN("Node heartbeat timeout: logical_id={} (ip={}), entering DEFERRED_RESHARD", + logical_id, entry.current_ip_port); + cm_node_manager_ptr_->SetNodeStatus( + logical_id, NodeStatus::DEFERRED_RESHARD, now, NodeStatus::RUNNING); + + } else if (entry.status == NodeStatus::DEAD) { + // Already dead, skip + } + + } else if (entry_opt && !FLAGS_cm_deferred_reshard_enabled) { + // logical_node_id available, deferred reshard disabled + // Immediate DEAD + reshard + auto& entry = *entry_opt; + if (entry.status == NodeStatus::DEAD) { + // Already dead, skip + } else { + // Pre-check: is reshard feasible? + auto alive = cm_node_manager_ptr_->GetAllNodeAddress(true); + size_t alive_after = alive.size() > 0 ? alive.size() - 1 : 0; + if (alive_after >= FLAGS_dataserver_min_num) { + MLOG_WARN("Node heartbeat timeout: logical_id={} (ip={}), marking DEAD", + logical_id, entry.current_ip_port); + // CAS: only mark DEAD if the node hasn't been concurrently restored + auto ret = cm_node_manager_ptr_->SetNodeStatus( + logical_id, NodeStatus::DEAD, {}, entry.status); + if (ret == CommonErr::OK) { + dead_dataservers.emplace_back(entry.current_ip_port); + } + } else { + // Not enough alive nodes for reshard — enter DEFERRED_RESHARD as safety net + MLOG_WARN("Node heartbeat timeout: logical_id={} (ip={}), alive_after={} < min={}. " + "Entering DEFERRED_RESHARD as safety net to avoid orphaned shards.", + logical_id, entry.current_ip_port, alive_after, FLAGS_dataserver_min_num); + cm_node_manager_ptr_->SetNodeStatus( + logical_id, NodeStatus::DEFERRED_RESHARD, now, entry.status); + } + } + } else { - // mark server as dead state - cm_node_manager_ptr_->UpdateNodeStatus(it->first, NodeStatus::DEAD); - // record all dead servers in current scan round to trigger shard manager update in batch - dead_dataservers.emplace_back(it->first); - // FIXME(ytji): we still keep the heartbeat records for the node, - // it = uomap_locked->erase(it); + // Legacy path (no logical_node_id): immediate reshard + if (cm_node_manager_ptr_->QueryNodeStatus(hb_key) == NodeStatus::DEAD) { + ++it; + continue; + } + // Pre-check: is reshard feasible? + auto alive = cm_node_manager_ptr_->GetAllNodeAddress(true); + size_t alive_after = alive.size() > 0 ? alive.size() - 1 : 0; + if (alive_after >= FLAGS_dataserver_min_num) { + MLOG_ERROR("Dataserver node({}) heartbeat timeout({} secs), mark it as DEAD", + hb_key, FLAGS_cm_heartbeat_timeout_inSecs); + cm_node_manager_ptr_->UpdateNodeStatus(hb_key, NodeStatus::DEAD); + dead_dataservers.emplace_back(hb_key); + } else { + MLOG_WARN("Dataserver node({}) heartbeat timeout, but alive_after={} < min={}. " + "Skipping DEAD to avoid orphaned shards.", + hb_key, alive_after, FLAGS_dataserver_min_num); + } } } ++it; @@ -119,6 +248,9 @@ void ClusterManagerHBMonitor::HandleNodeFailure(const std::vector & MLOG_WARN("Handling failure of {} nodes", dead_node_addresses.size()); auto alive_servers = cm_node_manager_ptr_->GetAllNodeAddress(true /* alive only */); + if (alive_servers.empty()) { + return; + } // rebalance shards after node failure error_code_t ret = cm_shard_manager_ptr_->RebalanceShardsAfterNodeFailure(dead_node_addresses, alive_servers); diff --git a/src/cluster_manager/cm_hb_monitor.h b/src/cluster_manager/cm_hb_monitor.h index 18c37ba..8e885d8 100644 --- a/src/cluster_manager/cm_hb_monitor.h +++ b/src/cluster_manager/cm_hb_monitor.h @@ -38,8 +38,16 @@ class ClusterManagerHBMonitor : public std::enable_shared_from_this>; folly::Synchronized ds_hb_records_; diff --git a/src/cluster_manager/cm_node_manager.cc b/src/cluster_manager/cm_node_manager.cc index 3a212da..4c9ada9 100644 --- a/src/cluster_manager/cm_node_manager.cc +++ b/src/cluster_manager/cm_node_manager.cc @@ -21,6 +21,8 @@ DECLARE_string(dataserver_svc_name); DECLARE_string(dataserver_port_name); DECLARE_uint32(rpc_timeout_inSecs); DECLARE_uint32(dataserver_resource_interval_inSecs); +DECLARE_bool(cm_deferred_reshard_enabled); +DECLARE_uint32(cm_deferred_reshard_window_inSecs); DECLARE_LOG_MODULE("cluster_manager"); @@ -206,5 +208,161 @@ std::unordered_map> return resource_map; } +// logical_node_id based methods + +void ClusterManagerNodeManager::migrateNodeIp(const std::string& logical_id, + NodeEntry& entry, + const std::string& new_ip_port) { + addr_to_logical_.erase(entry.current_ip_port); + addr_to_logical_.insert_or_assign(new_ip_port, logical_id); + node_status_map_.erase(entry.current_ip_port); + node_info_map_.erase(entry.current_ip_port); + AddNode(new_ip_port); + entry.current_ip_port = new_ip_port; +} + +HandshakeResult ClusterManagerNodeManager::ProcessHandshake( + const std::string& logical_id, + const std::string& new_ip_port, + const std::vector& reported_shards) { + + HandshakeResult result; + + auto it = logical_node_table_.find(logical_id); + if (it == logical_node_table_.end()) { + // Case 1: logical_id not seen before → new node registration + NodeEntry entry; + entry.logical_node_id = logical_id; + entry.current_ip_port = new_ip_port; + entry.status = NodeStatus::RUNNING; + logical_node_table_.insert_or_assign(logical_id, entry); + addr_to_logical_.insert_or_assign(new_ip_port, logical_id); + + // Also register in legacy ip:port maps + AddNode(new_ip_port); + + result.action = HandshakeResult::Action::NEW_NODE; + result.shards_to_assign = reported_shards; + MLOG_INFO("New node registered: logical_id={} ip={}", logical_id, new_ip_port); + return result; + } + + // ConcurrentHashMap iterators yield const refs — copy, mutate, assign back + NodeEntry entry = it->second; + + switch (entry.status) { + case NodeStatus::DEFERRED_RESHARD: { + // Case 2: replacement DS registered while waiting — in-place IP update + result.action = HandshakeResult::Action::DEFERRED_RESHARD_REPLACE; + result.old_ip_port = entry.current_ip_port; + migrateNodeIp(logical_id, entry, new_ip_port); + entry.status = NodeStatus::RUNNING; + entry.deferred_reshard_since = {}; + logical_node_table_.assign(logical_id, entry); + MLOG_INFO("Node replacement: logical_id={} old_ip={} new_ip={}", + logical_id, result.old_ip_port, new_ip_port); + return result; + } + + case NodeStatus::RUNNING: { + result.action = HandshakeResult::Action::IP_UPDATE; + result.old_ip_port = entry.current_ip_port; + if (entry.current_ip_port == new_ip_port) { + // Case 3: DS restarted before HB timeout, same IP. Treat as IP_UPDATE so the + // handler looks up shards via GetShardsOwnedByNode and returns them to the DS. + // BatchAssignRoutingTable with same addr is a no-op on the routing table. + MLOG_INFO("Node re-registration (fast restart): logical_id={} ip={}", logical_id, new_ip_port); + } else { + // Case 4: IP changed while still RUNNING (rare IP drift) + migrateNodeIp(logical_id, entry, new_ip_port); + logical_node_table_.assign(logical_id, entry); + MLOG_INFO("Node IP update: logical_id={} old_ip={} new_ip={}", + logical_id, result.old_ip_port, new_ip_port); + } + return result; + } + + case NodeStatus::DEAD: + default: { + // Case 5: node was DEAD (reshard already happened), treat as new + migrateNodeIp(logical_id, entry, new_ip_port); + entry.status = NodeStatus::RUNNING; + entry.deferred_reshard_since = {}; + logical_node_table_.assign(logical_id, entry); + result.action = HandshakeResult::Action::NEW_NODE; + result.shards_to_assign = reported_shards; + MLOG_INFO("Node rejoin after DEAD: logical_id={} ip={}", logical_id, new_ip_port); + return result; + } + } +} + +std::optional ClusterManagerNodeManager::GetNodeEntry(const std::string& logical_id) const { + auto it = logical_node_table_.find(logical_id); + if (it == logical_node_table_.end()) { + return std::nullopt; + } + return it->second; +} + +error_code_t ClusterManagerNodeManager::OnHeartbeat(const std::string& logical_id, + const std::string& ip_port) { + auto it = logical_node_table_.find(logical_id); + if (it != logical_node_table_.end()) { + auto entry = it->second; + if (entry.current_ip_port != ip_port) { + migrateNodeIp(logical_id, entry, ip_port); + logical_node_table_.assign(logical_id, entry); + } + } + return CommonErr::OK; +} + +error_code_t ClusterManagerNodeManager::SetNodeStatus( + const std::string& logical_id, + NodeStatus status, + std::chrono::steady_clock::time_point ts, + std::optional expected_status) { + + auto it = logical_node_table_.find(logical_id); + if (it == logical_node_table_.end()) { + MLOG_WARN("SetNodeStatus: logical_id={} not found", logical_id); + return CommonErr::TargetNotFound; + } + + auto entry = it->second; + + // Compare-and-set: if caller specified an expected status, abort if it no + // longer matches (e.g. ProcessHandshake already moved the node back to RUNNING). + if (expected_status.has_value() && entry.status != expected_status.value()) { + MLOG_WARN("SetNodeStatus: logical_id={} status mismatch (expected={} actual={}), skipping", + logical_id, + common::NodeStatusToString(expected_status.value()), + common::NodeStatusToString(entry.status)); + return CommonErr::OK; + } + + entry.status = status; + if (status == NodeStatus::DEFERRED_RESHARD) { + entry.deferred_reshard_since = ts; + } + logical_node_table_.assign(logical_id, entry); + + // Keep legacy map in sync + UpdateNodeStatus(entry.current_ip_port, status); + + MLOG_INFO("SetNodeStatus: logical_id={} ip={} status={}", + logical_id, entry.current_ip_port, common::NodeStatusToString(status)); + return CommonErr::OK; +} + +std::string ClusterManagerNodeManager::ResolveLogicalId(const std::string& ip_port) const { + auto it = addr_to_logical_.find(ip_port); + if (it == addr_to_logical_.end()) { + return ""; + } + return it->second; +} + } // namespace cm } // namespace simm diff --git a/src/cluster_manager/cm_node_manager.h b/src/cluster_manager/cm_node_manager.h index 38e1dc8..e38321f 100644 --- a/src/cluster_manager/cm_node_manager.h +++ b/src/cluster_manager/cm_node_manager.h @@ -1,7 +1,9 @@ #pragma once #include +#include #include +#include #include #include #include @@ -27,6 +29,22 @@ namespace cm { using NodeStatus = common::NodeStatus; +// Per-logical-node entry tracking stable identity and current physical address +struct NodeEntry { + std::string logical_node_id; // unique key: "namespace/pod-name" in K8s, or --ds_logical_node_id flag value + std::string current_ip_port; // current physical address (may change on restart) + NodeStatus status{NodeStatus::UNKNOWN}; + std::chrono::steady_clock::time_point deferred_reshard_since; // when DEFERRED_RESHARD started +}; + +// Handshake processing result +struct HandshakeResult { + enum class Action { NEW_NODE, IP_UPDATE, DEFERRED_RESHARD_REPLACE, REJECT }; + Action action{Action::REJECT}; + std::string old_ip_port; // previous IP (for DEFERRED_RESHARD_REPLACE / IP_UPDATE) + std::vector shards_to_assign; // shards CM decides to assign +}; + class ClusterManagerNodeManager : public std::enable_shared_from_this { public: /** @@ -77,10 +95,42 @@ class ClusterManagerNodeManager : public std::enable_shared_from_this> GetAllNodeResource(); + // logical_node_id based interfaces + + // Process a Handshake from DS with logical_node_id. + // CM looks up the logical_id in its state table and decides the action. + HandshakeResult ProcessHandshake(const std::string& logical_id, + const std::string& new_ip_port, + const std::vector& reported_shards); + + // Get a node entry by logical_node_id (for HBMonitor state checks) + std::optional GetNodeEntry(const std::string& logical_id) const; + + // Called on each heartbeat: update logical_id → ip_port mapping + error_code_t OnHeartbeat(const std::string& logical_id, const std::string& ip_port); + + // Set node status by logical_node_id (used by HBMonitor). + // If expected_status is provided, the update is only applied when the current + // status matches expected_status (compare-and-set), preventing stale-snapshot + // races between BgHBScanLoop and concurrent ProcessHandshake calls. + error_code_t SetNodeStatus(const std::string& logical_id, NodeStatus status, + std::chrono::steady_clock::time_point ts = {}, + std::optional expected_status = std::nullopt); + + // Resolve logical_node_id from ip:port (reverse lookup) + std::string ResolveLogicalId(const std::string& ip_port) const; + private: std::shared_ptr getNodeResource(const std::string &addr_str); void updateAllNodeResource(); + // Migrate a node's physical address from old_ip_port to new_ip_port: + // updates addr_to_logical_, node_status_map_, node_info_map_, logical_node_table_, + // and registers the new address via AddNode(). Caller must pass the mutable + // NodeEntry copy that will be written back to logical_node_table_. + void migrateNodeIp(const std::string& logical_id, NodeEntry& entry, + const std::string& new_ip_port); + private: // map to record dataservers' resource stats // key : address string(ip:port) @@ -92,6 +142,12 @@ class ClusterManagerNodeManager : public std::enable_shared_from_this node_status_map_; + // logical_node_id based tables + // Primary table: logical_node_id → NodeEntry + folly::ConcurrentHashMap logical_node_table_; + // Reverse lookup: ip:port → logical_node_id + folly::ConcurrentHashMap addr_to_logical_; + sicl::rpc::SiRPC *rpc_client_{nullptr}; std::thread *resource_thread_{nullptr}; std::atomic resource_thread_stop_{false}; diff --git a/src/cluster_manager/cm_rpc_handler.cc b/src/cluster_manager/cm_rpc_handler.cc index 00a7c82..950b5d5 100644 --- a/src/cluster_manager/cm_rpc_handler.cc +++ b/src/cluster_manager/cm_rpc_handler.cc @@ -16,6 +16,7 @@ #include "proto/ds_cm_rpcs.pb.h" DECLARE_uint32(shard_total_num); +DECLARE_bool(cm_deferred_reshard_enabled); DECLARE_LOG_MODULE("cluster_manager"); @@ -77,23 +78,83 @@ void NewNodeHandshakeHandler::Work(const std::shared_ptr auto req = dynamic_cast(request); auto resp = std::make_shared(); simm::common::NodeAddress node_addr = {req->node().ip(), req->node().port()}; + std::string addr_str = node_addr.toString(); + const std::string& logical_id = req->logical_node_id(); error_code_t ret = CommonErr::OK; - if (simm::common::ModuleServiceState::GetInstance().GracePeriodFinished()) { + + if (!logical_id.empty() && simm::common::ModuleServiceState::GetInstance().GracePeriodFinished()) { + // Post-grace-period with logical_node_id: process handshake via node manager + std::vector reported_shards(req->shard_ids().begin(), req->shard_ids().end()); + auto result = node_manager_->ProcessHandshake(logical_id, addr_str, reported_shards); + + switch (result.action) { + case HandshakeResult::Action::DEFERRED_RESHARD_REPLACE: { + // Replacement scenario: collect shards from old IP, assign to new IP + auto shards = shard_manager_->GetShardsOwnedByNode(result.old_ip_port); + shard_manager_->BatchAssignRoutingTable( + shards, std::make_shared(node_addr)); + hb_monitor_->OnDeferredReshardResolved(logical_id); + result.shards_to_assign = std::move(shards); + MLOG_INFO("Node handshake replace complete: {} -> {} for {} shards", + result.old_ip_port, addr_str, result.shards_to_assign.size()); + break; + } + case HandshakeResult::Action::IP_UPDATE: { + // IP changed while RUNNING: update routing table entries to new IP + auto shards = shard_manager_->GetShardsOwnedByNode(result.old_ip_port); + shard_manager_->BatchAssignRoutingTable( + shards, std::make_shared(node_addr)); + hb_monitor_->OnDeferredReshardResolved(logical_id); + result.shards_to_assign = std::move(shards); + MLOG_INFO("Node handshake IP update: {} -> {} for {} shards", + result.old_ip_port, addr_str, result.shards_to_assign.size()); + break; + } + case HandshakeResult::Action::NEW_NODE: { + // Post-grace-period scale-out: a brand-new logical_node_id not seen before. + // Shard assignment for scale-out is not yet supported; no shards assigned here. + // TODO: implement scale-out shard assignment + MLOG_WARN("New node registered post-grace-period (scale-out not yet supported): logical_id={} ip={}", + logical_id, addr_str); + break; + } + default: + ret = CommonErr::CmRegisterNewNodeFailed; + break; + } + + // Fill assigned_shard_ids in response + for (auto sid : result.shards_to_assign) { + resp->add_assigned_shard_ids(sid); + } + + } else if (simm::common::ModuleServiceState::GetInstance().GracePeriodFinished()) { MLOG_INFO("Grace period is already finished, new dataserver({}) will be waited for joining the cluster", - node_addr.toString()); - // already out of grace period, new dataserver nodes will be hold for - // one timewindow, and be added in batch after current timewindow finishes + addr_str); + // Post-grace-period without logical_node_id: legacy behavior } else { - // still in grace period, just add nodes - // FIXME(ytji): needn't to use NodeAddress as intermediary struct - MLOG_INFO("Still in Grace period new dataserver({}) will be added into cluster", node_addr.toString()); - ret = node_manager_->AddNode(node_addr.toString()); - shard_manager_->BatchAssignRoutingTable(std::vector(req->shard_ids().begin(), req->shard_ids().end()), - std::make_shared(node_addr)); + // Still in grace period: register node and assign shards + MLOG_INFO("Still in Grace period new dataserver({}) will be added into cluster", addr_str); + std::vector reported_shards(req->shard_ids().begin(), req->shard_ids().end()); + shard_manager_->BatchAssignRoutingTable(reported_shards, + std::make_shared(node_addr)); + + if (!logical_id.empty()) { + // ProcessHandshake (Case 1) calls AddNode internally, so no separate AddNode needed. + node_manager_->ProcessHandshake(logical_id, addr_str, reported_shards); + } else { + // No logical_node_id: legacy registration path (no deferred reshard support) + ret = node_manager_->AddNode(addr_str); + } + + // Fill assigned_shard_ids in response (echo back what was assigned) + for (auto sid : reported_shards) { + resp->add_assigned_shard_ids(sid); + } } if (ret != CommonErr::OK) { - MLOG_ERROR("Failed to register new node({}) into cluster, ret:{}", node_addr.toString(), ret); + MLOG_ERROR("Failed to register new node({}) into cluster, ret:{}", addr_str, ret); } resp->set_ret_code(ret); simm::common::Metrics::Instance("cluster_manager") @@ -116,9 +177,16 @@ void NodeHeartBeatHandler::Work(const std::shared_ptr ctx auto resp = std::make_shared(); error_code_t ret = CommonErr::OK; std::string addr_str = req->node().ip() + ":" + std::to_string(req->node().port()); + const std::string& logical_id = req->logical_node_id(); if (simm::utils::IsValidV4IPAddr(req->node().ip()) && simm::utils::IsValidPortNum(req->node().port())) { - ret = hb_monitor_->OnRecvNodeHeartbeat(addr_str); + if (!logical_id.empty()) { + // heartbeat keyed by logical_node_id + ret = hb_monitor_->OnRecvNodeHeartbeat(logical_id, addr_str); + } else { + // Legacy: heartbeat keyed by ip:port + ret = hb_monitor_->OnRecvNodeHeartbeat(addr_str); + } } else { MLOG_WARN("NodeHeartBeatHandler::Work, invalid node address received:{}", addr_str); ret = CommonErr::InvalidArgument; diff --git a/src/cluster_manager/cm_rpc_handler.h b/src/cluster_manager/cm_rpc_handler.h index f53c7f5..2512e69 100644 --- a/src/cluster_manager/cm_rpc_handler.h +++ b/src/cluster_manager/cm_rpc_handler.h @@ -29,8 +29,10 @@ class NewNodeHandshakeHandler : public sicl::rpc::HandlerBase { sicl::rpc::SiRPC* service, google::protobuf::Message* request, std::shared_ptr node_manager, - std::shared_ptr shard_manager): - HandlerBase(service, request), node_manager_(node_manager), shard_manager_(shard_manager) {} + std::shared_ptr shard_manager, + std::shared_ptr hb_monitor): + HandlerBase(service, request), node_manager_(node_manager), + shard_manager_(shard_manager), hb_monitor_(hb_monitor) {} virtual void Work(const std::shared_ptr ctx, const std::shared_ptr conn, @@ -39,6 +41,7 @@ class NewNodeHandshakeHandler : public sicl::rpc::HandlerBase { private: std::shared_ptr node_manager_{nullptr}; std::shared_ptr shard_manager_{nullptr}; + std::shared_ptr hb_monitor_{nullptr}; }; class NodeHeartBeatHandler : public sicl::rpc::HandlerBase { diff --git a/src/cluster_manager/cm_service.cc b/src/cluster_manager/cm_service.cc index 4576cb6..bbae00e 100644 --- a/src/cluster_manager/cm_service.cc +++ b/src/cluster_manager/cm_service.cc @@ -146,7 +146,7 @@ error_code_t ClusterManagerService::StartRPCServices() { inter_rpc_service_->RegisterHandler( static_cast(simm::cm::ClusterManagerRpcType::RPC_NEW_NODE_HANDSHAKE), new NewNodeHandshakeHandler( - inter_rpc_service_.get(), new NewNodeHandShakeRequestPB, node_manager_, shard_manager_)); + inter_rpc_service_.get(), new NewNodeHandShakeRequestPB, node_manager_, shard_manager_, hb_monitor_)); inter_rpc_service_->RegisterHandler( static_cast(simm::cm::ClusterManagerRpcType::RPC_NODE_HEARTBEAT), new NodeHeartBeatHandler(inter_rpc_service_.get(), new DataServerHeartBeatRequestPB, hb_monitor_)); diff --git a/src/cluster_manager/cm_shard_manager.cc b/src/cluster_manager/cm_shard_manager.cc index 4e50fc7..55cbff3 100644 --- a/src/cluster_manager/cm_shard_manager.cc +++ b/src/cluster_manager/cm_shard_manager.cc @@ -218,6 +218,19 @@ error_code_t ClusterManagerShardManager::ReassignOrphanedShards( return CommonErr::OK; } +std::vector ClusterManagerShardManager::GetShardsOwnedByNode(const std::string& ip_port) { + std::vector result; + for (const auto& entry : mShardRoutingTable) { + if (entry.second) { + std::string node_addr = entry.second->node_ip_ + ":" + std::to_string(entry.second->node_port_); + if (node_addr == ip_port) { + result.push_back(entry.first); + } + } + } + return result; +} + error_code_t ClusterManagerShardManager::MarkShardsUnavailableForNodes( const std::vector &target_node_addresses) { if (target_node_addresses.empty()) { diff --git a/src/cluster_manager/cm_shard_manager.h b/src/cluster_manager/cm_shard_manager.h index d8241ab..0b84a3e 100644 --- a/src/cluster_manager/cm_shard_manager.h +++ b/src/cluster_manager/cm_shard_manager.h @@ -72,6 +72,9 @@ class ClusterManagerShardManager : public std::enable_shared_from_this &orphaned_shards, const std::vector> &alive_servers); + // Get all shards currently owned by a specific node (identified by ip:port) + std::vector GetShardsOwnedByNode(const std::string& ip_port); + // Mark shards routed to target nodes as unavailable (nullptr). error_code_t MarkShardsUnavailableForNodes(const std::vector &target_node_addresses); diff --git a/src/common/base/common_types.h b/src/common/base/common_types.h index 4635985..5a32f2b 100644 --- a/src/common/base/common_types.h +++ b/src/common/base/common_types.h @@ -16,6 +16,7 @@ namespace common { #define FOREACH_NODESTATUS(C) \ C(UNKNOWN) \ C(RUNNING) \ + C(DEFERRED_RESHARD) \ C(DEAD) #define NODESTATUS_ENUM(name) name, diff --git a/src/data_server/ds_flags.cc b/src/data_server/ds_flags.cc index 76b5ae9..a7d58cd 100644 --- a/src/data_server/ds_flags.cc +++ b/src/data_server/ds_flags.cc @@ -12,12 +12,12 @@ DEFINE_bool(ds_process_exit_cm_disconnection, true, "Exit data server gracefully after repeated CM heartbeat failures; if false, trigger CM re-register logic"); DEFINE_uint32(cm_connect_retry_interval_sec, 1, "Duration seconds between attempts to reconnect to Cluster Manager"); -DEFINE_uint32(ds_free_memory_usable_ratio, 80, "Percentage of free memory can be used by data server"); +DEFINE_uint32(ds_free_memory_usable_ratio, 90, "Percentage of free memory can be used by data server"); DEFINE_int32(ds_initial_blocks, 3, "Initial number of cache blocks to pre-allocate"); DEFINE_int32(ds_bg_evict_thread_num, 3, "background chunk level eviction thread num"); DEFINE_uint32(ds_bg_evict_interval_ms, 500, "background chunk level eviction checking interval"); DEFINE_uint32(ds_bg_evict_prefetch_factor, 1, "background chunk level eviction prefetch factor"); -DEFINE_double(ds_bg_evict_trigger_threshold, 0.85, "background chunk level eviction trigger threshold"); +DEFINE_double(ds_bg_evict_trigger_threshold, 0.95, "background chunk level eviction trigger threshold"); DEFINE_uint32(ds_bg_evict_slab_class_num, 3, "background chunk level eviction for how much slab class number"); DEFINE_uint64(ds_bg_evict_chunk_cooldown_ms, 3600000, @@ -28,3 +28,7 @@ DEFINE_uint64(memory_limit_bytes, 0, "Memory limit for data servers, 0 means to DEFINE_uint32(busy_wait_timeout_us, 10000, "Duration microseconds to wait for certain conditions"); // Log related flags DEFINE_string(ds_log_file, "/var/log/simm/simm_ds.log", "simm data server log file path & name"); +// logical node identity flag +DEFINE_string(ds_logical_node_id, "", + "Logical node ID for this DS; in K8s, auto-detected from POD_NAME env var. " + "Used by CM for Deferred Reshard node identity tracking"); diff --git a/src/data_server/kv_cache_pool.cc b/src/data_server/kv_cache_pool.cc index bf68cc8..c98b538 100644 --- a/src/data_server/kv_cache_pool.cc +++ b/src/data_server/kv_cache_pool.cc @@ -88,6 +88,8 @@ int KVCachePool::init(uint64_t bound_bytes, evictor_ = evictor; max_memory_block_num_ = bound_bytes / BLOCK_SIZE; if (initial_blocks < 0 || (size_t)initial_blocks > max_memory_block_num_) { + MLOG_WARN("KVCachePool::init initial_blocks:{} is greater than max_memory_block_num_:{}, set to upper bound", + initial_blocks, max_memory_block_num_); initial_blocks = max_memory_block_num_; } cache_memory_bound_bytes_ = max_memory_block_num_ * BLOCK_SIZE; @@ -124,13 +126,15 @@ int KVCachePool::init(uint64_t bound_bytes, return -1; } } - // init chunks - for (size_t i = 0; i < CHUNKS_PER_BLOCK; i++) { - size_t slab_idx = i % SlabClassCount; - auto sc = all_slabs[slab_idx]; - bool res = init_chunk(0, i, sc); - if (!res) { - MLOG_ERROR("Init cache chunk failed, res:{}", res); + // init chunks for the first block + if (max_memory_block_num_ != 1) { + for (size_t i = 0; i < CHUNKS_PER_BLOCK; i++) { + size_t slab_idx = i % SlabClassCount; + auto sc = all_slabs[slab_idx]; + bool res = init_chunk(0, i, sc); + if (!res) { + MLOG_ERROR("Init cache chunk failed, res:{}", res); + } } } @@ -244,7 +248,8 @@ int KVCachePool::init(uint64_t bound_bytes, // clean empty chunk which cooldowns for specific period if (first_check) { auto curr_tag = utils::current_microseconds(); - for (size_t block_idx = 1; block_idx < max_memory_block_num_; block_idx++) { + size_t start_block_idx = (max_memory_block_num_ == 1) ? 0 : 1; + for (size_t block_idx = start_block_idx; block_idx < max_memory_block_num_; block_idx++) { for (size_t chunk_idx = 0; chunk_idx < CHUNKS_PER_BLOCK; chunk_idx++) { auto assign_cnt = chunk_assign_cnt_[block_idx][chunk_idx].load(std::memory_order_acquire); if (assign_cnt == 0) { diff --git a/src/data_server/kv_rpc_service.cc b/src/data_server/kv_rpc_service.cc index 77d2f6b..26e3259 100644 --- a/src/data_server/kv_rpc_service.cc +++ b/src/data_server/kv_rpc_service.cc @@ -112,6 +112,7 @@ error_code_t KVRpcService::Init() { cache_pool_ = std::make_unique(); cache_evictor_ = std::make_unique(this); uint64_t system_level_free_memory = FLAGS_memory_limit_bytes; + uint64_t cache_bound_bytes = system_level_free_memory; if (system_level_free_memory == 0) { auto result = utils::GetMemoryFreeToUse(); if (result < 0) { @@ -120,7 +121,7 @@ error_code_t KVRpcService::Init() { } system_level_free_memory = static_cast(result); } - auto cache_bound_bytes = system_level_free_memory * FLAGS_ds_free_memory_usable_ratio / 100; + cache_bound_bytes = cache_bound_bytes == 0 ? (system_level_free_memory * FLAGS_ds_free_memory_usable_ratio / 100) : cache_bound_bytes; int ret = cache_pool_->init(cache_bound_bytes, io_service->GetMempool(), cache_evictor_.get(), FLAGS_ds_initial_blocks); if (ret) { @@ -134,6 +135,7 @@ error_code_t KVRpcService::Init() { shard_used_bytes_[i].store(0, std::memory_order_relaxed); } SetCMAddressFromK8S(); + InitLogicalNodeId(); initialized_ = true; } return CommonErr::OK; // Success @@ -308,6 +310,38 @@ void KVRpcService::SetCMAddressFromK8S() { gflags::SetCommandLineOption("cm_primary_node_port", std::to_string(cm_addr->node_port_).c_str()); } +void KVRpcService::InitLogicalNodeId() { + const char* pod_name = std::getenv("POD_NAME"); + if (pod_name && std::strlen(pod_name) > 0) { + // K8s scenario: build "namespace/pod-name" for global uniqueness across namespaces. + // POD_NAMESPACE is injected via the downward API alongside POD_NAME. + // If POD_NAMESPACE is absent (unusual misconfiguration), fall back to bare pod-name + // with a warning — uniqueness is only guaranteed within a single namespace in that case. + const char* pod_namespace = std::getenv("POD_NAMESPACE"); + if (pod_namespace && std::strlen(pod_namespace) > 0) { + logical_node_id_ = std::string(pod_namespace) + "/" + pod_name; + } else { + MLOG_WARN("POD_NAMESPACE env var not set; logical_node_id will be bare pod-name '{}'. " + "Uniqueness is only guaranteed within a single namespace.", + pod_name); + logical_node_id_ = pod_name; + } + } else if (!FLAGS_ds_logical_node_id.empty()) { + logical_node_id_ = FLAGS_ds_logical_node_id; // Non-K8s: manual flag + } else { + // No logical_node_id available: fail fast so that K8s restarts the pod with + // the correct environment (POD_NAME injected via downward API). + // Falling back to ip:port is intentionally NOT done: if the DS later restarts + // with POD_NAME set, the logical_id changes, the CM cannot match the old entry + // in DEFERRED_RESHARD state, and client IO will fail for the entire deferred + // window before a fallback reshard is triggered. + MLOG_CRITICAL("DS logical_node_id not set: neither POD_NAME env var nor " + "--ds_logical_node_id flag is provided. Aborting."); + std::abort(); + } + MLOG_INFO("DS logical_node_id: {}", logical_node_id_); +} + void KVRpcService::KeepAlive() { while (is_running_) { while (is_running_ && !is_registered_) { @@ -353,12 +387,24 @@ void KVRpcService::RegisterOnCluster() { NewNodeHandShakeRequestPB handshake_req; handshake_req.mutable_node()->set_ip(ip); handshake_req.mutable_node()->set_port(port); + handshake_req.set_logical_node_id(logical_node_id_); auto handshake_rsp = std::make_shared(); auto handshake_done = [handshake_rsp, this](const google::protobuf::Message *rsp, const std::shared_ptr ctx) { if (!ctx->Failed()) { auto response = dynamic_cast(rsp); if (response->ret_code() == CommonErr::OK) { + // Initialize KVHashTables for shards assigned by CM + if (response->assigned_shard_ids_size() > 0) { + std::lock_guard table_lock(all_table_mtx_); + for (auto shard_id : response->assigned_shard_ids()) { + if (!all_tables_[shard_id]) { + auto* table = new KVHashTable(); + table->Init(shard_id); + all_tables_[shard_id] = table; + } + } + } is_registered_ = true; register_condv_.post(); } else { @@ -393,6 +439,7 @@ void KVRpcService::RegisterToRestartedManager() { NewNodeHandShakeRequestPB handshake_req; handshake_req.mutable_node()->set_ip(ip); handshake_req.mutable_node()->set_port(port); + handshake_req.set_logical_node_id(logical_node_id_); for (const auto [shard_id, table] : all_tables_) { if (table && !table->Empty()) { handshake_req.mutable_shard_ids()->Add(shard_id); @@ -440,6 +487,7 @@ void KVRpcService::HeartBeatToCluster() { DataServerHeartBeatRequestPB heartbeat_req; heartbeat_req.mutable_node()->set_ip(ip); heartbeat_req.mutable_node()->set_port(port); + heartbeat_req.set_logical_node_id(logical_node_id_); auto heartbeat_rsp = std::make_shared(); auto heartbeat_done = [heartbeat_rsp, this](const google::protobuf::Message *rsp, const std::shared_ptr ctx) { @@ -583,9 +631,8 @@ error_code_t KVRpcService::KVPut(std::shared_ptr ctx, std::lock_guard table_lock(all_table_mtx_); table = all_tables_[shard_id]; if (!table) { - table = new KVHashTable(); - table->Init(shard_id); - all_tables_[shard_id] = table; + MLOG_ERROR("KVPut: shard {} table not initialized (not assigned by CM)", shard_id); + return DsErr::KeyNotFound; } } std::unique_ptr key_meta(object_pool_->AcquireKey(), diff --git a/src/data_server/kv_rpc_service.h b/src/data_server/kv_rpc_service.h index 9498967..3a44606 100644 --- a/src/data_server/kv_rpc_service.h +++ b/src/data_server/kv_rpc_service.h @@ -39,6 +39,7 @@ DECLARE_int32(heartbeat_cooldown_sec); DECLARE_string(cm_primary_node_ip); DECLARE_int32(cm_rpc_inter_port); DECLARE_uint32(busy_wait_timeout_us); +DECLARE_string(ds_logical_node_id); namespace simm { namespace ds { @@ -85,6 +86,7 @@ class KVRpcService { error_code_t RegisterHandlers(); void SetCMAddressFromK8S(); + void InitLogicalNodeId(); void KeepAlive(); void RegisterOnCluster(); void RegisterToRestartedManager(); @@ -117,6 +119,7 @@ class KVRpcService { std::function cluster_disconnect_handler_{}; std::string local_ip_; + std::string logical_node_id_; // stable node identity across restarts std::deque> shard_used_bytes_; friend class KVCacheEvictor; diff --git a/src/proto/ds_cm_rpcs.proto b/src/proto/ds_cm_rpcs.proto index fd44c9a..8fd0ce4 100644 --- a/src/proto/ds_cm_rpcs.proto +++ b/src/proto/ds_cm_rpcs.proto @@ -27,16 +27,19 @@ import "common.proto"; message NewNodeHandShakeRequestPB { proto.common.NodeAddressPB node = 1; repeated uint32 shard_ids = 2; + string logical_node_id = 3; // stable node identity (K8s: namespace/pod-name), CM uses as primary key } // CM -> SDS message NewNodeHandShakeResponsePB { sint32 ret_code = 1; + repeated uint32 assigned_shard_ids = 2; // shards CM assigns to this DS (deferred reshard or initial) } // ********** DataServer Heartbeat RPC ********** // SDS -> CM message DataServerHeartBeatRequestPB { proto.common.NodeAddressPB node = 1; + string logical_node_id = 2; // same as in Handshake, CM keys heartbeat records by this } // CM -> SDS message DataServerHeartBeatResponsePB { diff --git a/tests/client/test_clnt_messenger.cc b/tests/client/test_clnt_messenger.cc index da003e9..34ad329 100644 --- a/tests/client/test_clnt_messenger.cc +++ b/tests/client/test_clnt_messenger.cc @@ -17,6 +17,8 @@ DECLARE_bool(clnt_use_k8s); DECLARE_string(cm_primary_node_ip); DECLARE_int32(cm_rpc_inter_port); DECLARE_uint32(shard_total_num); +DECLARE_uint32(clnt_deferred_reshard_wait_inSecs); +DECLARE_uint32(clnt_cm_addr_check_interval_inSecs); namespace simm { namespace clnt { @@ -214,6 +216,9 @@ class ClientMessengerTestPeer { messenger.test_get_cm_address_hook_ = nullptr; messenger.test_route_query_hook_ = nullptr; messenger.test_build_connection_hook_ = nullptr; + // clear dead-since state between tests + std::lock_guard lg(messenger.ds_dead_since_mtx_); + messenger.ds_dead_since_.clear(); } static void InstallDsContext(const std::string &addr, @@ -332,6 +337,81 @@ class ClientMessengerTestPeer { static error_code_t BuildConnectionNoWait(const std::string &addr) { return ClientMessenger::Instance().build_connection(addr, ClientMessenger::BuildConnWaitMode::kNoWait); } + + // Inject a dead-since timestamp for a DS, simulating it having been seen as dead at a given + // time. Used to control the deferred-reshard wait window in failover thread tests. + static void InjectDeadSince(const std::string &addr, + std::chrono::steady_clock::time_point tp) { + auto &messenger = ClientMessenger::Instance(); + std::lock_guard lg(messenger.ds_dead_since_mtx_); + messenger.ds_dead_since_[addr] = tp; + } + + // Clear the dead-since record for a DS. + static void ClearDeadSince(const std::string &addr) { + auto &messenger = ClientMessenger::Instance(); + std::lock_guard lg(messenger.ds_dead_since_mtx_); + messenger.ds_dead_since_.erase(addr); + } + + // Check if a dead-since record exists for a DS. + static bool HasDeadSince(const std::string &addr) { + auto &messenger = ClientMessenger::Instance(); + std::lock_guard lg(messenger.ds_dead_since_mtx_); + return messenger.ds_dead_since_.count(addr) > 0; + } + + // Wake up the failover thread immediately (simulate a scan cycle without waiting for the + // clnt_cm_addr_check_interval_inSecs sleep). + static void WakeFailoverThread() { + auto &messenger = ClientMessenger::Instance(); + messenger.failover_condv_.notify_all(); + } + + // Run one failover scan iteration synchronously (the DS-inactive branch only). + // cm_addr_ is set to match the CM hook so the CM-change branch is skipped. + // Returns true if ReInit() was triggered. + static bool RunOneFailoverCycle() { + auto &messenger = ClientMessenger::Instance(); + messenger.cm_addr_ = "10.0.0.100:9000"; + + bool should_reinit = false; + for (auto [addr, ds_ctx] : messenger.ds_conn_ctxs_) { + if (!ds_ctx->active.load()) { + { + std::lock_guard lg(messenger.ds_dead_since_mtx_); + if (!messenger.ds_dead_since_.count(addr)) { + messenger.ds_dead_since_[addr] = std::chrono::steady_clock::now(); + } + } + + if (CommonErr::OK == messenger.build_connection(addr)) { + std::lock_guard lg(messenger.ds_dead_since_mtx_); + messenger.ds_dead_since_.erase(addr); + continue; + } + + std::chrono::duration dur; + { + std::lock_guard lg(messenger.ds_dead_since_mtx_); + dur = std::chrono::steady_clock::now() - messenger.ds_dead_since_[addr]; + } + + if (dur > std::chrono::seconds(FLAGS_clnt_deferred_reshard_wait_inSecs)) { + { + std::lock_guard lg(messenger.ds_dead_since_mtx_); + messenger.ds_dead_since_.erase(addr); + } + should_reinit = true; + } + } + } + + if (should_reinit) { + messenger.ReInit(); + } + return should_reinit; + } }; class ClientMessengerUnitTest : public ::testing::Test { @@ -619,6 +699,138 @@ TEST_F(ClientMessengerUnitTest, GetCmAddressUsesFlagToSkipK8SLookup) { FLAGS_cm_rpc_inter_port = old_cm_port; } +// ───────────────────────────────────────────────────────────────────────────── +// Deferred reshard window tests: verify the failover thread's DS-dead tracking +// and ReInit-trigger logic for the "DS IP changed after restart" scenario. +// +// The failover thread logic has two key branches we test here: +// 1. DS inactive + reconnect fails + window NOT expired → should_reinit = false +// 2. DS inactive + reconnect fails + window expired → should_reinit = true → ReInit() +// 3. DS inactive + reconnect succeeds → dead-since cleared, no ReInit +// +// We test these by driving the failover decision logic directly (via RunOneFailoverCycle) +// rather than relying on timing of the background thread. +// ───────────────────────────────────────────────────────────────────────────── + +// Expose a test-only hook to run one failover scan cycle synchronously. +// This mirrors the body of the failover thread loop, minus the sleep. +// We add this to ClientMessengerTestPeer rather than production code. + +// Helper: manually execute one failover scan iteration (the DS-scan branch only, +// CM-address-change check is skipped by keeping cm_addr_ matching the hook). +static bool RunOneFailoverCycle(FakeSiRPC * /*unused*/) { + return ClientMessengerTestPeer::RunOneFailoverCycle(); +} + +// Test: DS goes inactive, reconnect fails, window NOT expired → ReInit not triggered. +// After backdating dead-since past the window, next cycle triggers ReInit and picks +// up the new IP from CM. +TEST_F(ClientMessengerUnitTest, DeferredWindowExpiryTriggersReInitWithNewIP) { + const std::string old_addr = "10.0.0.1:1001"; + const std::string new_addr = "10.0.0.99:1001"; + const uint32_t window_secs = 2; + + const auto saved_window = FLAGS_clnt_deferred_reshard_wait_inSecs; + FLAGS_clnt_deferred_reshard_wait_inSecs = window_secs; + + auto *fake_rpc = InstallFakeRpcClient(); + + // DS starts active on old_addr, owns shard 0 + ClientMessengerTestPeer::InstallDsContext(old_addr, true, 1, {0}, + std::make_shared("old")); + // DS goes inactive (pod crash) + ClientMessengerTestPeer::MarkConnectionActive(old_addr, false, 2); + + // CM hook must be set before RunOneFailoverCycle (ReInit calls get_cm_address via hook) + ClientMessengerTestPeer::SetGetCmAddressHook([]() { return std::string("10.0.0.100:9000"); }); + + // build_connection for old_addr always fails — old IP unreachable. + // build_connection for new_addr succeeds — DS restarted on new IP. + // Both go through test_build_connection_hook_ which takes priority over rpc_client_->connect. + fake_rpc->SetConnectHandler([]() -> std::shared_ptr { return nullptr; }); + + std::atomic reinit_count{0}; + // CM now knows the new IP (DS has already re-registered with new IP) + ClientMessengerTestPeer::SetRouteQueryHook([&](const std::string &) { + reinit_count.fetch_add(1); + return std::make_pair(CommonErr::OK, BuildRoutingResponse({{new_addr, {0}}})); + }); + // Use build_connection hook so both old and new addr go through the same path + ClientMessengerTestPeer::SetBuildConnectionHook([&](const std::string &addr) { + if (addr == new_addr) { + ClientMessengerTestPeer::MarkConnectionActive(addr, true, 10); + return CommonErr::OK; + } + // old_addr (and any other addr) — still unreachable + return ClntErr::BuildConnectionFailed; + }); + + // ── Cycle 1: dead-since just now → still within window ── + ClientMessengerTestPeer::InjectDeadSince(old_addr, std::chrono::steady_clock::now()); + bool did_reinit = RunOneFailoverCycle(fake_rpc); + + EXPECT_FALSE(did_reinit) << "ReInit must not fire within deferred window"; + EXPECT_EQ(reinit_count.load(), 0) << "route_query hook must not be called within window"; + EXPECT_EQ(ClientMessengerTestPeer::ShardOwner(0), old_addr) + << "Shard 0 must still point to old addr while within window"; + EXPECT_TRUE(ClientMessengerTestPeer::HasDeadSince(old_addr)) + << "dead-since entry must persist while in window"; + + // ── Cycle 2: backdate dead-since past window → ReInit fires ── + ClientMessengerTestPeer::InjectDeadSince( + old_addr, + std::chrono::steady_clock::now() - std::chrono::seconds(window_secs + 1)); + did_reinit = RunOneFailoverCycle(fake_rpc); + + EXPECT_TRUE(did_reinit) << "ReInit must fire after window expires"; + EXPECT_EQ(reinit_count.load(), 1) << "route_query hook must be called exactly once"; + EXPECT_FALSE(ClientMessengerTestPeer::HasDeadSince(old_addr)) + << "dead-since entry must be cleared after ReInit"; + EXPECT_EQ(ClientMessengerTestPeer::ShardOwner(0), new_addr) + << "Shard 0 must be rerouted to new DS IP after ReInit"; + + FLAGS_clnt_deferred_reshard_wait_inSecs = saved_window; +} + +// Test: DS goes inactive but recovers on the SAME IP within the deferred window. +// build_connection succeeds → dead-since cleared, ReInit never triggered. +TEST_F(ClientMessengerUnitTest, DeferredWindowReconnectSameIPClearsDeadSince) { + const std::string addr = "10.0.0.1:1001"; + const uint32_t window_secs = 10; + + const auto saved_window = FLAGS_clnt_deferred_reshard_wait_inSecs; + FLAGS_clnt_deferred_reshard_wait_inSecs = window_secs; + + auto *fake_rpc = InstallFakeRpcClient(); + + ClientMessengerTestPeer::InstallDsContext(addr, false, 1, {0}, nullptr); + + std::atomic reinit_count{0}; + ClientMessengerTestPeer::SetRouteQueryHook([&](const std::string &) { + reinit_count.fetch_add(1); + return std::make_pair(CommonErr::OK, BuildRoutingResponse({{addr, {0}}})); + }); + + // DS recovered on same IP — reconnect succeeds + fake_rpc->SetConnectHandler([&]() { + return std::make_shared("recovered"); + }); + + // Inject a fresh dead-since (well within the window) + ClientMessengerTestPeer::InjectDeadSince(addr, std::chrono::steady_clock::now()); + + bool did_reinit = RunOneFailoverCycle(fake_rpc); + + EXPECT_FALSE(did_reinit) << "ReInit must not fire when DS recovers on same IP"; + EXPECT_EQ(reinit_count.load(), 0) << "route_query hook must not be called on same-IP recovery"; + EXPECT_FALSE(ClientMessengerTestPeer::HasDeadSince(addr)) + << "dead-since must be cleared after successful reconnect"; + EXPECT_TRUE(ClientMessengerTestPeer::IsDsActive(addr)) + << "DS must be active after successful reconnect"; + + FLAGS_clnt_deferred_reshard_wait_inSecs = saved_window; +} + } // namespace clnt } // namespace simm diff --git a/tests/cluster_manager/test_cm_deferred_reshard.cc b/tests/cluster_manager/test_cm_deferred_reshard.cc new file mode 100644 index 0000000..d0629d7 --- /dev/null +++ b/tests/cluster_manager/test_cm_deferred_reshard.cc @@ -0,0 +1,617 @@ +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include "cluster_manager/cm_hb_monitor.h" +#include "cluster_manager/cm_node_manager.h" +#include "cluster_manager/cm_rpc_handler.h" +#include "cluster_manager/cm_service.h" +#include "cluster_manager/cm_shard_manager.h" +#include "common/base/common_types.h" +#include "common/errcode/errcode_def.h" +#include "common/logging/logging.h" +#include "proto/ds_cm_rpcs.pb.h" + +DECLARE_LOG_MODULE("deferred_reshard_test"); + +DECLARE_uint32(shard_total_num); +DECLARE_int32(cm_rpc_inter_port); +DECLARE_uint32(cm_cluster_init_grace_period_inSecs); +DECLARE_uint32(cm_heartbeat_timeout_inSecs); +DECLARE_uint32(cm_heartbeat_bg_scan_interval_inSecs); +DECLARE_bool(cm_deferred_reshard_enabled); +DECLARE_uint32(cm_deferred_reshard_window_inSecs); +DECLARE_string(cm_log_file); + +namespace simm { +namespace cm { + +namespace { + +template +bool WaitUntil(Predicate pred, + std::chrono::milliseconds timeout, + std::chrono::milliseconds poll_interval = std::chrono::milliseconds(100)) { + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + if (pred()) return true; + std::this_thread::sleep_for(poll_interval); + } + return pred(); +} + +} // namespace + +// ── Mock DS with logical_node_id support ──────────────────────────────────── +// Key design: SiRPC client is created once in Start() and kept alive until Stop(). +// Handshake is synchronous. Heartbeats are async but client lifetime outlives callbacks. + +class MockDeferredDS { + public: + MockDeferredDS(const std::string& ip, int port, const std::string& logical_id, + std::chrono::milliseconds hb_interval = std::chrono::milliseconds(500)) + : ip_(ip), port_(port), logical_id_(logical_id), hb_interval_(hb_interval) {} + + ~MockDeferredDS() { Stop(); } + + void Start() { + if (running_.exchange(true)) return; + stop_.store(false); + + // Create a long-lived RPC client for this DS + sicl::rpc::SiRPC* raw = nullptr; + sicl::rpc::SiRPC::newInstance(raw, false); + rpc_client_ = raw; + + SendHandshakeSync(); + hb_thread_ = std::jthread([this]() { HBLoop(); }); + MLOG_INFO("[MockDS] Started: logical_id={} ip={}:{}", logical_id_, ip_, port_); + } + + void Stop() { + if (!running_.exchange(false)) return; + stop_.store(true); + if (hb_thread_.joinable()) hb_thread_.join(); + // Give async callbacks time to complete before deleting client + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + if (rpc_client_) { + delete rpc_client_; + rpc_client_ = nullptr; + } + MLOG_INFO("[MockDS] Stopped: logical_id={} ip={}:{}", logical_id_, ip_, port_); + } + + void RestartWithNewIp(const std::string& new_ip) { + Stop(); + ip_ = new_ip; + Start(); + } + + std::string GetAddr() const { return ip_ + ":" + std::to_string(port_); } + std::string GetLogicalId() const { return logical_id_; } + + private: + void SendHandshakeSync() { + sicl::rpc::RpcContext* ctx_p = nullptr; + sicl::rpc::RpcContext::newInstance(ctx_p); + auto ctx = std::shared_ptr(ctx_p); + ctx->set_timeout(sicl::transport::TimerTick::TIMER_3S); + + NewNodeHandShakeRequestPB req; + req.mutable_node()->set_ip(ip_); + req.mutable_node()->set_port(port_); + req.set_logical_node_id(logical_id_); + + folly::Baton<> done; + auto* resp = new NewNodeHandShakeResponsePB(); + std::string lid = logical_id_; + auto cb = [resp, lid, &done](const google::protobuf::Message* rsp, + const std::shared_ptr ctx) { + if (ctx->ErrorCode() == sicl::transport::Result::SICL_SUCCESS) { + auto* response = dynamic_cast(rsp); + MLOG_INFO("[MockDS] Handshake OK: logical_id={} ret_code={} assigned_shards={}", + lid, response->ret_code(), response->assigned_shard_ids_size()); + } else { + MLOG_ERROR("[MockDS] Handshake FAILED: logical_id={} err={}", lid, ctx->ErrorCode()); + } + delete resp; + done.post(); + }; + + rpc_client_->SendRequest("127.0.0.1", FLAGS_cm_rpc_inter_port, + static_cast(ClusterManagerRpcType::RPC_NEW_NODE_HANDSHAKE), + req, resp, ctx, cb); + done.try_wait_for(std::chrono::seconds(5)); + } + + void HBLoop() { + while (!stop_.load()) { + SendHeartbeat(); + std::this_thread::sleep_for(hb_interval_); + } + } + + void SendHeartbeat() { + sicl::rpc::RpcContext* ctx_p = nullptr; + sicl::rpc::RpcContext::newInstance(ctx_p); + auto ctx = std::shared_ptr(ctx_p); + ctx->set_timeout(sicl::transport::TimerTick::TIMER_3S); + + DataServerHeartBeatRequestPB req; + req.mutable_node()->set_ip(ip_); + req.mutable_node()->set_port(port_); + req.set_logical_node_id(logical_id_); + + auto* resp = new DataServerHeartBeatResponsePB(); + auto cb = [resp](const google::protobuf::Message*, const std::shared_ptr) { + delete resp; + }; + + rpc_client_->SendRequest("127.0.0.1", FLAGS_cm_rpc_inter_port, + static_cast(ClusterManagerRpcType::RPC_NODE_HEARTBEAT), + req, resp, ctx, cb); + } + + std::string ip_; + int port_; + std::string logical_id_; + std::chrono::milliseconds hb_interval_; + std::atomic running_{false}; + std::atomic stop_{false}; + std::jthread hb_thread_; + sicl::rpc::SiRPC* rpc_client_{nullptr}; +}; + +// ── Test fixture ──────────────────────────────────────────────────────────── + +class DeferredReshardTest : public ::testing::Test { + protected: + void SetUp() override { + simm::logging::LogConfig log_cfg{FLAGS_cm_log_file, "INFO"}; + simm::logging::LoggerManager::Instance().UpdateConfig("cluster_manager", log_cfg); + simm::logging::LoggerManager::Instance().UpdateConfig("deferred_reshard_test", log_cfg); + + saved_grace_ = FLAGS_cm_cluster_init_grace_period_inSecs; + saved_hb_timeout_ = FLAGS_cm_heartbeat_timeout_inSecs; + saved_hb_scan_ = FLAGS_cm_heartbeat_bg_scan_interval_inSecs; + saved_deferred_enabled_ = FLAGS_cm_deferred_reshard_enabled; + saved_deferred_window_ = FLAGS_cm_deferred_reshard_window_inSecs; + + FLAGS_cm_cluster_init_grace_period_inSecs = 2; + FLAGS_cm_heartbeat_timeout_inSecs = 2; + FLAGS_cm_heartbeat_bg_scan_interval_inSecs = 1; + FLAGS_cm_deferred_reshard_enabled = true; + FLAGS_cm_deferred_reshard_window_inSecs = 10; + } + + void TearDown() override { + FLAGS_cm_cluster_init_grace_period_inSecs = saved_grace_; + FLAGS_cm_heartbeat_timeout_inSecs = saved_hb_timeout_; + FLAGS_cm_heartbeat_bg_scan_interval_inSecs = saved_hb_scan_; + FLAGS_cm_deferred_reshard_enabled = saved_deferred_enabled_; + FLAGS_cm_deferred_reshard_window_inSecs = saved_deferred_window_; + } + + struct CMContext { + std::shared_ptr shard_mgr; + std::shared_ptr node_mgr; + std::shared_ptr hb_mon; + std::unique_ptr service; + std::atomic started{false}; + std::atomic done{false}; + folly::Baton<> done_baton; + }; + + std::unique_ptr StartCM(folly::CPUThreadPoolExecutor& executor, uint32_t ds_num, + const std::string& ip_prefix, int base_port) { + auto ctx = std::make_unique(); + ctx->shard_mgr = std::make_shared(); + ctx->node_mgr = std::make_shared(); + ctx->hb_mon = std::make_shared(ctx->node_mgr, ctx->shard_mgr); + ctx->service = std::make_unique(ctx->shard_mgr, ctx->node_mgr, ctx->hb_mon); + + auto* raw = ctx.get(); + executor.add([raw, ds_num, ip_prefix, base_port]() { + simm::common::ModuleServiceState::GetInstance().Reset(FLAGS_cm_cluster_init_grace_period_inSecs); + + std::vector> servers; + for (uint32_t i = 0; i < ds_num; ++i) { + auto addr = std::make_shared( + ip_prefix + std::to_string(i + 1), base_port + static_cast(i)); + servers.push_back(addr); + raw->node_mgr->AddNode(addr->toString()); + } + + raw->service->Init(); + raw->shard_mgr->InitShardRoutingTable(servers); + raw->service->Start(); + raw->started.store(true); + + while (!raw->done.load()) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + + raw->service->Stop(); + raw->done_baton.post(); + }); + + EXPECT_TRUE(WaitUntil([&]() { return ctx->started.load(); }, std::chrono::seconds(15))); + return ctx; + } + + void StopCM(std::unique_ptr& ctx) { + ctx->done.store(true); + ctx->done_baton.wait(); + } + + private: + uint32_t saved_grace_{}, saved_hb_timeout_{}, saved_hb_scan_{}, saved_deferred_window_{}; + bool saved_deferred_enabled_{}; +}; + +// ═════════════════════════════════════════════════════════════════════════════ +// TEST 1: DS crashes → DEFERRED_RESHARD → new DS with same logical_id +// → in-place IP replacement, shards stay, no reshard +// ═════════════════════════════════════════════════════════════════════════════ + +TEST_F(DeferredReshardTest, ReplacementWithinWindow) { + constexpr uint32_t kDSNum = 3; + const std::string ip_prefix = "10.100.0."; + const int base_port = 50000; + + folly::CPUThreadPoolExecutor executor(8); + auto cm = StartCM(executor, kDSNum, ip_prefix, base_port); + + // Start 3 DS + std::vector> ds_list; + for (uint32_t i = 0; i < kDSNum; ++i) { + ds_list.push_back(std::make_unique( + ip_prefix + std::to_string(i + 1), base_port + i, + "simm-ds-" + std::to_string(i))); + ds_list.back()->Start(); + } + + // Wait for registration + stable heartbeats + std::this_thread::sleep_for(std::chrono::seconds(FLAGS_cm_cluster_init_grace_period_inSecs + 2)); + + // Record DS-0's initial shard ownership + std::string ds0_old_addr = ds_list[0]->GetAddr(); + auto ds0_initial_shards = cm->shard_mgr->GetShardsOwnedByNode(ds0_old_addr); + MLOG_INFO("DS-0 ({}) owns {} shards initially", ds0_old_addr, ds0_initial_shards.size()); + ASSERT_GT(ds0_initial_shards.size(), 0u); + + // ── Kill DS-0 ── + MLOG_INFO("=== Killing DS-0 ==="); + ds_list[0]->Stop(); + + // Wait for HB timeout → should enter DEFERRED_RESHARD + std::this_thread::sleep_for(std::chrono::seconds( + FLAGS_cm_heartbeat_timeout_inSecs + FLAGS_cm_heartbeat_bg_scan_interval_inSecs + 1)); + + auto entry = cm->node_mgr->GetNodeEntry("simm-ds-0"); + ASSERT_TRUE(entry.has_value()) << "simm-ds-0 should exist in logical_node_table"; + EXPECT_EQ(entry->status, NodeStatus::DEFERRED_RESHARD); + MLOG_INFO("PASS: DS-0 is in DEFERRED_RESHARD state"); + + // Shards should NOT have moved + auto shards_still = cm->shard_mgr->GetShardsOwnedByNode(ds0_old_addr); + EXPECT_EQ(shards_still.size(), ds0_initial_shards.size()); + MLOG_INFO("PASS: {} shards still on old IP (no reshard)", shards_still.size()); + + // ── Restart DS-0 with NEW IP ── + std::string ds0_new_ip = "10.100.0.99"; + std::string ds0_new_addr = ds0_new_ip + ":" + std::to_string(base_port); + MLOG_INFO("=== Restarting DS-0 with new IP {} ===", ds0_new_addr); + ds_list[0]->RestartWithNewIp(ds0_new_ip); + + std::this_thread::sleep_for(std::chrono::seconds(2)); + + // Verify: back to RUNNING with new IP + entry = cm->node_mgr->GetNodeEntry("simm-ds-0"); + ASSERT_TRUE(entry.has_value()); + EXPECT_EQ(entry->status, NodeStatus::RUNNING); + EXPECT_EQ(entry->current_ip_port, ds0_new_addr); + MLOG_INFO("PASS: DS-0 back to RUNNING with new IP {}", ds0_new_addr); + + // Verify: shards moved to new IP + auto shards_new = cm->shard_mgr->GetShardsOwnedByNode(ds0_new_addr); + EXPECT_EQ(shards_new.size(), ds0_initial_shards.size()); + MLOG_INFO("PASS: {} shards now on new IP", shards_new.size()); + + // Verify: no shards on old IP + auto shards_old = cm->shard_mgr->GetShardsOwnedByNode(ds0_old_addr); + EXPECT_EQ(shards_old.size(), 0u); + MLOG_INFO("PASS: 0 shards on old IP"); + + // Verify: total shards intact + auto all = cm->shard_mgr->QueryAllShardRoutingInfos(); + EXPECT_EQ(all.size(), FLAGS_shard_total_num); + MLOG_INFO("PASS: Total shards = {}", all.size()); + + for (auto& ds : ds_list) ds->Stop(); + StopCM(cm); + executor.join(); +} + +// ═════════════════════════════════════════════════════════════════════════════ +// TEST 2: DS crashes, NO replacement → window expires → fallback reshard +// ═════════════════════════════════════════════════════════════════════════════ + +TEST_F(DeferredReshardTest, FallbackToReshardOnWindowExpiry) { + FLAGS_cm_deferred_reshard_window_inSecs = 4; // short window for fast test + + constexpr uint32_t kDSNum = 4; // need >= dataserver_min_num+1 so reshard has enough alive servers + const std::string ip_prefix = "10.200.0."; + const int base_port = 51000; + + folly::CPUThreadPoolExecutor executor(8); + auto cm = StartCM(executor, kDSNum, ip_prefix, base_port); + + std::vector> ds_list; + for (uint32_t i = 0; i < kDSNum; ++i) { + ds_list.push_back(std::make_unique( + ip_prefix + std::to_string(i + 1), base_port + i, + "fallback-ds-" + std::to_string(i))); + ds_list.back()->Start(); + } + + std::this_thread::sleep_for(std::chrono::seconds(FLAGS_cm_cluster_init_grace_period_inSecs + 2)); + + std::string ds0_addr = ds_list[0]->GetAddr(); + auto ds0_shards = cm->shard_mgr->GetShardsOwnedByNode(ds0_addr); + ASSERT_GT(ds0_shards.size(), 0u); + + // Kill DS-0 and do NOT restart + MLOG_INFO("=== Killing DS-0, will NOT restart ==="); + ds_list[0]->Stop(); + + // Wait for: HB timeout → DEFERRED_RESHARD → window expiry → DEAD + reshard + auto total_wait = FLAGS_cm_heartbeat_timeout_inSecs + FLAGS_cm_deferred_reshard_window_inSecs + + FLAGS_cm_heartbeat_bg_scan_interval_inSecs * 2 + 2; + MLOG_INFO("Waiting {}s for full cycle...", total_wait); + std::this_thread::sleep_for(std::chrono::seconds(total_wait)); + + // Verify: DS-0 should be DEAD + auto entry = cm->node_mgr->GetNodeEntry("fallback-ds-0"); + ASSERT_TRUE(entry.has_value()); + EXPECT_EQ(entry->status, NodeStatus::DEAD); + MLOG_INFO("PASS: DS-0 is DEAD after window expiry"); + + // Verify: shards redistributed + auto shards_dead = cm->shard_mgr->GetShardsOwnedByNode(ds0_addr); + EXPECT_EQ(shards_dead.size(), 0u); + MLOG_INFO("PASS: 0 shards on dead DS"); + + uint32_t alive_total = 0; + for (uint32_t i = 1; i < kDSNum; ++i) { + alive_total += cm->shard_mgr->GetShardsOwnedByNode(ds_list[i]->GetAddr()).size(); + } + EXPECT_EQ(alive_total, FLAGS_shard_total_num); + MLOG_INFO("PASS: All {} shards on alive nodes", alive_total); + + for (auto& ds : ds_list) ds->Stop(); + StopCM(cm); + executor.join(); +} + +// ═════════════════════════════════════════════════════════════════════════════ +// TEST 3: Deferred Reshard disabled → immediate DEAD + reshard (legacy) +// ═════════════════════════════════════════════════════════════════════════════ + +TEST_F(DeferredReshardTest, DisabledFallsBackToImmediate) { + FLAGS_cm_deferred_reshard_enabled = false; + + constexpr uint32_t kDSNum = 4; // need >= dataserver_min_num+1 + const std::string ip_prefix = "10.88.0."; + const int base_port = 52000; + + folly::CPUThreadPoolExecutor executor(8); + auto cm = StartCM(executor, kDSNum, ip_prefix, base_port); + + std::vector> ds_list; + for (uint32_t i = 0; i < kDSNum; ++i) { + ds_list.push_back(std::make_unique( + ip_prefix + std::to_string(i + 1), base_port + i, + "disabled-ds-" + std::to_string(i))); + ds_list.back()->Start(); + } + + std::this_thread::sleep_for(std::chrono::seconds(FLAGS_cm_cluster_init_grace_period_inSecs + 2)); + + std::string ds0_addr = ds_list[0]->GetAddr(); + ASSERT_GT(cm->shard_mgr->GetShardsOwnedByNode(ds0_addr).size(), 0u); + + MLOG_INFO("=== Killing DS-0 with deferred reshard DISABLED ==="); + ds_list[0]->Stop(); + + // With deferred disabled, should go straight to DEAD after HB timeout + std::this_thread::sleep_for(std::chrono::seconds( + FLAGS_cm_heartbeat_timeout_inSecs + FLAGS_cm_heartbeat_bg_scan_interval_inSecs + 2)); + + auto shards_dead = cm->shard_mgr->GetShardsOwnedByNode(ds0_addr); + EXPECT_EQ(shards_dead.size(), 0u); + MLOG_INFO("PASS: Immediate reshard, 0 shards on dead DS"); + + for (auto& ds : ds_list) ds->Stop(); + StopCM(cm); + executor.join(); +} + +// ═════════════════════════════════════════════════════════════════════════════ +// TEST 4: Kill enough DS so alive < dataserver_min_num +// → Window expires but nodes stay DEFERRED_RESHARD (no orphaned shards) +// → Restart one DS → deferred reshard replace succeeds +// ═════════════════════════════════════════════════════════════════════════════ + +TEST_F(DeferredReshardTest, InsufficientNodesKeepsDeferredState) { + // dataserver_min_num defaults to 3 + // Start 3 DS, kill 2 → only 1 alive → reshard not feasible → should stay DEFERRED_RESHARD + FLAGS_cm_deferred_reshard_window_inSecs = 3; // short window + + constexpr uint32_t kDSNum = 3; + const std::string ip_prefix = "10.77.0."; + const int base_port = 53000; + + folly::CPUThreadPoolExecutor executor(8); + auto cm = StartCM(executor, kDSNum, ip_prefix, base_port); + + std::vector> ds_list; + for (uint32_t i = 0; i < kDSNum; ++i) { + ds_list.push_back(std::make_unique( + ip_prefix + std::to_string(i + 1), base_port + i, + "safe-ds-" + std::to_string(i))); + ds_list.back()->Start(); + } + + std::this_thread::sleep_for(std::chrono::seconds(FLAGS_cm_cluster_init_grace_period_inSecs + 2)); + + // Record initial shards + std::string ds0_addr = ds_list[0]->GetAddr(); + std::string ds1_addr = ds_list[1]->GetAddr(); + auto ds0_shards = cm->shard_mgr->GetShardsOwnedByNode(ds0_addr); + auto ds1_shards = cm->shard_mgr->GetShardsOwnedByNode(ds1_addr); + ASSERT_GT(ds0_shards.size(), 0u); + ASSERT_GT(ds1_shards.size(), 0u); + MLOG_INFO("DS-0 ({}) owns {} shards, DS-1 ({}) owns {} shards", + ds0_addr, ds0_shards.size(), ds1_addr, ds1_shards.size()); + + // Kill DS-0 and DS-1 simultaneously — only DS-2 survives (1 alive < min 3) + MLOG_INFO("=== Killing DS-0 and DS-1 simultaneously ==="); + ds_list[0]->Stop(); + ds_list[1]->Stop(); + + // Wait for HB timeout → both enter DEFERRED_RESHARD + std::this_thread::sleep_for(std::chrono::seconds( + FLAGS_cm_heartbeat_timeout_inSecs + FLAGS_cm_heartbeat_bg_scan_interval_inSecs + 1)); + + auto e0 = cm->node_mgr->GetNodeEntry("safe-ds-0"); + auto e1 = cm->node_mgr->GetNodeEntry("safe-ds-1"); + ASSERT_TRUE(e0.has_value()); + ASSERT_TRUE(e1.has_value()); + EXPECT_EQ(e0->status, NodeStatus::DEFERRED_RESHARD); + EXPECT_EQ(e1->status, NodeStatus::DEFERRED_RESHARD); + MLOG_INFO("PASS: Both DS-0 and DS-1 are in DEFERRED_RESHARD"); + + // Wait for window to expire + MLOG_INFO("Waiting for deferred window ({} s) to expire...", FLAGS_cm_deferred_reshard_window_inSecs); + std::this_thread::sleep_for(std::chrono::seconds( + FLAGS_cm_deferred_reshard_window_inSecs + FLAGS_cm_heartbeat_bg_scan_interval_inSecs + 1)); + + // Key assertion: both should STILL be DEFERRED_RESHARD (not DEAD) + // because marking them DEAD would leave only 1 alive node < dataserver_min_num=3 + e0 = cm->node_mgr->GetNodeEntry("safe-ds-0"); + e1 = cm->node_mgr->GetNodeEntry("safe-ds-1"); + ASSERT_TRUE(e0.has_value()); + ASSERT_TRUE(e1.has_value()); + EXPECT_EQ(e0->status, NodeStatus::DEFERRED_RESHARD) + << "DS-0 should stay DEFERRED_RESHARD (not DEAD) because reshard not feasible"; + EXPECT_EQ(e1->status, NodeStatus::DEFERRED_RESHARD) + << "DS-1 should stay DEFERRED_RESHARD (not DEAD) because reshard not feasible"; + MLOG_INFO("PASS: Both still DEFERRED_RESHARD after window expiry (safe degradation)"); + + // Shards should still be on old IPs (not redistributed, not orphaned) + auto s0 = cm->shard_mgr->GetShardsOwnedByNode(ds0_addr); + auto s1 = cm->shard_mgr->GetShardsOwnedByNode(ds1_addr); + EXPECT_EQ(s0.size(), ds0_shards.size()); + EXPECT_EQ(s1.size(), ds1_shards.size()); + MLOG_INFO("PASS: Shards still on old IPs (not orphaned)"); + + // Now restart DS-0 with new IP → should deferred replace successfully + std::string ds0_new_ip = "10.77.0.99"; + MLOG_INFO("=== Restarting DS-0 with new IP {} ===", ds0_new_ip); + ds_list[0]->RestartWithNewIp(ds0_new_ip); + std::this_thread::sleep_for(std::chrono::seconds(2)); + + e0 = cm->node_mgr->GetNodeEntry("safe-ds-0"); + ASSERT_TRUE(e0.has_value()); + EXPECT_EQ(e0->status, NodeStatus::RUNNING); + MLOG_INFO("PASS: DS-0 back to RUNNING via deferred reshard replace"); + + // Verify shards moved to new IP + std::string ds0_new_addr = ds0_new_ip + ":" + std::to_string(base_port); + auto new_shards = cm->shard_mgr->GetShardsOwnedByNode(ds0_new_addr); + EXPECT_EQ(new_shards.size(), ds0_shards.size()); + MLOG_INFO("PASS: DS-0's {} shards now on new IP {}", new_shards.size(), ds0_new_addr); + + for (auto& ds : ds_list) ds->Stop(); + StopCM(cm); + executor.join(); +} + +// ═════════════════════════════════════════════════════════════════════════════ +// TEST 5: DS restarts before HB timeout (same logical_id, same IP) +// → ProcessHandshake should return IP_UPDATE action +// → shards inherited and returned to DS (not empty) +// ═════════════════════════════════════════════════════════════════════════════ + +TEST(ProcessHandshakeUnitTest, ReregistrationInheritsShard) { + // Unit test: no CM service, directly tests ProcessHandshake + GetShardsOwnedByNode + auto shard_mgr = std::make_shared(); + auto node_mgr = std::make_shared(); + + const std::string logical_id = "simm-ds-0"; + const std::string ip_port = "192.168.1.10:40000"; + + // Step 1: Initial registration (simulate grace-period AddNode + BatchAssign) + node_mgr->AddNode(ip_port); + + // Assign shards 0..3 to this DS + std::vector expected_shards = {0, 1, 2, 3}; + auto addr = std::make_shared("192.168.1.10", 40000); + shard_mgr->BatchAssignRoutingTable(expected_shards, addr); + + // Also register in logical_node_table via ProcessHandshake (grace-period path) + node_mgr->ProcessHandshake(logical_id, ip_port, expected_shards); + + // Verify initial state + { + auto entry = node_mgr->GetNodeEntry(logical_id); + ASSERT_TRUE(entry.has_value()); + EXPECT_EQ(entry->status, NodeStatus::RUNNING); + EXPECT_EQ(entry->current_ip_port, ip_port); + } + + // Step 2: DS crashes and restarts — same logical_id, same IP, no shard_ids in request + std::vector empty_shards; + auto result = node_mgr->ProcessHandshake(logical_id, ip_port, empty_shards); + + // Key assertion: action should be IP_UPDATE (not NEW_NODE), with old_ip_port filled + EXPECT_EQ(result.action, HandshakeResult::Action::IP_UPDATE); + EXPECT_EQ(result.old_ip_port, ip_port); + MLOG_INFO("PASS: ProcessHandshake returned IP_UPDATE for re-registration"); + + // Step 3: Simulate what the handler does — query shards from ShardManager + auto inherited_shards = shard_mgr->GetShardsOwnedByNode(result.old_ip_port); + EXPECT_EQ(inherited_shards.size(), expected_shards.size()); + std::unordered_set expected_set(expected_shards.begin(), expected_shards.end()); + for (auto s : inherited_shards) { + EXPECT_TRUE(expected_set.count(s) > 0) << "Unexpected shard: " << s; + } + MLOG_INFO("PASS: {} shards inherited correctly via GetShardsOwnedByNode", inherited_shards.size()); + + // Step 4: BatchAssignRoutingTable with same addr is a no-op — verify routing table unchanged + shard_mgr->BatchAssignRoutingTable(inherited_shards, addr); + auto still_owned = shard_mgr->GetShardsOwnedByNode(ip_port); + EXPECT_EQ(still_owned.size(), expected_shards.size()); + MLOG_INFO("PASS: Routing table unchanged after re-registration (idempotent)"); +} + +} // namespace cm +} // namespace simm + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, true); + return RUN_ALL_TESTS(); +} diff --git a/tests/cluster_manager/test_cm_hb_monitor.cc b/tests/cluster_manager/test_cm_hb_monitor.cc index 42bf6f7..c877c1b 100644 --- a/tests/cluster_manager/test_cm_hb_monitor.cc +++ b/tests/cluster_manager/test_cm_hb_monitor.cc @@ -28,6 +28,7 @@ DECLARE_uint32(cm_heartbeat_records_perserver); DECLARE_uint32(cm_heartbeat_bg_scan_interval_inSecs); DECLARE_uint32(cm_heartbeat_timeout_inSecs); DECLARE_uint32(cm_cluster_init_grace_period_inSecs); +DECLARE_uint32(dataserver_min_num); DECLARE_string(cm_log_file); namespace simm { @@ -360,10 +361,13 @@ TEST_F(ClusterManagerHBMonitorTest, TestHBMonitor) { TEST_F(ClusterManagerHBMonitorTest, TestRestartAfterStopResumesHeartbeatScanning) { auto old_timeout = FLAGS_cm_heartbeat_timeout_inSecs; auto old_scan_interval = FLAGS_cm_heartbeat_bg_scan_interval_inSecs; + auto old_min_num = FLAGS_dataserver_min_num; auto restore_flags = folly::makeGuard([&]() { FLAGS_cm_heartbeat_timeout_inSecs = old_timeout; FLAGS_cm_heartbeat_bg_scan_interval_inSecs = old_scan_interval; + FLAGS_dataserver_min_num = old_min_num; }); + FLAGS_dataserver_min_num = 0; // single-node cluster: 0 alive after death is acceptable FLAGS_cm_heartbeat_timeout_inSecs = 1; FLAGS_cm_heartbeat_bg_scan_interval_inSecs = 1;