main
use std::collections::BTreeMap;
use std::pin::Pin;
use std::sync::Arc;
use proto::grpc::health::v1::health_check_response::ServingStatus;
use proto::grpc::health::v1::health_server::{Health as HealthTrait, HealthServer};
use proto::grpc::health::v1::{HealthCheckRequest, HealthCheckResponse};
use tonic::{Request, Response, Status};
pub const HEALTH_SERVICE: &str = "grpc.health.v1.Health";
pub const REFLECTION_SERVICE: &str = "grpc.reflection.v1.ServerReflection";
/// 标准 gRPC 健康检查：整体（""）与每个已注册 service 都返回 SERVING。
#[derive(Clone)]
pub struct HealthService {
statuses: Arc<BTreeMap<String, i32>>,
}
impl HealthService {
pub fn new(service_names: &[&'static str]) -> Self {
let mut statuses = BTreeMap::new();
statuses.insert(String::new(), ServingStatus::Serving as i32);
for name in service_names {
statuses.insert((*name).to_string(), ServingStatus::Serving as i32);
}
Self {
statuses: Arc::new(statuses),
}
}
pub fn server(self) -> HealthServer<Self> {
HealthServer::new(self)
}
fn status(&self, service: &str) -> i32 {
self.statuses
.get(service)
.copied()
.unwrap_or(ServingStatus::ServiceUnknown as i32)
}
}
#[tonic::async_trait]
impl HealthTrait for HealthService {
type WatchStream =
Pin<Box<dyn tokio_stream::Stream<Item = Result<HealthCheckResponse, Status>> + Send>>;
async fn check(
&self,
request: Request<HealthCheckRequest>,
) -> Result<Response<HealthCheckResponse>, Status> {
let service = request.into_inner().service;
let status = self.status(&service);
if status == ServingStatus::ServiceUnknown as i32 {
return Err(Status::not_found("unknown service"));
}
Ok(Response::new(HealthCheckResponse { status }))
}
async fn watch(
&self,
request: Request<HealthCheckRequest>,
) -> Result<Response<Self::WatchStream>, Status> {
let service = request.into_inner().service;
let status = self.status(&service);
if status == ServingStatus::ServiceUnknown as i32 {
return Err(Status::not_found("unknown service"));
}
let (tx, rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let _ = tx.send(Ok(HealthCheckResponse { status })).await;
// 状态不变：保持流打开直到客户端取消。
futures_util::future::pending::<()>().await;
});
Ok(Response::new(Box::pin(
tokio_stream::wrappers::ReceiverStream::new(rx),
)))
}
}