rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-24 17:00:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: make grpc service able to be added dynamically (#3160)

Browse Source
This commit is contained in:
LFC
2024-01-15 12:33:27 +08:00
committed by GitHub
parent ca4d690424
commit 93f28c2a37
16 changed files with 245 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
src/servers/src/grpc.rs
View File

@@ -23,70 +23,38 @@ pub mod region_server;
use std::net::SocketAddr;
use api::v1::greptime_database_server::GreptimeDatabaseServer;
use api::v1::health_check_server::{HealthCheck, HealthCheckServer};
use api::v1::prometheus_gateway_server::{PrometheusGateway, PrometheusGatewayServer};
#[cfg(feature = "testing")]
use api::v1::region::region_server::Region;
use api::v1::region::region_server::RegionServer;
use api::v1::{HealthCheckRequest, HealthCheckResponse};
#[cfg(feature = "testing")]
use arrow_flight::flight_service_server::FlightService;
use arrow_flight::flight_service_server::FlightServiceServer;
use async_trait::async_trait;
use auth::UserProviderRef;
use common_grpc::channel_manager::{
DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE, DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
};
use common_telemetry::logging::info;
use common_telemetry::{error, warn};
use futures::FutureExt;
use opentelemetry_proto::tonic::collector::metrics::v1::metrics_service_server::MetricsServiceServer;
use opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::TraceServiceServer;
use snafu::{ensure, OptionExt, ResultExt};
use tokio::net::TcpListener;
use tokio::sync::oneshot::{self, Receiver, Sender};
use tokio::sync::Mutex;
use tonic::transport::server::TcpIncoming;
use tonic::transport::server::{Routes, TcpIncoming};
use tonic::{Request, Response, Status};
use tonic_reflection::server::{ServerReflection, ServerReflectionServer};
use tower::ServiceBuilder;
use self::authorize::AuthMiddlewareLayer;
use self::flight::{FlightCraftRef, FlightCraftWrapper};
use self::otlp::OtlpService;
use self::prom_query_gateway::PrometheusGatewayService;
use self::region_server::RegionServerRequestHandler;
use crate::error::{
AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, TcpBindSnafu, TcpIncomingSnafu,
};
use crate::grpc::database::DatabaseService;
use crate::grpc::greptime_handler::GreptimeRequestHandler;
use crate::prometheus_handler::PrometheusHandlerRef;
use crate::query_handler::OpenTelemetryProtocolHandlerRef;
use crate::server::Server;
type TonicResult<T> = std::result::Result<T, Status>;
pub struct GrpcServer {
config: GrpcServerConfig,
// states
shutdown_tx: Mutex<Option<Sender<()>>>,
/// gRPC serving state receiver. Only present if the gRPC server is started.
/// Used to wait for the server to stop, performing the old blocking fashion.
serve_state: Mutex<Option<Receiver<Result<()>>>>,
user_provider: Option<UserProviderRef>,
// handlers
/// Handler for [DatabaseService] service.
database_handler: Option<GreptimeRequestHandler>,
/// Handler for Prometheus-compatible PromQL queries ([PrometheusGateway]). Only present for frontend server.
prometheus_handler: Option<PrometheusHandlerRef>,
/// Handler for [FlightService](arrow_flight::flight_service_server::FlightService).
flight_handler: Option<FlightCraftRef>,
/// Handler for [RegionServer].
region_server_handler: Option<RegionServerRequestHandler>,
/// Handler for OpenTelemetry Protocol (OTLP) requests.
otlp_handler: Option<OpenTelemetryProtocolHandlerRef>,
routes: Mutex<Option<Routes>>,
}
/// Grpc Server configuration
@@ -108,16 +76,6 @@ impl Default for GrpcServerConfig {
}
impl GrpcServer {
#[cfg(feature = "testing")]
pub fn create_flight_service(&self) -> FlightServiceServer<impl FlightService> {
FlightServiceServer::new(FlightCraftWrapper(self.flight_handler.clone().unwrap()))
}
#[cfg(feature = "testing")]
pub fn create_region_service(&self) -> RegionServer<impl Region> {
RegionServer::new(self.region_server_handler.clone().unwrap())
}
pub fn create_healthcheck_service(&self) -> HealthCheckServer<impl HealthCheck> {
HealthCheckServer::new(HealthCheckHandler)
}
@@ -132,16 +90,6 @@ impl GrpcServer {
.unwrap()
}
pub fn create_prom_query_gateway_service(
&self,
handler: PrometheusHandlerRef,
) -> PrometheusGatewayServer<impl PrometheusGateway> {
PrometheusGatewayServer::new(PrometheusGatewayService::new(
handler,
self.user_provider.clone(),
))
}
pub async fn wait_for_serve(&self) -> Result<()> {
let mut serve_state = self.serve_state.lock().await;
let rx = serve_state.take().context(InternalSnafu {
@@ -188,8 +136,17 @@ impl Server for GrpcServer {
}
async fn start(&self, addr: SocketAddr) -> Result<SocketAddr> {
let max_recv_message_size = self.config.max_recv_message_size;
let max_send_message_size = self.config.max_send_message_size;
let routes = {
let mut routes = self.routes.lock().await;
let Some(routes) = routes.take() else {
return AlreadyStartedSnafu {
server: self.name(),
}
.fail();
};
routes
};
let (tx, rx) = oneshot::channel();
let (incoming, addr) = {
let mut shutdown_tx = self.shutdown_tx.lock().await;
@@ -211,52 +168,10 @@ impl Server for GrpcServer {
(incoming, addr)
};
let mut builder = tonic::transport::Server::builder()
let builder = tonic::transport::Server::builder()
.add_routes(routes)
.add_service(self.create_healthcheck_service())
.add_service(self.create_reflection_service());
if let Some(database_handler) = &self.database_handler {
builder = builder.add_service(
GreptimeDatabaseServer::new(DatabaseService::new(database_handler.clone()))
.max_decoding_message_size(max_recv_message_size)
.max_encoding_message_size(max_send_message_size),
)
}
if let Some(prometheus_handler) = &self.prometheus_handler {
builder = builder
.add_service(self.create_prom_query_gateway_service(prometheus_handler.clone()))
}
if let Some(otlp_handler) = &self.otlp_handler {
let trace_server = ServiceBuilder::new()
.layer(AuthMiddlewareLayer::with(self.user_provider.clone()))
.service(TraceServiceServer::new(OtlpService::new(
otlp_handler.clone(),
)));
builder = builder.add_service(trace_server);
let metrics_server = ServiceBuilder::new()
.layer(AuthMiddlewareLayer::with(self.user_provider.clone()))
.service(MetricsServiceServer::new(OtlpService::new(
otlp_handler.clone(),
)));
builder = builder.add_service(metrics_server);
}
if let Some(flight_handler) = &self.flight_handler {
builder = builder.add_service(
FlightServiceServer::new(FlightCraftWrapper(flight_handler.clone()))
.max_decoding_message_size(max_recv_message_size)
.max_encoding_message_size(max_send_message_size),
)
}
if let Some(region_server_handler) = &self.region_server_handler {
builder = builder.add_service(
RegionServer::new(region_server_handler.clone())
.max_decoding_message_size(max_recv_message_size)
.max_encoding_message_size(max_send_message_size),
);
}
let (serve_state_tx, serve_state_rx) = oneshot::channel();
let mut serve_state = self.serve_state.lock().await;

124
src/servers/src/grpc/builder.rs
View File

@@ -14,101 +14,127 @@
use std::sync::Arc;
use api::v1::greptime_database_server::GreptimeDatabaseServer;
use api::v1::prometheus_gateway_server::PrometheusGatewayServer;
use api::v1::region::region_server::RegionServer;
use arrow_flight::flight_service_server::FlightServiceServer;
use auth::UserProviderRef;
use common_runtime::Runtime;
use opentelemetry_proto::tonic::collector::metrics::v1::metrics_service_server::MetricsServiceServer;
use opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::TraceServiceServer;
use tokio::sync::Mutex;
use tonic::transport::server::RoutesBuilder;
use tower::ServiceBuilder;
use super::flight::FlightCraftRef;
use super::flight::{FlightCraftRef, FlightCraftWrapper};
use super::region_server::{RegionServerHandlerRef, RegionServerRequestHandler};
use super::{GrpcServer, GrpcServerConfig};
use crate::grpc::authorize::AuthMiddlewareLayer;
use crate::grpc::database::DatabaseService;
use crate::grpc::greptime_handler::GreptimeRequestHandler;
use crate::grpc::otlp::OtlpService;
use crate::grpc::prom_query_gateway::PrometheusGatewayService;
use crate::prometheus_handler::PrometheusHandlerRef;
use crate::query_handler::OpenTelemetryProtocolHandlerRef;
macro_rules! add_service {
($builder: ident, $service: expr) => {
$builder.routes_builder.add_service(
$service
.max_decoding_message_size($builder.config.max_recv_message_size)
.max_encoding_message_size($builder.config.max_send_message_size),
)
};
}
pub struct GrpcServerBuilder {
config: Option<GrpcServerConfig>,
database_handler: Option<GreptimeRequestHandler>,
prometheus_handler: Option<PrometheusHandlerRef>,
flight_handler: Option<FlightCraftRef>,
region_server_handler: Option<RegionServerHandlerRef>,
otlp_handler: Option<OpenTelemetryProtocolHandlerRef>,
user_provider: Option<UserProviderRef>,
config: GrpcServerConfig,
runtime: Arc<Runtime>,
routes_builder: RoutesBuilder,
}
impl GrpcServerBuilder {
pub fn new(runtime: Arc<Runtime>) -> Self {
pub fn new(config: GrpcServerConfig, runtime: Arc<Runtime>) -> Self {
Self {
config: None,
database_handler: None,
prometheus_handler: None,
flight_handler: None,
region_server_handler: None,
otlp_handler: None,
user_provider: None,
config,
runtime,
routes_builder: RoutesBuilder::default(),
}
}
pub fn config(mut self, config: GrpcServerConfig) -> Self {
self.config = Some(config);
self
}
pub fn option_config(mut self, config: Option<GrpcServerConfig>) -> Self {
self.config = config;
self
}
pub fn runtime(&self) -> &Arc<Runtime> {
&self.runtime
}
/// Add handler for [DatabaseService] service.
pub fn database_handler(mut self, database_handler: GreptimeRequestHandler) -> Self {
self.database_handler = Some(database_handler);
add_service!(
self,
GreptimeDatabaseServer::new(DatabaseService::new(database_handler))
);
self
}
pub fn prometheus_handler(mut self, prometheus_handler: PrometheusHandlerRef) -> Self {
self.prometheus_handler = Some(prometheus_handler);
/// Add handler for Prometheus-compatible PromQL queries ([PrometheusGateway]).
pub fn prometheus_handler(
mut self,
prometheus_handler: PrometheusHandlerRef,
user_provider: Option<UserProviderRef>,
) -> Self {
add_service!(
self,
PrometheusGatewayServer::new(PrometheusGatewayService::new(
prometheus_handler,
user_provider,
))
);
self
}
/// Add handler for [FlightService](arrow_flight::flight_service_server::FlightService).
pub fn flight_handler(mut self, flight_handler: FlightCraftRef) -> Self {
self.flight_handler = Some(flight_handler);
add_service!(
self,
FlightServiceServer::new(FlightCraftWrapper(flight_handler.clone()))
);
self
}
/// Add handler for [RegionServer].
pub fn region_server_handler(mut self, region_server_handler: RegionServerHandlerRef) -> Self {
self.region_server_handler = Some(region_server_handler);
let handler = RegionServerRequestHandler::new(region_server_handler, self.runtime.clone());
add_service!(self, RegionServer::new(handler));
self
}
pub fn otlp_handler(mut self, otlp_handler: OpenTelemetryProtocolHandlerRef) -> Self {
self.otlp_handler = Some(otlp_handler);
/// Add handler for OpenTelemetry Protocol (OTLP) requests.
pub fn otlp_handler(
mut self,
otlp_handler: OpenTelemetryProtocolHandlerRef,
user_provider: Option<UserProviderRef>,
) -> Self {
let trace_server = ServiceBuilder::new()
.layer(AuthMiddlewareLayer::with(user_provider.clone()))
.service(TraceServiceServer::new(OtlpService::new(
otlp_handler.clone(),
)));
self.routes_builder.add_service(trace_server);
let metrics_server = ServiceBuilder::new()
.layer(AuthMiddlewareLayer::with(user_provider))
.service(MetricsServiceServer::new(OtlpService::new(otlp_handler)));
self.routes_builder.add_service(metrics_server);
self
}
pub fn user_provider(mut self, user_provider: Option<UserProviderRef>) -> Self {
self.user_provider = user_provider;
self
pub fn routes_builder_mut(&mut self) -> &mut RoutesBuilder {
&mut self.routes_builder
}
pub fn build(self) -> GrpcServer {
let config = self.config.unwrap_or_default();
let runtime = self.runtime;
let region_server_handler = self
.region_server_handler
.map(|handler| RegionServerRequestHandler::new(handler, runtime));
GrpcServer {
config,
prometheus_handler: self.prometheus_handler,
flight_handler: self.flight_handler,
region_server_handler,
database_handler: self.database_handler,
otlp_handler: self.otlp_handler,
user_provider: self.user_provider,
routes: Mutex::new(Some(self.routes_builder.routes())),
shutdown_tx: Mutex::new(None),
serve_state: Mutex::new(None),
}