use crate::config::GrpcConfig;
use crate::error::Result;
use crate::state::AppState;
use std::net::SocketAddr;
use tonic::server::NamedService;
use tonic::transport::Server;
#[derive(Debug)]
pub struct GrpcServer {
config: GrpcConfig,
}
impl GrpcServer {
pub fn new(config: GrpcConfig) -> Self {
Self { config }
}
pub fn build(self) -> Result<Server> {
let server = Server::builder()
.max_frame_size(Some(self.config.max_message_size_bytes() as u32))
.timeout(self.config.timeout())
.tcp_keepalive(Some(std::time::Duration::from_secs(60)));
Ok(server)
}
pub fn socket_addr(&self, http_port: u16) -> SocketAddr {
let port = self.config.effective_port(http_port);
SocketAddr::from(([0, 0, 0, 0], port))
}
}
pub struct GrpcServicesBuilder {
routes: tonic::service::Routes,
reflection_enabled: bool,
health_enabled: bool,
file_descriptor_sets: Vec<&'static [u8]>,
}
impl GrpcServicesBuilder {
pub fn new() -> Self {
Self {
routes: tonic::service::Routes::default(),
reflection_enabled: false,
health_enabled: false,
file_descriptor_sets: Vec::new(),
}
}
pub fn with_reflection(mut self) -> Self {
self.reflection_enabled = true;
self
}
pub fn with_health(mut self) -> Self {
self.health_enabled = true;
self
}
pub fn add_file_descriptor_set(mut self, file_descriptor_set: &'static [u8]) -> Self {
self.file_descriptor_sets.push(file_descriptor_set);
self
}
pub fn add_service<S>(mut self, service: S) -> Self
where
S: tower::Service<
http::Request<tonic::body::Body>,
Response = http::Response<tonic::body::Body>,
Error = std::convert::Infallible,
> + NamedService
+ Clone
+ Send
+ Sync
+ 'static,
S::Future: Send + 'static,
{
self.routes = self.routes.add_service(service);
self
}
pub fn build(mut self, state: Option<AppState>) -> tonic::service::Routes {
if self.health_enabled {
if let Some(app_state) = state.clone() {
let health_service = crate::grpc::HealthService::new(app_state);
let health_server =
tonic_health::pb::health_server::HealthServer::new(health_service);
self.routes = self.routes.add_service(health_server);
tracing::info!("gRPC health service enabled");
} else {
tracing::warn!(
"Health service enabled but no AppState provided, skipping health service"
);
}
}
if self.reflection_enabled {
if self.file_descriptor_sets.is_empty() {
tracing::warn!("Reflection enabled but no file descriptor sets registered. Use add_file_descriptor_set() to register services.");
} else {
let mut reflection_builder = tonic_reflection::server::Builder::configure();
for file_descriptor_set in self.file_descriptor_sets {
reflection_builder = reflection_builder
.register_encoded_file_descriptor_set(file_descriptor_set);
}
match reflection_builder.build_v1() {
Ok(reflection_service) => {
self.routes = self.routes.add_service(reflection_service);
tracing::info!("gRPC reflection service enabled");
}
Err(e) => {
tracing::error!("Failed to build reflection service: {}", e);
}
}
}
}
self.routes
}
}
impl Default for GrpcServicesBuilder {
fn default() -> Self {
Self::new()
}
}