ec0bdb8
use std::net::SocketAddr;
use tonic::transport::Server;
use tonic::transport::server::Router;
use tower::layer::util::{Identity, Stack};
use crate::error::InternalError;
use crate::health::HealthService;
use crate::interceptor::Gate;
use crate::reflection::ReflectionService;
type Layer = Stack<oltp::GrpcRequestLog, Identity>;
/// 统一装配后的内部 RPC 路由:health + reflection + 日志层已就绪,
/// 业务 crate 只需继续 `add_service`。
pub type InternalRouter = Router<Layer>;
#[derive(Debug, Clone)]
pub struct ServerConfig {
pub listen_address: String,
pub rpc_key: String,
pub service_names: Vec<&'static str>,
}
impl ServerConfig {
pub fn new(
listen_address: impl Into<String>,
rpc_key: impl Into<String>,
service_names: Vec<&'static str>,
) -> Self {
Self {
listen_address: listen_address.into(),
rpc_key: rpc_key.into(),
service_names,
}
}
pub fn gate(&self) -> Gate {
Gate::new(self.rpc_key.clone())
}
pub fn listen_address(&self) -> &str {
&self.listen_address
}
}
pub fn router(config: &ServerConfig) -> Result<InternalRouter, InternalError> {
let gate = Gate::new(config.rpc_key.clone());
if !gate.is_armed() {
return Err(InternalError::Config(
"internal rpc key must not be blank".to_string(),
));
}
let health = HealthService::new(&config.service_names).server();
let reflection = ReflectionService::new(&config.service_names)?.server();
Ok(Server::builder()
.layer(oltp::GrpcRequestLog)
.add_service(health)
.add_service(reflection))
}
pub fn parse_address(listen: &str) -> Result<SocketAddr, InternalError> {
listen
.parse()
.map_err(|error| InternalError::Address(format!("{listen}: {error}")))
}
/// 启动内部 RPC 服务端;`shutdown` 完成即优雅停止。
pub async fn serve(
router: InternalRouter,
listen: &str,
shutdown: impl std::future::Future<Output = ()> + Send + 'static,
) -> Result<(), InternalError> {
let address = parse_address(listen)?;
router
.serve_with_shutdown(address, shutdown)
.await
.map_err(InternalError::Transport)
}