3788a30
use std::convert::Infallible;
use std::net::SocketAddr;
use axum::response::IntoResponse;
use http::Request;
use tonic::body::Body;
use tonic::server::NamedService;
use tonic::transport::Server;
use tonic::transport::server::Router;
use tower::Service;
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>;
enum RouterState {
Pending(Server<Layer>),
Ready(Router<Layer>),
}
/// 统一装配后的内部 RPC 路由：记录业务 crate 注册的 service 名，
/// health / reflection 在 [`serve`] 时据此生成，避免手工清单漂移。
pub struct InternalRouter {
state: RouterState,
services: Vec<&'static str>,
}
impl InternalRouter {
pub fn add_service<S>(mut self, service: S) -> Self
where
S: Service<Request<Body>, Error = Infallible>
+ NamedService
+ Clone
+ Send
+ Sync
+ 'static,
S::Response: IntoResponse,
S::Future: Send + 'static,
{
self.services.push(S::NAME);
self.state = match self.state {
RouterState::Pending(mut server) => RouterState::Ready(server.add_service(service)),
RouterState::Ready(router) => RouterState::Ready(router.add_service(service)),
};
self
}
}
#[derive(Debug, Clone)]
pub struct ServerConfig {
pub listen_address: String,
pub rpc_key: String,
}
impl ServerConfig {
pub fn new(listen_address: impl Into<String>, rpc_key: impl Into<String>) -> Self {
Self {
listen_address: listen_address.into(),
rpc_key: rpc_key.into(),
}
}
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(),
));
}
Ok(InternalRouter {
state: RouterState::Pending(Server::builder().layer(oltp::GrpcRequestLog)),
services: Vec::new(),
})
}
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 InternalRouter { state, services } = router;
let health = HealthService::new(&services).server();
let reflection = ReflectionService::new(&services)?.server();
let router = match state {
RouterState::Pending(mut server) => server.add_service(health).add_service(reflection),
RouterState::Ready(router) => router.add_service(health).add_service(reflection),
};
let address = parse_address(listen)?;
router
.serve_with_shutdown(address, shutdown)
.await
.map_err(InternalError::Transport)
}