|
| 1 | +#include "kv_cache_manager/service/grpc_service/optimizer_event_service_grpc.h" |
| 2 | + |
| 3 | +#include <chrono> |
| 4 | +#include <utility> |
| 5 | + |
| 6 | +#include "kv_cache_manager/common/logger.h" |
| 7 | +#include "kv_cache_manager/common/request_context.h" |
| 8 | +#include "kv_cache_manager/config/instance_group.h" |
| 9 | +#include "kv_cache_manager/config/instance_info.h" |
| 10 | +#include "kv_cache_manager/config/registry_manager.h" |
| 11 | +#include "kv_cache_manager/event/optimizer_stream/subscription_event_sink.h" |
| 12 | + |
| 13 | +namespace kv_cache_manager { |
| 14 | + |
| 15 | +OptimizerEventServiceGRpc::OptimizerEventServiceGRpc(std::shared_ptr<SubscriptionEventSink> sink, |
| 16 | + std::shared_ptr<RegistryManager> registry_manager) |
| 17 | + : sink_(std::move(sink)), registry_manager_(std::move(registry_manager)) {} |
| 18 | + |
| 19 | +grpc::Status OptimizerEventServiceGRpc::GetConfiguration(grpc::ServerContext *, |
| 20 | + const proto::optimizer::KvcmConfigurationRequest *request, |
| 21 | + proto::optimizer::KvcmConfigurationResponse *response) { |
| 22 | + auto *status = response->mutable_header()->mutable_status(); |
| 23 | + if (!registry_manager_) { |
| 24 | + status->set_code(proto::optimizer::SERVICE_NOT_READY); |
| 25 | + status->set_message("KVCM registry manager is unavailable"); |
| 26 | + return grpc::Status::OK; |
| 27 | + } |
| 28 | + |
| 29 | + RequestContext request_context(request->trace_id()); |
| 30 | + const auto [group_ec, instance_groups] = registry_manager_->ListInstanceGroup(&request_context); |
| 31 | + if (group_ec != EC_OK) { |
| 32 | + status->set_code(proto::optimizer::INTERNAL_ERROR); |
| 33 | + status->set_message("Failed to list KVCM instance groups"); |
| 34 | + return grpc::Status::OK; |
| 35 | + } |
| 36 | + |
| 37 | + for (const auto &instance_group : instance_groups) { |
| 38 | + auto *group = response->add_instance_groups(); |
| 39 | + group->set_name(instance_group->name()); |
| 40 | + group->set_capacity_bytes(instance_group->quota().capacity()); |
| 41 | + |
| 42 | + const auto [instance_ec, instances] = |
| 43 | + registry_manager_->ListInstanceInfo(&request_context, instance_group->name()); |
| 44 | + if (instance_ec != EC_OK) { |
| 45 | + response->clear_instance_groups(); |
| 46 | + response->clear_instances(); |
| 47 | + status->set_code(proto::optimizer::INTERNAL_ERROR); |
| 48 | + status->set_message("Failed to list KVCM instances"); |
| 49 | + return grpc::Status::OK; |
| 50 | + } |
| 51 | + for (const auto &instance_info : instances) { |
| 52 | + auto *instance = response->add_instances(); |
| 53 | + instance->set_instance_group_name(instance_info->instance_group_name()); |
| 54 | + instance->set_instance_id(instance_info->instance_id()); |
| 55 | + instance->set_block_size(instance_info->block_size()); |
| 56 | + for (const auto &spec_info : instance_info->location_spec_infos()) { |
| 57 | + auto *spec = instance->add_location_spec_infos(); |
| 58 | + spec->set_name(spec_info.name()); |
| 59 | + spec->set_size(spec_info.size()); |
| 60 | + } |
| 61 | + for (const auto &spec_group : instance_info->location_spec_groups()) { |
| 62 | + auto *group_config = instance->add_location_spec_groups(); |
| 63 | + group_config->set_name(spec_group.name()); |
| 64 | + for (const auto &spec_name : spec_group.spec_names()) { |
| 65 | + group_config->add_spec_names(spec_name); |
| 66 | + } |
| 67 | + } |
| 68 | + } |
| 69 | + } |
| 70 | + |
| 71 | + status->set_code(proto::optimizer::OK); |
| 72 | + return grpc::Status::OK; |
| 73 | +} |
| 74 | + |
| 75 | +grpc::Status |
| 76 | +OptimizerEventServiceGRpc::SubscribeEvents(grpc::ServerContext *context, |
| 77 | + const proto::optimizer::OptimizerEventSubscriptionRequest *request, |
| 78 | + grpc::ServerWriter<proto::optimizer::TraceQueryRequest> *writer) { |
| 79 | + if (!sink_ || sink_->stopped()) { |
| 80 | + return grpc::Status(grpc::StatusCode::UNAVAILABLE, "optimizer event publisher is unavailable"); |
| 81 | + } |
| 82 | + auto subscription = sink_->Subscribe(request->consumer_id()); |
| 83 | + if (!subscription) { |
| 84 | + return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED, "optimizer subscriber limit reached"); |
| 85 | + } |
| 86 | + |
| 87 | + KVCM_LOG_INFO("OptimizerEventServiceGRpc: stream opened, consumer_id=%s, peer=%s", |
| 88 | + subscription->consumer_id().c_str(), |
| 89 | + context->peer().c_str()); |
| 90 | + while (!context->IsCancelled()) { |
| 91 | + proto::optimizer::TraceQueryRequest event; |
| 92 | + const auto result = subscription->WaitNext(&event, std::chrono::milliseconds(100)); |
| 93 | + if (result == SubscriptionEventSink::Subscription::WaitResult::kTimeout) { |
| 94 | + continue; |
| 95 | + } |
| 96 | + if (result == SubscriptionEventSink::Subscription::WaitResult::kClosed) { |
| 97 | + break; |
| 98 | + } |
| 99 | + if (!writer->Write(event)) { |
| 100 | + break; |
| 101 | + } |
| 102 | + } |
| 103 | + sink_->Unsubscribe(subscription); |
| 104 | + KVCM_LOG_INFO("OptimizerEventServiceGRpc: stream closed, consumer_id=%s", subscription->consumer_id().c_str()); |
| 105 | + return grpc::Status::OK; |
| 106 | +} |
| 107 | + |
| 108 | +} // namespace kv_cache_manager |
0 commit comments