[Feat]: implement deferred reshard with logical node identity and …

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)");

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<double> 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<int>(dur.count()), FLAGS_clnt_deferred_reshard_wait_inSecs);
               }
+              // else: still within wait window, keep retrying next cycle
             }
           }
         }

src/client/clnt_messenger.h

@@ -217,6 +217,11 @@ class ClientMessenger {
   std::atomic<bool> 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<std::string, std::chrono::steady_clock::time_point> ds_dead_since_;
+
   std::unique_ptr<simm::trace::TraceServer> trace_server_{nullptr};
 
 #if defined(SIMM_UNIT_TEST)

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,

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<std::string> 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<std::string> &
   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);

src/cluster_manager/cm_hb_monitor.h

@@ -38,16 +38,24 @@ class ClusterManagerHBMonitor : public std::enable_shared_from_this<ClusterManag
   error_code_t Start();
   error_code_t Stop();
 
+  // Legacy: heartbeat keyed by ip:port (backward compat for old DS without logical_node_id)
   error_code_t OnRecvNodeHeartbeat(const std::string &node_addr_str);
 
+  // heartbeat keyed by logical_node_id + ip:port
+  error_code_t OnRecvNodeHeartbeat(const std::string &logical_node_id, const std::string &ip_port);
+
+  // Called when a DEFERRED_RESHARD is successfully resolved (new DS registered).
+  // Resets stale heartbeat records so the new DS starts fresh.
+  void OnDeferredReshardResolved(const std::string &logical_node_id);
+
  private:
   void BgHBScanLoop();
 
   void HandleNodeFailure(const std::vector<std::string> &dead_dataservers);
 
  private:
   // record all dataservers' heartbeat timestamps
-  // key : ds address string(ip:port)
+  // key : logical_node_id (preferred) or ds address string(ip:port) for legacy
   // val : deque of latest N(default is 100) heartbeat timestamps
   using InnerUOMap = std::unordered_map<std::string, std::deque<simm::common::NodeHeartbeatTs>>;
   folly::Synchronized<InnerUOMap> ds_hb_records_;