servers/grpc/
builder.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::convert::Infallible;
16
17use api::v1::frontend::frontend_server::FrontendServer;
18use api::v1::greptime_database_server::GreptimeDatabaseServer;
19use api::v1::prometheus_gateway_server::PrometheusGatewayServer;
20use api::v1::region::region_server::RegionServer;
21use arrow_flight::flight_service_server::FlightServiceServer;
22use auth::UserProviderRef;
23use axum::extract::Request;
24use axum::response::IntoResponse;
25use axum::routing::Route;
26use common_grpc::error::{InvalidConfigFilePathSnafu, Result};
27use common_runtime::Runtime;
28use common_telemetry::warn;
29use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsServiceServer;
30use snafu::ResultExt;
31use tokio::sync::Mutex;
32use tonic::codec::CompressionEncoding;
33use tonic::codegen::Service;
34use tonic::service::RoutesBuilder;
35use tonic::service::interceptor::InterceptedService;
36use tonic::transport::{Identity, ServerTlsConfig};
37use tower::{Layer, ServiceBuilder};
38
39use crate::grpc::database::DatabaseService;
40use crate::grpc::flight::{FlightCraftRef, FlightCraftWrapper};
41use crate::grpc::frontend_grpc_handler::FrontendGrpcHandler;
42use crate::grpc::greptime_handler::GreptimeRequestHandler;
43use crate::grpc::memory_limit::{MemoryLimiterExtensionLayer, MemoryLimiterExtensionService};
44use crate::grpc::prom_query_gateway::PrometheusGatewayService;
45use crate::grpc::region_server::{RegionServerHandlerRef, RegionServerRequestHandler};
46use crate::grpc::{GrpcServer, GrpcServerConfig};
47use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler};
48use crate::prometheus_handler::PrometheusHandlerRef;
49use crate::query_handler::OpenTelemetryProtocolHandlerRef;
50use crate::request_memory_limiter::ServerMemoryLimiter;
51use crate::tls::TlsOption;
52
53/// Add a gRPC service (`service`) to a `builder`([RoutesBuilder]).
54/// This macro will automatically add some gRPC properties to the service.
55#[macro_export]
56macro_rules! add_service {
57    ($builder: ident, $service: expr) => {
58        let max_recv_message_size = $builder.config().max_recv_message_size;
59        let max_send_message_size = $builder.config().max_send_message_size;
60
61        use tonic::codec::CompressionEncoding;
62        let service_builder = $service
63            .max_decoding_message_size(max_recv_message_size)
64            .max_encoding_message_size(max_send_message_size)
65            .accept_compressed(CompressionEncoding::Gzip)
66            .accept_compressed(CompressionEncoding::Zstd)
67            .send_compressed(CompressionEncoding::Gzip)
68            .send_compressed(CompressionEncoding::Zstd);
69
70        // Apply memory limiter layer
71        use $crate::grpc::memory_limit::MemoryLimiterExtensionLayer;
72        let service_with_limiter = $crate::tower::ServiceBuilder::new()
73            .layer(MemoryLimiterExtensionLayer::new(
74                $builder.memory_limiter().clone(),
75                max_recv_message_size,
76            ))
77            .service(service_builder);
78
79        $builder
80            .routes_builder_mut()
81            .add_service(service_with_limiter);
82    };
83}
84
85pub struct GrpcServerBuilder {
86    name: Option<String>,
87    config: GrpcServerConfig,
88    runtime: Runtime,
89    routes_builder: RoutesBuilder,
90    tls_config: Option<ServerTlsConfig>,
91    otel_arrow_service: Option<
92        MemoryLimiterExtensionService<
93            InterceptedService<
94                ArrowMetricsServiceServer<OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>>,
95                HeaderInterceptor,
96            >,
97        >,
98    >,
99    memory_limiter: ServerMemoryLimiter,
100    cors_allowed_origins: Option<Vec<String>>,
101}
102
103impl GrpcServerBuilder {
104    pub fn new(config: GrpcServerConfig, runtime: Runtime) -> Self {
105        // Create a default unlimited limiter (can be overridden with with_memory_limiter)
106        let memory_limiter = ServerMemoryLimiter::default();
107
108        Self {
109            name: None,
110            config,
111            runtime,
112            routes_builder: RoutesBuilder::default(),
113            tls_config: None,
114            otel_arrow_service: None,
115            memory_limiter,
116            cors_allowed_origins: None,
117        }
118    }
119
120    /// Enables CORS for gRPC-Web clients in browsers. An empty list allows any origin.
121    pub fn with_cors(mut self, allowed_origins: Vec<String>) -> Self {
122        self.cors_allowed_origins = Some(allowed_origins);
123        self
124    }
125
126    /// Set a global memory limiter for all server protocols.
127    pub fn with_memory_limiter(mut self, limiter: ServerMemoryLimiter) -> Self {
128        self.memory_limiter = limiter;
129        self
130    }
131
132    pub fn config(&self) -> &GrpcServerConfig {
133        &self.config
134    }
135
136    pub fn runtime(&self) -> &Runtime {
137        &self.runtime
138    }
139
140    pub fn memory_limiter(&self) -> &ServerMemoryLimiter {
141        &self.memory_limiter
142    }
143
144    pub fn name(self, name: Option<String>) -> Self {
145        Self { name, ..self }
146    }
147
148    /// Add handler for [DatabaseService] service.
149    pub fn database_handler(mut self, database_handler: GreptimeRequestHandler) -> Self {
150        add_service!(
151            self,
152            GreptimeDatabaseServer::new(DatabaseService::new(database_handler))
153        );
154        self
155    }
156
157    /// Add handler for Prometheus-compatible PromQL queries ([PrometheusGateway]).
158    pub fn prometheus_handler(
159        mut self,
160        prometheus_handler: PrometheusHandlerRef,
161        user_provider: Option<UserProviderRef>,
162    ) -> Self {
163        add_service!(
164            self,
165            PrometheusGatewayServer::new(PrometheusGatewayService::new(
166                prometheus_handler,
167                user_provider,
168            ))
169        );
170        self
171    }
172
173    /// Add handler for [FlightService](arrow_flight::flight_service_server::FlightService).
174    pub fn flight_handler(mut self, flight_handler: FlightCraftRef) -> Self {
175        add_service!(
176            self,
177            FlightServiceServer::new(FlightCraftWrapper(flight_handler.clone()))
178        );
179        self
180    }
181
182    /// Add handler for Frontend gRPC service.
183    pub fn frontend_grpc_handler(mut self, handler: FrontendGrpcHandler) -> Self {
184        add_service!(self, FrontendServer::new(handler));
185        self
186    }
187
188    /// Add handler for [OtelArrowService].
189    pub fn otel_arrow_handler(
190        mut self,
191        handler: OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef>,
192    ) -> Self {
193        let mut server = ArrowMetricsServiceServer::new(handler);
194        server = server
195            .max_decoding_message_size(self.config.max_recv_message_size)
196            .max_encoding_message_size(self.config.max_send_message_size)
197            .accept_compressed(CompressionEncoding::Zstd)
198            .send_compressed(CompressionEncoding::Zstd);
199        let svc = InterceptedService::new(server, HeaderInterceptor {});
200        // Same pre-decode memory admission as `add_service!`.
201        let svc = ServiceBuilder::new()
202            .layer(MemoryLimiterExtensionLayer::new(
203                self.memory_limiter.clone(),
204                self.config.max_recv_message_size,
205            ))
206            .service(svc);
207        self.otel_arrow_service = Some(svc);
208        self
209    }
210
211    /// Add handler for [RegionServer].
212    pub fn region_server_handler(mut self, region_server_handler: RegionServerHandlerRef) -> Self {
213        let handler = RegionServerRequestHandler::new(region_server_handler, self.runtime.clone());
214        add_service!(self, RegionServer::new(handler));
215        self
216    }
217
218    pub fn routes_builder_mut(&mut self) -> &mut RoutesBuilder {
219        &mut self.routes_builder
220    }
221
222    pub fn with_tls_config(mut self, tls_option: TlsOption) -> Result<Self> {
223        // tonic does not support watching for tls config changes
224        // so we don't support it either for now
225        if tls_option.watch {
226            warn!("Certificates watch and reloading for gRPC is NOT supported at the moment");
227        }
228        self.tls_config = if tls_option.should_force_tls() {
229            let cert = std::fs::read_to_string(tls_option.cert_path)
230                .context(InvalidConfigFilePathSnafu)?;
231            let key =
232                std::fs::read_to_string(tls_option.key_path).context(InvalidConfigFilePathSnafu)?;
233            let identity = Identity::from_pem(cert, key);
234            Some(ServerTlsConfig::new().identity(identity))
235        } else {
236            None
237        };
238        Ok(self)
239    }
240
241    pub fn add_layer<L>(self, layer: L) -> Self
242    where
243        L: Layer<Route> + Clone + Send + Sync + 'static,
244        L::Service: Service<Request> + Clone + Send + Sync + 'static,
245        <L::Service as Service<Request>>::Response: IntoResponse + 'static,
246        <L::Service as Service<Request>>::Error: Into<Infallible> + 'static,
247        <L::Service as Service<Request>>::Future: Send + 'static,
248    {
249        let routes = self.routes_builder.routes();
250        let router = routes.into_axum_router();
251        let router = router.layer(layer);
252        Self {
253            routes_builder: RoutesBuilder::from(router),
254            ..self
255        }
256    }
257
258    pub fn build(self) -> GrpcServer {
259        GrpcServer {
260            routes: Mutex::new(Some(self.routes_builder.routes())),
261            shutdown_tx: Mutex::new(None),
262            serve_state: Mutex::new(None),
263            tls_config: self.tls_config,
264            otel_arrow_service: Mutex::new(self.otel_arrow_service),
265            bind_addr: None,
266            name: self.name,
267            config: self.config,
268            cors_allowed_origins: self.cors_allowed_origins,
269        }
270    }
271}
servers/grpc/builder.rs

servers/grpc/
builder.rs
