[data_storage] avoid false NODE_NOT_REGISTERED during heartbeat (#302)

Author: wangxiyu191

docs/design/report_event_snapshot_uri_version.md

@@ -384,8 +384,14 @@ snapshot replace + commit/abort
   lifecycle 写入；BatchMerge 的 block-create 与 targeted-location 两个 RMW 阶段之间同样释放并
   重新获取 lease；
 - mutation 在已经持有 metadata 锁时只做上述非阻塞的 per-reporter lifecycle lease 获取，
-  避免与 cleanup 的 `lifecycle -> metadata` 顺序形成锁序反转；不同 reporter 使用独立 fence，
-  不会因其他 host 的 HEARTBEAT/REGISTER 产生假失败；
+  避免与 cleanup 的 `lifecycle -> metadata` 顺序形成锁序反转。lifecycle fence 不存在、reporter
+  已注销、generation 不匹配，或 `try_lock` 与 REGISTER、HOST_DOWN、unavailable recovery 等
+  unique lifecycle writer 冲突时，返回 `NODE_NOT_REGISTERED`；不同 reporter 使用独立 fence；
+- 已注册且可用的 steady HEARTBEAT 不改变 registration/generation，因此持有 shared lifecycle
+  lease 完成心跳时间、状态与指标发布；同 generation 的 ADD/DELETE 可以同时取得 shared lease，
+  REGISTER、HOST_DOWN 和 unavailable recovery 等 lifecycle transition 则持有 unique lease；
+- 同一 reporter 的 HEARTBEAT 应保持 single-flight。曾评估在等待 status 锁前释放 node-table 锁，
+  但这会允许重叠 HEARTBEAT 乱序发布完整 status snapshot，因此未采用；
 - liveness unregister 的 generation 比较与节点删除在同一把锁内完成；
 - 显式 HOST_DOWN 的 generation 捕获与节点删除同样在同一把锁内完成，Heartbeat/REGISTER
   只能在线性化的 HOST_DOWN 之前或之后生效，不能在中间恢复后又被旧请求删除；

kv_cache_manager/data_storage/event_report_backend.cc

@@ -511,11 +511,29 @@ ErrorCode EventReportBackend::OnHeartbeat(const std::string &instance_id,
     }
     const ReporterSnapshotKey reporter_key{instance_id, host_ip_port};
     const auto lifecycle_fence = GetOrCreateLifecycleFence(reporter_key);
+
+    // Fast path: a steady HEARTBEAT does not change lifecycle state. Use a
+    // shared lease so same-generation ADD/DELETE can proceed; it still pins
+    // NodeInfo and excludes REGISTER/HOST_DOWN writers.
+    {
+        std::shared_lock<std::shared_mutex> lifecycle_lock(lifecycle_fence->mutex);
+        if (!AcceptingReports()) {
+            return EC_INSTANCE_NOT_EXIST;
+        }
+        std::unique_lock<std::shared_mutex> nodes_lock(nodes_mutex_);
+        if (TryPublishSteadyHeartbeatLocked(reporter_key, *lifecycle_fence, system_status, nodes_lock)) {
+            return EC_OK;
+        }
+    }
+
+    // Slow path: HEARTBEAT may create or recover a node and change lifecycle
+    // state. Drop the shared lease, acquire it exclusively, and revalidate
+    // because shared_mutex has no atomic lock upgrade.
     std::unique_lock<std::shared_mutex> lifecycle_lock(lifecycle_fence->mutex);
     if (!AcceptingReports()) {
         return EC_INSTANCE_NOT_EXIST;
     }
-    std::unique_lock<std::shared_mutex> lock(nodes_mutex_);
+    std::unique_lock<std::shared_mutex> nodes_lock(nodes_mutex_);
     auto &host_map = instance_nodes_[instance_id];
     auto it = host_map.find(host_ip_port);
     if (it == host_map.end()) {
@@ -563,16 +581,47 @@ ErrorCode EventReportBackend::OnHeartbeat(const std::string &instance_id,
     }
     lifecycle_fence->generation = node_generation_[instance_id][host_ip_port];
     lifecycle_fence->registered = true;
+    PublishHeartbeatStatus(info, system_status, nodes_lock);
+    return EC_OK;
+}
+
+bool EventReportBackend::TryPublishSteadyHeartbeatLocked(const ReporterSnapshotKey &reporter_key,
+                                                         const LifecycleFence &lifecycle_fence,
+                                                         const std::map<std::string, std::string> &system_status,
+                                                         std::unique_lock<std::shared_mutex> &nodes_lock) {
+    if (!lifecycle_fence.registered) {
+        return false;
+    }
+    const auto instance_it = instance_nodes_.find(reporter_key.instance_id);
+    const auto generation_it = node_generation_.find(reporter_key.instance_id);
+    if (instance_it == instance_nodes_.end() || generation_it == node_generation_.end()) {
+        return false;
+    }
+    const auto node_it = instance_it->second.find(reporter_key.host_ip_port);
+    const auto host_generation_it = generation_it->second.find(reporter_key.host_ip_port);
+    if (node_it == instance_it->second.end() || !node_it->second || host_generation_it == generation_it->second.end() ||
+        lifecycle_fence.generation != host_generation_it->second ||
+        !node_it->second->available.load(std::memory_order_relaxed)) {
+        return false;
+    }
+    auto &info = *node_it->second;
+    info.last_heartbeat_ms.store(NowMillis(), std::memory_order_release);
+    PublishHeartbeatStatus(info, system_status, nodes_lock);
+    return true;
+}
+
+void EventReportBackend::PublishHeartbeatStatus(NodeInfo &info,
+                                                const std::map<std::string, std::string> &system_status,
+                                                std::unique_lock<std::shared_mutex> &nodes_lock) {
     std::unique_lock<std::mutex> status_lock(info.status_mutex);
     const std::map<std::string, std::string> previous_system_status = info.last_system_status;
     info.last_system_status = system_status;
     const auto metrics_tags = info.metrics_tags;
 
-    // The per-reporter lifecycle writer keeps NodeInfo alive and prevents
-    // HOST_DOWN/REGISTER from crossing gauge publication. The status lock
-    // serializes SetNodeUnavailable's gauge reset. Release the global node
-    // table lock so metric work for one reporter does not block all others.
-    lock.unlock();
+    // Lock handoff: status_mutex now serializes this publication with
+    // SetNodeUnavailable's gauge reset, while the lifecycle lease pins
+    // NodeInfo. Release the global node-table lock before slower metric work.
+    nodes_lock.unlock();
     if (metrics_registry_) {
         const auto parse_gauge = [](const std::string &value, double &out) {
             if (value.empty()) {
@@ -607,7 +656,6 @@ ErrorCode EventReportBackend::OnHeartbeat(const std::string &instance_id,
             }
         }
     }
-    return EC_OK;
 }
 
 void EventReportBackend::SetNodeUnavailable(const std::string &instance_id, const std::string &host_ip_port) {
@@ -1126,6 +1174,9 @@ ErrorCode EventReportBackend::AcquireLifecycleMutationLease(const ReporterSnapsh
     if (!lifecycle_fence) {
         return EC_NODE_NOT_REGISTERED;
     }
+    // ADD/DELETE mutate metadata, not reporter lifecycle. This shared lease
+    // pins and validates the current registration/generation while the
+    // metadata RMW is in progress.
     auto lease = std::make_shared<std::shared_lock<std::shared_mutex>>(lifecycle_fence->mutex, std::try_to_lock);
     if (!lease->owns_lock()) {
         // Do not wait behind a lifecycle writer while the caller may already
@@ -1151,10 +1202,11 @@ ErrorCode EventReportBackend::CommitSnapshotVersionIfGeneration(const ReporterSn
         return EC_NODE_NOT_REGISTERED;
     }
     // Commit runs after the metadata RMW has released its shard locks, so it
-    // can safely wait for a transient HEARTBEAT writer. Using the mutation
+    // can safely wait for a transient lifecycle writer. Using the mutation
     // path's try-lock here would turn harmless lock contention into a failed
-    // snapshot. REGISTER/HOST_DOWN still serialize first and are rejected by
-    // the generation/registered check below.
+    // snapshot. REGISTER/HOST_DOWN and a lifecycle-changing HEARTBEAT still
+    // serialize first and are rejected by the generation/registered check
+    // below.
     std::shared_lock<std::shared_mutex> lifecycle_lease(lifecycle_fence->mutex);
     if (!AcceptingReports()) {
         return EC_INSTANCE_NOT_EXIST;
@@ -1197,9 +1249,9 @@ ErrorCode EventReportBackend::AcquireSnapshotCleanupLease(const ReporterSnapshot
         return EC_MISMATCH;
     }
     // Cleanup acquires this lease before taking metadata locks. It can block
-    // behind a short-lived HEARTBEAT/REGISTER writer without creating the
-    // lifecycle->metadata / metadata->lifecycle inversion that forces delta
-    // mutations to use try_lock.
+    // behind a short-lived REGISTER or lifecycle-changing HEARTBEAT writer
+    // without creating the lifecycle->metadata / metadata->lifecycle
+    // inversion that forces delta mutations to use try_lock.
     auto lease = std::make_shared<std::shared_lock<std::shared_mutex>>(lifecycle_fence->mutex);
     if (!AcceptingReports() || !lifecycle_fence->registered || lifecycle_fence->generation != expected_generation) {
         return EC_MISMATCH;

kv_cache_manager/data_storage/event_report_backend.h

@@ -188,6 +188,21 @@ class EventReportBackend : public DataStorageBackend {
         MetricsTags metrics_tags;
     };
 
+    // Caller holds the reporter lifecycle lease and nodes_mutex_. Publishes
+    // and returns true only when HEARTBEAT does not change lifecycle. Success
+    // releases nodes_lock; failure leaves it held.
+    bool TryPublishSteadyHeartbeatLocked(const ReporterSnapshotKey &reporter_key,
+                                         const LifecycleFence &lifecycle_fence,
+                                         const std::map<std::string, std::string> &system_status,
+                                         std::unique_lock<std::shared_mutex> &nodes_lock);
+
+    // Caller holds the reporter lifecycle lease and nodes_mutex_. This method
+    // hands protection to status_mutex, releases the global node-table lock,
+    // and publishes status and metrics while the lifecycle lease pins NodeInfo.
+    void PublishHeartbeatStatus(NodeInfo &info,
+                                const std::map<std::string, std::string> &system_status,
+                                std::unique_lock<std::shared_mutex> &nodes_lock);
+
     void LivenessCheckerLoop();
     void ClearNodeGauges(const NodeInfo &info);
     static int64_t NowMillis() {

kv_cache_manager/data_storage/test/event_report_backend_test.cc

@@ -246,6 +246,7 @@ TEST_F(EventReportBackendTest, OnHeartbeatRefreshesAndRevivesNode) {
     EventReportBackend backend(metrics_registry_);
     ASSERT_EQ(EC_OK, backend.Open(MakeConfig(/*hb*/ 200, /*grace*/ 5000, /*tick*/ 50), "trace"));
     ASSERT_EQ(EC_OK, backend.RegisterNode("test_inst", "10.0.0.3:8080", {"mem"}));
+    const uint64_t registered_generation = backend.GetNodeGeneration("test_inst", "10.0.0.3:8080");
 
     int64_t initial_hb = 0;
     {
@@ -258,6 +259,7 @@ TEST_F(EventReportBackendTest, OnHeartbeatRefreshesAndRevivesNode) {
 
     std::this_thread::sleep_for(20ms);
     ASSERT_EQ(EC_OK, backend.OnHeartbeat("test_inst", "10.0.0.3:8080", {{"version", "er-0.18"}}));
+    ASSERT_EQ(registered_generation, backend.GetNodeGeneration("test_inst", "10.0.0.3:8080"));
     {
         auto &host_map = backend.instance_nodes_["test_inst"];
         auto it = host_map.find("10.0.0.3:8080");
@@ -268,6 +270,7 @@ TEST_F(EventReportBackendTest, OnHeartbeatRefreshesAndRevivesNode) {
     backend.SetNodeUnavailable("test_inst", "10.0.0.3:8080");
     ASSERT_FALSE(backend.IsNodeAvailable("test_inst", "10.0.0.3:8080"));
     ASSERT_EQ(EC_OK, backend.OnHeartbeat("test_inst", "10.0.0.3:8080", {}));
+    ASSERT_GT(backend.GetNodeGeneration("test_inst", "10.0.0.3:8080"), registered_generation);
     {
         auto &host_map = backend.instance_nodes_["test_inst"];
         auto it = host_map.find("10.0.0.3:8080");
@@ -283,6 +286,50 @@ TEST_F(EventReportBackendTest, OnHeartbeatRefreshesAndRevivesNode) {
     ASSERT_EQ(EC_OK, backend.Close());
 }
 
+TEST_F(EventReportBackendTest, SteadyHeartbeatDoesNotRejectConcurrentLifecycleMutationLease) {
+    EventReportBackend backend(metrics_registry_);
+    const ReporterSnapshotKey reporter_key{"heartbeat-mutation-lease", "10.0.0.32:8080"};
+    ASSERT_EQ(EC_OK, backend.RegisterNode(reporter_key.instance_id, reporter_key.host_ip_port, {"mem"}));
+    const uint64_t generation = backend.GetNodeGeneration(reporter_key.instance_id, reporter_key.host_ip_port);
+
+    auto &node = *backend.instance_nodes_[reporter_key.instance_id][reporter_key.host_ip_port];
+    node.last_heartbeat_ms.store(0, std::memory_order_relaxed);
+
+    // Pin HEARTBEAT after it refreshes the timestamp but before it publishes
+    // system status. At that point it still holds its lifecycle lease, so the
+    // concurrent mutation below exercises the exact production overlap
+    // without relying on scheduler timing or sleeps.
+    std::unique_lock<std::mutex> status_gate(node.status_mutex);
+    auto heartbeat = std::async(std::launch::async, [&] {
+        return backend.OnHeartbeat(reporter_key.instance_id, reporter_key.host_ip_port, {{"load", "1"}});
+    });
+
+    const auto deadline = std::chrono::steady_clock::now() + 1s;
+    while (node.last_heartbeat_ms.load(std::memory_order_acquire) == 0 && std::chrono::steady_clock::now() < deadline) {
+        std::this_thread::yield();
+    }
+    const bool heartbeat_reached_status_gate = node.last_heartbeat_ms.load(std::memory_order_acquire) != 0;
+    if (!heartbeat_reached_status_gate) {
+        status_gate.unlock();
+        EXPECT_EQ(EC_OK, heartbeat.get());
+        FAIL() << "heartbeat did not reach the status publication gate";
+    }
+
+    // A steady heartbeat is a lifecycle reader, not an unfenced operation:
+    // lifecycle writers must still wait until its status publication ends.
+    const auto lifecycle_fence = backend.GetOrCreateLifecycleFence(reporter_key);
+    std::unique_lock<std::shared_mutex> lifecycle_writer(lifecycle_fence->mutex, std::try_to_lock);
+    EXPECT_FALSE(lifecycle_writer.owns_lock());
+
+    EventReportBackend::LifecycleMutationLease mutation_lease;
+    const ErrorCode mutation_ec = backend.AcquireLifecycleMutationLease(reporter_key, generation, mutation_lease);
+
+    mutation_lease.reset();
+    status_gate.unlock();
+    EXPECT_EQ(EC_OK, heartbeat.get());
+    EXPECT_EQ(EC_OK, mutation_ec);
+}
+
 TEST_F(EventReportBackendTest, DataMutationsDoNotRefreshHeartbeat) {
     EventReportBackend backend(metrics_registry_);
     ASSERT_EQ(EC_OK, backend.Open(MakeConfig(/*hb*/ 5000, /*grace*/ 10000, /*tick*/ 50), "trace"));