[event] harden optimizer event stream lifecycle

Author: Tyndalllll

AGENTS.md

@@ -19,11 +19,12 @@ flowchart TD
     data_storage --> common["common"]
 
     manager --> event["event"]
+    event --> protocol["protocol（proto/grpc）"]
     manager --> metrics["metrics"]
     service --> metrics
     service --> config
     service --> data_storage
-    manager --> protocol["protocol（proto/grpc）"]
+    manager --> protocol
 
     %% 有意的反向边（近似环，改动需谨慎）
     common -. request_context .-> metrics

docs/design/module_architecture.md

@@ -43,8 +43,8 @@ KVCache Manager 采用中心化部署，负责 KVCache 的全局元数据管理
 |---|---|---|
 | **common** | `common/` | 基础设施层：日志、JSON、错误码、Redis 客户端、`RequestContext`（逐请求追踪上下文）、`concurrent_hash_map`、`lru_cache`、`loop_thread`、服务发现、崩溃处理等。几乎所有 C++ 模块都依赖它。 |
 | **metrics** | `metrics/` | 可观测性。`MetricsRegistry`/`MetricsCollector` 收集指标，多种 reporter（kmonitor/local/logging/dummy）上报，`PrometheusExporter` 通过 HTTP 暴露。 |
-| **event** | `event/` | 轻量事件总线。`EventManager` 将领域事件（如 cache 回收事件）分发给注册的 `EventPublisher`（默认 `LogEventPublisher`）。 |
-| **protocol** | `protocol/protobuf/` | gRPC/proto 契约。定义 meta/admin/debug/kv_meta 服务，生成 C++ 与 Python 桩。 |
+| **event** | `event/` | 轻量事件总线。`EventManager` 将领域事件分发给注册的 `EventPublisher`；除默认日志发布器外，`OptimizerEventPublisher` 会将缓存读取事件转换为 protocol 中的 `TraceQueryRequest`，再经 `SubscriptionEventSink` 交给 service 层的 gRPC 流。 |
+| **protocol** | `protocol/protobuf/` | gRPC/proto 契约。定义 meta/admin/debug/kv_meta 以及 optimizer 事件流服务，生成 C++ 与 Python 桩。 |
 
 ### 客户端与连接器
 
@@ -83,7 +83,7 @@ client 通过 `InitParams.role_type` 区分角色：**SCHEDULER**（调度节点
 service → manager → meta → config → data_storage → common
 ```
 
-- `common`、`protocol` 是最底层的通用模块，被各层广泛依赖。
+- `common`、`protocol` 是最底层的通用模块，被各层广泛依赖；`event` 的 optimizer 发布链路直接使用 `protocol` 定义的 `TraceQueryRequest`。
 - `metrics`、`event` 是通用支撑模块，被 `manager` 与 `service` 复用。
 - `service` 在启动时实例化 `CacheManager`（注入 `MetricsRegistry` + `RegistryManager`），并通过 `config` 的 `LeaderElector` 门控 recover/cleanup。
 
@@ -138,6 +138,7 @@ flowchart TD
 
     %% 通用支撑依赖
     manager --> event
+    event --> protocol
     manager --> metrics
     service --> metrics
     service --> event

kv_cache_manager/event/optimizer_event_publisher.cc

@@ -67,9 +67,8 @@ bool OptimizerEventPublisher::Stop() {
 }
 
 void OptimizerEventPublisher::WorkerThread() {
-    // Single threaded on purpose: the replay requires per-instance events to
-    // arrive in non-decreasing timestamp order, and one worker draining one
-    // queue preserves that ordering.
+    // A single worker keeps conversion and delivery off serving threads
+    // without adding synchronization between multiple consumers of the queue.
     while (running_) {
         BasicWait();
 

kv_cache_manager/event/optimizer_event_publisher.h

@@ -28,9 +28,8 @@ class TraceQueryRequest;
 // the event is dropped rather than the caller blocked: these are analysis
 // samples, and a slow or absent consumer must never slow down serving.
 //
-// Dropping biases the replayed hit rate downwards (an unseen request never
-// touches the simulated LRU). That is accepted, but it means the drop count
-// has to be readable before anyone trusts a capacity curve.
+// Dropping samples may slightly affect replay accuracy, which is accepted for
+// this best-effort analysis path.
 class OptimizerEventPublisher : public EventPublisher {
 public:
     OptimizerEventPublisher(std::shared_ptr<EventSink> sink, const OptimizerEventPublisherConfig &config);

kv_cache_manager/event/optimizer_stream/subscription_event_sink.cc

@@ -58,7 +58,7 @@ SubscriptionEventSink::~SubscriptionEventSink() { Stop(); }
 
 std::shared_ptr<SubscriptionEventSink::Subscription> SubscriptionEventSink::Subscribe(const std::string &consumer_id) {
     std::lock_guard<std::mutex> lock(subscriptions_mutex_);
-    if (stopped_ || subscriptions_.size() >= config_.max_subscribers()) {
+    if (stopped_ || !accepting_subscriptions_ || subscriptions_.size() >= config_.max_subscribers()) {
         return nullptr;
     }
     auto subscription = std::shared_ptr<Subscription>(
@@ -87,15 +87,33 @@ void SubscriptionEventSink::Unsubscribe(const std::shared_ptr<Subscription> &sub
                   subscribers);
 }
 
-bool SubscriptionEventSink::Send(const proto::optimizer::TraceQueryRequest &event) {
-    if (stopped_) {
-        dropped_.fetch_add(1);
-        return false;
+void SubscriptionEventSink::EnableSubscriptions() {
+    std::lock_guard<std::mutex> lock(subscriptions_mutex_);
+    if (!stopped_) {
+        accepting_subscriptions_ = true;
     }
+}
 
+void SubscriptionEventSink::DisableSubscriptions() {
     std::vector<std::shared_ptr<Subscription>> subscriptions;
     {
         std::lock_guard<std::mutex> lock(subscriptions_mutex_);
+        accepting_subscriptions_ = false;
+        subscriptions.swap(subscriptions_);
+    }
+    for (const auto &subscription : subscriptions) {
+        subscription->Close();
+    }
+}
+
+bool SubscriptionEventSink::Send(const proto::optimizer::TraceQueryRequest &event) {
+    std::vector<std::shared_ptr<Subscription>> subscriptions;
+    {
+        std::lock_guard<std::mutex> lock(subscriptions_mutex_);
+        if (stopped_ || !accepting_subscriptions_) {
+            dropped_.fetch_add(1);
+            return false;
+        }
         subscriptions = subscriptions_;
     }
     if (subscriptions.empty()) {
@@ -118,6 +136,7 @@ void SubscriptionEventSink::Stop() {
     if (stopped_.exchange(true)) {
         return;
     }
+    accepting_subscriptions_ = false;
     std::vector<std::shared_ptr<Subscription>> subscriptions;
     {
         std::lock_guard<std::mutex> lock(subscriptions_mutex_);

kv_cache_manager/event/optimizer_stream/subscription_event_sink.h

@@ -56,6 +56,10 @@ class SubscriptionEventSink final : public EventSink {
     std::shared_ptr<Subscription> Subscribe(const std::string &consumer_id);
     void Unsubscribe(const std::shared_ptr<Subscription> &subscription);
 
+    void EnableSubscriptions();
+    void DisableSubscriptions();
+    bool accepting_subscriptions() const { return accepting_subscriptions_.load(); }
+
     bool Send(const proto::optimizer::TraceQueryRequest &event) override;
     void Stop() override;
     std::size_t DroppedCount() const;
@@ -67,6 +71,7 @@ class SubscriptionEventSink final : public EventSink {
     OptimizerEventPublisherConfig config_;
     mutable std::mutex subscriptions_mutex_;
     std::vector<std::shared_ptr<Subscription>> subscriptions_;
+    std::atomic<bool> accepting_subscriptions_{true};
     std::atomic<bool> stopped_{false};
     std::atomic<std::size_t> dropped_{0};
 };

kv_cache_manager/event/test/subscription_event_sink_test.cc

@@ -77,6 +77,27 @@ TEST_F(SubscriptionEventSinkTest, TestRejectsExcessSubscribers) {
     EXPECT_EQ(nullptr, sink.Subscribe("second"));
 }
 
+TEST_F(SubscriptionEventSinkTest, TestDisableClosesSubscribersAndEnableAcceptsNewOnes) {
+    SubscriptionEventSink sink(MakeSinkConfig(1, 1));
+    auto subscription = sink.Subscribe("optimizer");
+    ASSERT_NE(nullptr, subscription);
+
+    sink.DisableSubscriptions();
+    proto::optimizer::TraceQueryRequest event;
+    EXPECT_EQ(SubscriptionEventSink::Subscription::WaitResult::kClosed,
+              subscription->WaitNext(&event, std::chrono::milliseconds(10)));
+    EXPECT_EQ(nullptr, sink.Subscribe("disabled"));
+    EXPECT_FALSE(sink.Send(MakeRequest("disabled")));
+
+    sink.EnableSubscriptions();
+    auto resumed = sink.Subscribe("resumed");
+    ASSERT_NE(nullptr, resumed);
+    EXPECT_TRUE(sink.Send(MakeRequest("resumed")));
+    EXPECT_EQ(SubscriptionEventSink::Subscription::WaitResult::kEvent,
+              resumed->WaitNext(&event, std::chrono::milliseconds(10)));
+    EXPECT_EQ("resumed", event.trace_id());
+}
+
 TEST_F(SubscriptionEventSinkTest, TestStopWakesSubscriberAndIsIdempotent) {
     SubscriptionEventSink sink(MakeSinkConfig(1, 1));
     auto subscription = sink.Subscribe("optimizer");

kv_cache_manager/service/grpc_service/BUILD

@@ -36,6 +36,7 @@ cc_library(
         "//kv_cache_manager/common",
         "//kv_cache_manager/common:logger",
         "//kv_cache_manager/config",
+        "//kv_cache_manager/config:leader_elector_base",
         "//kv_cache_manager/event/optimizer_stream:subscription_event_sink",
         "//kv_cache_manager/protocol/protobuf:service_cc_grpc",
         "//kv_cache_manager/service:grpc_wrapper",

kv_cache_manager/service/grpc_service/optimizer_event_service_grpc.cc

@@ -7,22 +7,46 @@
 #include "kv_cache_manager/common/request_context.h"
 #include "kv_cache_manager/config/instance_group.h"
 #include "kv_cache_manager/config/instance_info.h"
+#include "kv_cache_manager/config/leader_elector.h"
 #include "kv_cache_manager/config/registry_manager.h"
 #include "kv_cache_manager/event/optimizer_stream/subscription_event_sink.h"
 
 namespace kv_cache_manager {
 
 OptimizerEventServiceGRpc::OptimizerEventServiceGRpc(std::shared_ptr<SubscriptionEventSink> sink,
-                                                     std::shared_ptr<RegistryManager> registry_manager)
-    : sink_(std::move(sink)), registry_manager_(std::move(registry_manager)) {}
+                                                     std::shared_ptr<RegistryManager> registry_manager,
+                                                     std::shared_ptr<LeaderElector> leader_elector)
+    : sink_(std::move(sink))
+    , registry_manager_(std::move(registry_manager))
+    , leader_elector_(std::move(leader_elector)) {
+    DisableSubscriptions();
+}
+
+void OptimizerEventServiceGRpc::EnableSubscriptions() {
+    if (sink_) {
+        sink_->EnableSubscriptions();
+    }
+}
+
+void OptimizerEventServiceGRpc::DisableSubscriptions() {
+    if (sink_) {
+        sink_->DisableSubscriptions();
+    }
+}
+
+bool OptimizerEventServiceGRpc::IsAvailable() const {
+    return leader_elector_ && leader_elector_->GetRoleState() == RoleState::LEADER &&
+           leader_elector_->IsStableState() && registry_manager_ && registry_manager_->IsRecoverComplete() && sink_ &&
+           sink_->accepting_subscriptions() && !sink_->stopped();
+}
 
 grpc::Status OptimizerEventServiceGRpc::GetConfiguration(grpc::ServerContext *,
                                                          const proto::optimizer::KvcmConfigurationRequest *request,
                                                          proto::optimizer::KvcmConfigurationResponse *response) {
     auto *status = response->mutable_header()->mutable_status();
-    if (!registry_manager_) {
+    if (!IsAvailable()) {
         status->set_code(proto::optimizer::SERVICE_NOT_READY);
-        status->set_message("KVCM registry manager is unavailable");
+        status->set_message("KVCM is unavailable");
         return grpc::Status::OK;
     }
 
@@ -76,24 +100,26 @@ grpc::Status
 OptimizerEventServiceGRpc::SubscribeEvents(grpc::ServerContext *context,
                                            const proto::optimizer::OptimizerEventSubscriptionRequest *request,
                                            grpc::ServerWriter<proto::optimizer::TraceQueryRequest> *writer) {
-    if (!sink_ || sink_->stopped()) {
-        return grpc::Status(grpc::StatusCode::UNAVAILABLE, "optimizer event publisher is unavailable");
+    if (!IsAvailable()) {
+        return grpc::Status(grpc::StatusCode::UNAVAILABLE, "KVCM is unavailable");
     }
     auto subscription = sink_->Subscribe(request->consumer_id());
     if (!subscription) {
-        return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED, "optimizer subscriber limit reached");
+        return grpc::Status(grpc::StatusCode::UNAVAILABLE, "KVCM is unavailable");
     }
 
     KVCM_LOG_INFO("OptimizerEventServiceGRpc: stream opened, consumer_id=%s, peer=%s",
                   subscription->consumer_id().c_str(),
                   context->peer().c_str());
+    bool subscription_closed = false;
     while (!context->IsCancelled()) {
         proto::optimizer::TraceQueryRequest event;
         const auto result = subscription->WaitNext(&event, std::chrono::milliseconds(100));
         if (result == SubscriptionEventSink::Subscription::WaitResult::kTimeout) {
             continue;
         }
         if (result == SubscriptionEventSink::Subscription::WaitResult::kClosed) {
+            subscription_closed = true;
             break;
         }
         if (!writer->Write(event)) {
@@ -102,6 +128,9 @@ OptimizerEventServiceGRpc::SubscribeEvents(grpc::ServerContext *context,
     }
     sink_->Unsubscribe(subscription);
     KVCM_LOG_INFO("OptimizerEventServiceGRpc: stream closed, consumer_id=%s", subscription->consumer_id().c_str());
+    if (subscription_closed) {
+        return grpc::Status(grpc::StatusCode::UNAVAILABLE, "KVCM is unavailable");
+    }
     return grpc::Status::OK;
 }
 

kv_cache_manager/service/grpc_service/optimizer_event_service_grpc.h

@@ -9,12 +9,17 @@ namespace kv_cache_manager {
 
 class SubscriptionEventSink;
 class RegistryManager;
+class LeaderElector;
 
 // gRPC service for optimizer event subscriptions.
 class OptimizerEventServiceGRpc final : public proto::optimizer::OptimizerEventStreamService::Service {
 public:
     OptimizerEventServiceGRpc(std::shared_ptr<SubscriptionEventSink> sink,
-                              std::shared_ptr<RegistryManager> registry_manager);
+                              std::shared_ptr<RegistryManager> registry_manager,
+                              std::shared_ptr<LeaderElector> leader_elector);
+
+    void EnableSubscriptions();
+    void DisableSubscriptions();
 
     grpc::Status GetConfiguration(grpc::ServerContext *context,
                                   const proto::optimizer::KvcmConfigurationRequest *request,
@@ -25,8 +30,11 @@ class OptimizerEventServiceGRpc final : public proto::optimizer::OptimizerEventS
                                  grpc::ServerWriter<proto::optimizer::TraceQueryRequest> *writer) override;
 
 private:
+    bool IsAvailable() const;
+
     std::shared_ptr<SubscriptionEventSink> sink_;
     std::shared_ptr<RegistryManager> registry_manager_;
+    std::shared_ptr<LeaderElector> leader_elector_;
 };
 
 } // namespace kv_cache_manager