1use std::net::SocketAddr;
16use std::sync::Arc;
17use std::time::Duration;
18
19use auth::UserProviderRef;
20use axum::extract::{Request, State};
21use axum::middleware::Next;
22use axum::response::IntoResponse;
23use common_base::Plugins;
24use common_config::Configurable;
25use common_telemetry::{info, warn};
26use meta_client::MetaClientOptions;
27use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled};
28use servers::error::Error as ServerError;
29use servers::grpc::builder::GrpcServerBuilder;
30use servers::grpc::flight::FlightCraftRef;
31use servers::grpc::frontend_grpc_handler::FrontendGrpcHandler;
32use servers::grpc::greptime_handler::GreptimeRequestHandler;
33use servers::grpc::{GrpcOptions, GrpcServer};
34use servers::http::event::LogValidatorRef;
35use servers::http::result::error_result::ErrorResponse;
36use servers::http::utils::router::RouterConfigurator;
37use servers::http::{HttpOptions, HttpServer, HttpServerBuilder};
38use servers::interceptor::LogIngestInterceptorRef;
39use servers::metrics_handler::MetricsHandler;
40use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
41use servers::otel_arrow::OtelArrowServiceHandler;
42use servers::postgres::PostgresServer;
43use servers::request_memory_limiter::ServerMemoryLimiter;
44use servers::server::{Server, ServerHandlers};
45use servers::tls::{ReloadableTlsServerConfig, maybe_watch_server_tls_config};
46use session::context::Channel;
47use snafu::ResultExt;
48use tonic::Status;
49
50use crate::error::{self, Result, StartServerSnafu, TomlFormatSnafu};
51use crate::frontend::FrontendOptions;
52use crate::instance::Instance;
53use crate::service_config::PromStoreOptions;
54
55pub struct Services<T>
56where
57 T: Into<FrontendOptions> + Configurable + Clone,
58{
59 opts: T,
60 instance: Arc<Instance>,
61 grpc_server_builder: Option<GrpcServerBuilder>,
62 http_server_builder: Option<HttpServerBuilder>,
63 plugins: Plugins,
64 flight_handler: Option<FlightCraftRef>,
65 internal_flight_handler: Option<FlightCraftRef>,
66 pub server_memory_limiter: ServerMemoryLimiter,
67}
68
69impl<T> Services<T>
70where
71 T: Into<FrontendOptions> + Configurable + Clone,
72{
73 pub fn new(opts: T, instance: Arc<Instance>, plugins: Plugins) -> Self {
74 let feopts = opts.clone().into();
75 instance.init_logical_batcher(&feopts);
76 let server_memory_limiter = ServerMemoryLimiter::new(
78 feopts.max_in_flight_write_bytes.as_bytes(),
79 feopts.write_bytes_exhausted_policy,
80 );
81
82 Self {
83 opts,
84 instance,
85 grpc_server_builder: None,
86 http_server_builder: None,
87 plugins,
88 flight_handler: None,
89 internal_flight_handler: None,
90 server_memory_limiter,
91 }
92 }
93
94 pub fn grpc_server_builder(
95 &self,
96 opts: &GrpcOptions,
97 request_memory_limiter: ServerMemoryLimiter,
98 ) -> Result<GrpcServerBuilder> {
99 let builder = GrpcServerBuilder::new(opts.as_config(), common_runtime::global_runtime())
100 .with_memory_limiter(request_memory_limiter)
101 .with_tls_config(opts.tls.clone())
102 .context(error::InvalidTlsConfigSnafu)?;
103 Ok(builder)
104 }
105
106 pub fn http_server_builder(
107 &self,
108 opts: &FrontendOptions,
109 request_memory_limiter: ServerMemoryLimiter,
110 ) -> HttpServerBuilder {
111 let mut builder = HttpServerBuilder::new(effective_http_options(opts))
112 .with_batching_protocols(opts.table_batcher_options().protocols.clone())
113 .with_logical_batching_protocols(opts.logical_batcher_options().protocols)
114 .with_memory_limiter(request_memory_limiter)
115 .with_sql_handler(self.instance.clone());
116
117 let validator = self.plugins.get::<LogValidatorRef>();
118 let ingest_interceptor = self.plugins.get::<LogIngestInterceptorRef<ServerError>>();
119 builder =
120 builder.with_log_ingest_handler(self.instance.clone(), validator, ingest_interceptor);
121 builder = builder.with_logs_handler(self.instance.clone());
122
123 if let Some(user_provider) = self.plugins.get::<UserProviderRef>() {
124 builder = builder.with_user_provider(user_provider);
125 }
126
127 if opts.opentsdb.enable {
128 builder = builder.with_opentsdb_handler(self.instance.clone());
129 }
130
131 if opts.influxdb.enable {
132 builder = builder.with_influxdb_handler(self.instance.clone());
133 }
134
135 let prom_store = effective_prom_store_options(opts);
136 if prom_store.enable {
137 let pending_rows_batcher = opts
138 .logical_batcher_options()
139 .protocols
140 .contains(&BatchingProtocol::Prom)
141 .then(|| self.instance.logical_batcher().cloned())
142 .flatten();
143 builder = builder
144 .with_prom_handler(
145 self.instance.clone(),
146 Some(self.instance.clone()),
147 opts.prom_store.with_metric_engine,
148 opts.prom_store.prom_validation_mode,
149 pending_rows_batcher,
150 )
151 .with_prometheus_handler(self.instance.clone());
152 }
153
154 if opts.otlp.enable {
155 builder = builder
156 .with_otlp_handler(self.instance.clone(), opts.prom_store.with_metric_engine);
157 }
158
159 if opts.jaeger.enable {
160 builder = builder.with_jaeger_handler(self.instance.clone());
161 }
162
163 builder = builder.with_dashboard_handler(self.instance.clone());
164
165 if let Some(configurator) = self.plugins.get::<RouterConfigurator>() {
166 info!("Adding extra router from plugins");
167 builder = builder.with_extra_router(configurator.router());
168 }
169
170 builder.add_layer(axum::middleware::from_fn_with_state(
171 self.instance.clone(),
172 async move |State(state): State<Arc<Instance>>, request: Request, next: Next| {
173 if state.is_suspended() {
174 return ErrorResponse::from_error(servers::error::SuspendedSnafu.build())
175 .into_response();
176 }
177 next.run(request).await
178 },
179 ))
180 }
181
182 pub fn with_grpc_server_builder(self, builder: GrpcServerBuilder) -> Self {
183 Self {
184 grpc_server_builder: Some(builder),
185 ..self
186 }
187 }
188
189 pub fn with_http_server_builder(self, builder: HttpServerBuilder) -> Self {
190 Self {
191 http_server_builder: Some(builder),
192 ..self
193 }
194 }
195
196 pub fn with_flight_handler(self, flight_handler: FlightCraftRef) -> Self {
197 Self {
198 flight_handler: Some(flight_handler),
199 ..self
200 }
201 }
202
203 pub fn with_internal_flight_handler(self, flight_handler: FlightCraftRef) -> Self {
204 Self {
205 internal_flight_handler: Some(flight_handler),
206 ..self
207 }
208 }
209
210 fn build_grpc_server(
211 &mut self,
212 grpc: &GrpcOptions,
213 meta_client: &Option<MetaClientOptions>,
214 name: Option<String>,
215 external: bool,
216 request_memory_limiter: ServerMemoryLimiter,
217 ) -> Result<GrpcServer> {
218 let mut builder = if let Some(builder) = self.grpc_server_builder.take() {
219 builder
220 } else {
221 self.grpc_server_builder(grpc, request_memory_limiter)?
222 };
223 if external && grpc.enable_cors {
225 builder = builder.with_cors(grpc.cors_allowed_origins.clone());
226 }
227
228 let user_provider = if external {
229 self.plugins.get::<UserProviderRef>()
230 } else {
231 None
233 };
234
235 let runtime = if meta_client.is_none() {
237 Some(builder.runtime().clone())
238 } else {
239 None
240 };
241
242 let greptime_request_handler = GreptimeRequestHandler::new(
243 self.instance.clone(),
244 user_provider.clone(),
245 runtime,
246 grpc.flight_compression,
247 );
248
249 let default_flight_handler = Arc::new(greptime_request_handler.clone()) as FlightCraftRef;
250 let flight_handler = if external {
251 self.flight_handler
252 .clone()
253 .unwrap_or(default_flight_handler)
254 } else {
255 self.internal_flight_handler
256 .clone()
257 .unwrap_or(default_flight_handler)
258 };
259
260 let grpc_server = builder
261 .name(name)
262 .database_handler(greptime_request_handler.clone())
263 .prometheus_handler(self.instance.clone(), user_provider.clone())
264 .otel_arrow_handler(OtelArrowServiceHandler::new(
265 self.instance.clone(),
266 user_provider.clone(),
267 ))
268 .flight_handler(flight_handler)
269 .add_layer(axum::middleware::from_fn_with_state(
270 self.instance.clone(),
271 move |State(state): State<Arc<Instance>>, mut request: Request, next: Next| async move {
272 if state.is_suspended() {
273 let status = Status::from(servers::error::SuspendedSnafu.build());
274 return status.into_http();
275 }
276 if !external {
278 request.extensions_mut().insert(Channel::Internal);
279 }
280 next.run(request).await
281 },
282 ));
283
284 let grpc_server = if !external {
285 let frontend_grpc_handler =
286 FrontendGrpcHandler::new(self.instance.process_manager().clone());
287 grpc_server.frontend_grpc_handler(frontend_grpc_handler)
288 } else {
289 grpc_server
290 }
291 .build();
292
293 Ok(grpc_server)
294 }
295
296 fn build_http_server(
297 &mut self,
298 opts: &FrontendOptions,
299 toml: String,
300 request_memory_limiter: ServerMemoryLimiter,
301 ) -> Result<(HttpServer, Option<HttpServer>)> {
302 let builder = if let Some(builder) = self.http_server_builder.take() {
303 builder
304 } else {
305 self.http_server_builder(opts, request_memory_limiter)
306 };
307
308 let (internal, api) = builder
312 .with_metrics_handler(MetricsHandler)
313 .with_greptime_config_options(toml)
314 .build_servers();
315 Ok((internal, api))
316 }
317
318 pub fn build(mut self) -> Result<ServerHandlers> {
319 let opts = self.opts.clone();
320 let instance = self.instance.clone();
321
322 let toml = opts.to_toml().context(TomlFormatSnafu)?;
323 let opts: FrontendOptions = opts.into();
324
325 let handlers = ServerHandlers::default();
326
327 let user_provider = self.plugins.get::<UserProviderRef>();
328
329 {
330 let grpc_addr = parse_addr(&opts.grpc.bind_addr)?;
332 let grpc_server = self.build_grpc_server(
333 &opts.grpc,
334 &opts.meta_client,
335 None,
336 true,
337 self.server_memory_limiter.clone(),
338 )?;
339 handlers.insert((Box::new(grpc_server), grpc_addr));
340 }
341
342 if let Some(internal_grpc) = &opts.internal_grpc {
343 let grpc_addr = parse_addr(&internal_grpc.bind_addr)?;
345 let grpc_server = self.build_grpc_server(
346 internal_grpc,
347 &opts.meta_client,
348 Some("INTERNAL_GRPC_SERVER".to_string()),
349 false,
350 self.server_memory_limiter.clone(),
351 )?;
352 handlers.insert((Box::new(grpc_server), grpc_addr));
353 }
354
355 {
356 let http_options = &opts.http;
359 let http_addr = parse_addr(&http_options.addr)?;
360 let (http_server, http_api_server) =
361 self.build_http_server(&opts, toml, self.server_memory_limiter.clone())?;
362 handlers.insert((Box::new(http_server), http_addr));
363
364 if let Some(http_api_server) = http_api_server {
365 let http_api_addr = parse_addr(&http_options.api_server_addr)?;
366 info!("HTTP API server is enabled at {}", http_api_addr);
367 handlers.insert((Box::new(http_api_server), http_api_addr));
368 }
369 }
370
371 let table_batcher = opts.table_batcher_options();
372 let batching_enabled = table_batcher.pending_rows_batching_enabled();
373 let mysql_batching =
374 batching_enabled && table_batcher.protocols.contains(&BatchingProtocol::Mysql);
375 let postgres_batching = batching_enabled
376 && table_batcher
377 .protocols
378 .contains(&BatchingProtocol::Postgres);
379
380 if opts.mysql.enable {
381 let opts = &opts.mysql;
383 let mysql_addr = parse_addr(&opts.addr)?;
384
385 let tls_server_config = Arc::new(
386 ReloadableTlsServerConfig::try_new(opts.tls.clone()).context(StartServerSnafu)?,
387 );
388
389 maybe_watch_server_tls_config(tls_server_config.clone()).context(StartServerSnafu)?;
391
392 let mysql_server = MysqlServer::create_server(
393 common_runtime::global_runtime(),
394 Arc::new(MysqlSpawnRef::new(instance.clone(), user_provider.clone())),
395 Arc::new(
396 MysqlSpawnConfig::new(
397 opts.tls.should_force_tls(),
398 tls_server_config,
399 opts.keep_alive.as_secs(),
400 opts.reject_no_database.unwrap_or(false),
401 opts.prepared_stmt_cache_size,
402 )
403 .with_batching_enabled(mysql_batching),
404 ),
405 Some(instance.process_manager().clone()),
406 );
407 handlers.insert((mysql_server, mysql_addr));
408 }
409
410 if opts.postgres.enable {
411 let opts = &opts.postgres;
413 let pg_addr = parse_addr(&opts.addr)?;
414
415 let tls_server_config = Arc::new(
416 ReloadableTlsServerConfig::try_new(opts.tls.clone()).context(StartServerSnafu)?,
417 );
418
419 maybe_watch_server_tls_config(tls_server_config.clone()).context(StartServerSnafu)?;
420
421 let pg_server = Box::new(
422 PostgresServer::new(
423 instance.clone(),
424 instance.clone(),
425 opts.tls.should_force_tls(),
426 tls_server_config,
427 opts.keep_alive.as_secs(),
428 common_runtime::global_runtime(),
429 user_provider.clone(),
430 Some(self.instance.process_manager().clone()),
431 )
432 .with_influxdb_handler(instance.clone())
433 .with_batching_enabled(postgres_batching),
434 ) as Box<dyn Server>;
435
436 handlers.insert((pg_server, pg_addr));
437 }
438
439 Ok(handlers)
440 }
441}
442
443fn effective_prom_store_options(opts: &FrontendOptions) -> PromStoreOptions {
445 let mut prom_store = opts.prom_store.clone();
446 let shared = opts.logical_batcher_options();
447 if shared.protocols.contains(&BatchingProtocol::Prom) {
448 prom_store.pending_rows_flush_interval = shared.pending_rows_flush_interval;
449 prom_store.max_batch_rows = shared.max_batch_rows;
450 prom_store.max_concurrent_flushes = shared.max_concurrent_flushes;
451 prom_store.worker_channel_capacity = shared.worker_channel_capacity;
452 prom_store.max_inflight_requests = shared.max_inflight_requests;
453 prom_store.flow_notification_queue_capacity = shared.flow_notification_queue_capacity;
454 } else {
455 prom_store.pending_rows_flush_interval = Duration::ZERO;
456 }
457 prom_store
458}
459
460pub fn effective_http_options(opts: &FrontendOptions) -> HttpOptions {
462 effective_http_options_with_sync(opts, pending_rows_batch_sync_enabled())
463}
464
465fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> HttpOptions {
466 let mut http = opts.http.clone();
467 let prom_store = effective_prom_store_options(opts);
468 let shared = opts.table_batcher_options();
469 let common_enabled = batch_sync
470 && shared.pending_rows_batching_enabled()
471 && shared.protocols.iter().any(|protocol| {
472 !matches!(
473 protocol,
474 BatchingProtocol::Mysql | BatchingProtocol::Postgres
475 ) && (*protocol != BatchingProtocol::Prom
476 || (prom_store.enable && !prom_store.with_metric_engine))
477 });
478 let common_interval = common_enabled.then_some(shared.pending_rows_flush_interval);
479 let prom_interval = (prom_store.pending_rows_batching_enabled() && batch_sync)
480 .then_some(prom_store.pending_rows_flush_interval);
481 let logical = opts.logical_batcher_options();
482 let otlp_interval = (batch_sync
483 && opts.otlp.enable
484 && opts.prom_store.with_metric_engine
485 && logical.protocols.contains(&BatchingProtocol::Otlp)
486 && logical.pending_rows_batching_enabled())
487 .then_some(logical.pending_rows_flush_interval);
488 let Some(flush_interval) = common_interval
489 .into_iter()
490 .chain(prom_interval)
491 .chain(otlp_interval)
492 .max()
493 else {
494 return http;
495 };
496 let fallback_timeout = flush_interval.saturating_add(Duration::from_secs(1));
497 if http.timeout.is_zero() || http.timeout > fallback_timeout {
498 return http;
499 }
500
501 let configured_timeout = http.timeout;
502 http.timeout = fallback_timeout;
503 warn!(
504 ?configured_timeout,
505 ?flush_interval,
506 ?fallback_timeout,
507 "HTTP request timeout is not longer than the pending-row timeout fallback; using the fallback"
508 );
509 http
510}
511
512fn parse_addr(addr: &str) -> Result<SocketAddr> {
513 addr.parse().context(error::ParseAddrSnafu { addr })
514}
515
516#[cfg(test)]
517mod tests {
518 use std::num::NonZeroUsize;
519 use std::sync::Arc;
520 use std::sync::atomic::{AtomicUsize, Ordering};
521 use std::time::Duration;
522
523 use api::v1::HealthCheckRequest;
524 use api::v1::health_check_client::HealthCheckClient;
525 use api::v1::meta::Role;
526 use arrow_flight::{FlightData, PutResult, Ticket};
527 use async_trait::async_trait;
528 use auth::{UserProviderRef, static_user_provider_from_option};
529 use client::{Client, Database};
530 use common_grpc::channel_manager::ChannelManager;
531 use meta_client::client::MetaClientBuilder;
532 use reqwest::header::{ACCESS_CONTROL_ALLOW_ORIGIN, ACCESS_CONTROL_REQUEST_METHOD, ORIGIN};
533 use servers::grpc::GRPC_SERVER;
534 use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream};
535 use tonic::{Code, Request, Response, Status, Streaming};
536
537 use crate::instance::builder::FrontendBuilder;
538 use crate::server::*;
539
540 #[tokio::test]
541 async fn test_logical_batcher_shared_without_prom_endpoint() {
542 use crate::service_config::BatcherOptions;
543 let mut options = FrontendOptions::default();
544 options.prom_store.enable = false;
545 options.pending_rows_batcher.logical_table = Some(BatcherOptions {
546 protocols: vec![BatchingProtocol::Otlp],
547 pending_rows_flush_interval: Duration::from_millis(5),
548 ..Default::default()
549 });
550 let meta_client = Arc::new(
551 MetaClientBuilder::new(0, Role::Frontend)
552 .enable_procedure()
553 .build(),
554 );
555 let instance = Arc::new(
556 FrontendBuilder::new_test(&options, meta_client)
557 .try_build()
558 .await
559 .unwrap(),
560 );
561 let mut disabled_options = options.clone();
562 disabled_options.otlp.enable = false;
563 instance.init_logical_batcher(&disabled_options);
564 assert!(instance.logical_batcher().is_none());
565
566 let services = Services::new(options.clone(), instance.clone(), Plugins::default());
567 let batcher = instance.logical_batcher().unwrap().clone();
568 instance.init_logical_batcher(&options);
569 assert!(Arc::ptr_eq(&batcher, instance.logical_batcher().unwrap()));
570 let weak = Arc::downgrade(&instance);
571 drop(services);
572 drop(instance);
573 assert!(
574 weak.upgrade().is_none(),
575 "batcher must not retain its schema owner"
576 );
577 }
578
579 #[test]
580 fn test_logical_batcher_http_timeout_and_prom_disable() {
581 use crate::service_config::pending_rows_batcher::BatcherOptions;
582 let mut opts = FrontendOptions::default();
583 opts.http.timeout = Duration::from_secs(1);
584 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2);
585 opts.pending_rows_batcher.logical_table = Some(BatcherOptions {
586 protocols: vec![BatchingProtocol::Otlp],
587 pending_rows_flush_interval: Duration::from_secs(5),
588 ..Default::default()
589 });
590 assert!(!effective_prom_store_options(&opts).pending_rows_batching_enabled());
591 opts.prom_store.enable = false;
594 assert_eq!(
595 effective_http_options_with_sync(&opts, true).timeout,
596 Duration::from_secs(6)
597 );
598 assert_eq!(
599 effective_http_options_with_sync(&opts, false).timeout,
600 opts.http.timeout
601 );
602 opts.prom_store.with_metric_engine = false;
603 assert_eq!(
604 effective_http_options_with_sync(&opts, false).timeout,
605 opts.http.timeout
606 );
607 opts.prom_store.with_metric_engine = true;
608 opts.pending_rows_batcher
609 .logical_table
610 .as_mut()
611 .unwrap()
612 .protocols
613 .clear();
614 assert_eq!(
615 effective_http_options_with_sync(&opts, true).timeout,
616 opts.http.timeout
617 );
618 }
619
620 #[test]
621 fn test_effective_prom_batching_controls() {
622 for (selected, shared_enabled, metric_engine, prom_enabled) in [
624 (true, true, true, true),
625 (false, true, true, true),
626 (true, false, true, true),
627 (true, true, false, true),
628 (true, true, true, false),
629 ] {
630 let mut opts = FrontendOptions::default();
631 opts.http.timeout = Duration::from_millis(1);
632 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2);
633 opts.prom_store.with_metric_engine = metric_engine;
634 opts.prom_store.enable = prom_enabled;
635 let shared = &mut opts.pending_rows_batcher.table;
636 shared.protocols = vec![if selected {
637 BatchingProtocol::Prom
638 } else {
639 BatchingProtocol::Influxdb
640 }];
641 shared.pending_rows_flush_interval = if shared_enabled {
642 Duration::from_secs(5)
643 } else {
644 Duration::ZERO
645 };
646 shared.max_batch_rows = 7;
647 shared.max_concurrent_flushes = 3;
648 shared.worker_channel_capacity = 11;
649 shared.max_inflight_requests = 13;
650 shared.flow_notification_queue_capacity = NonZeroUsize::new(17).unwrap();
651
652 let mut expected = opts.prom_store.clone();
653 if selected && shared_enabled {
654 expected.pending_rows_flush_interval = shared.pending_rows_flush_interval;
655 expected.max_batch_rows = shared.max_batch_rows;
656 expected.max_concurrent_flushes = shared.max_concurrent_flushes;
657 expected.worker_channel_capacity = shared.worker_channel_capacity;
658 expected.max_inflight_requests = shared.max_inflight_requests;
659 expected.flow_notification_queue_capacity = shared.flow_notification_queue_capacity;
660 }
661 let actual = effective_prom_store_options(&opts);
662 assert_eq!(actual, expected);
663 assert_eq!(
664 actual.pending_rows_batching_enabled(),
665 metric_engine && prom_enabled
666 );
667 }
668 }
669
670 #[test]
671 fn test_http_timeout_covers_synchronous_batchers() {
672 for (
674 protocols,
675 metric_engine,
676 batch_sync,
677 shared_secs,
678 legacy_secs,
679 timeout_secs,
680 expected_secs,
681 ) in [
682 (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 1),
683 (vec![BatchingProtocol::Mysql], false, false, 5, 0, 1, 1),
684 (vec![BatchingProtocol::Postgres], false, false, 5, 0, 1, 1),
685 (vec![BatchingProtocol::Mysql], false, true, 5, 0, 1, 1),
686 (vec![BatchingProtocol::Postgres], false, true, 5, 0, 1, 1),
687 (vec![BatchingProtocol::Prom], true, false, 5, 2, 1, 1),
688 (vec![BatchingProtocol::Prom], true, true, 5, 2, 1, 6),
689 (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 1),
690 (vec![BatchingProtocol::Influxdb], true, true, 5, 8, 1, 9),
691 (vec![BatchingProtocol::Influxdb], true, true, 8, 5, 1, 9),
692 (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 0, 0),
693 (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 10, 10),
694 (vec![BatchingProtocol::Prom], false, false, 0, 2, 1, 1),
695 (vec![], true, false, 5, 2, 1, 1),
696 (vec![], true, true, 5, 2, 1, 3),
697 ] {
698 let mut opts = FrontendOptions::default();
699 opts.http.timeout = Duration::from_secs(timeout_secs);
700 opts.prom_store.with_metric_engine = metric_engine;
701 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(legacy_secs);
702 opts.pending_rows_batcher.table.protocols = protocols;
703 opts.pending_rows_batcher.table.pending_rows_flush_interval =
704 Duration::from_secs(shared_secs);
705 assert_eq!(
706 effective_http_options_with_sync(&opts, batch_sync).timeout,
707 Duration::from_secs(expected_secs)
708 );
709 }
710 }
711
712 struct CountingFlightCraft {
713 inner: FlightCraftRef,
714 do_get_calls: AtomicUsize,
715 do_put_calls: AtomicUsize,
716 }
717
718 #[async_trait]
719 impl FlightCraft for CountingFlightCraft {
720 async fn do_get(
721 &self,
722 request: Request<Ticket>,
723 ) -> std::result::Result<Response<TonicStream<FlightData>>, Status> {
724 assert_eq!(
725 request.extensions().get::<Channel>(),
726 Some(&Channel::Internal)
727 );
728 self.do_get_calls.fetch_add(1, Ordering::SeqCst);
729 self.inner.do_get(request).await
730 }
731
732 async fn do_put(
733 &self,
734 request: Request<Streaming<FlightData>>,
735 ) -> std::result::Result<Response<TonicStream<PutResult>>, Status> {
736 assert_eq!(
737 request.extensions().get::<Channel>(),
738 Some(&Channel::Internal)
739 );
740 self.do_put_calls.fetch_add(1, Ordering::SeqCst);
741 self.inner.do_put(request).await
742 }
743 }
744
745 #[test]
746 fn test_effective_http_timeout_for_pending_rows() {
747 let cases = [
748 ("disabled timeout", 0, 5000, true, true, 0),
749 ("disabled prom store", 1000, 5000, false, true, 1000),
750 ("disabled metric engine", 1000, 5000, true, false, 1000),
751 ("disabled batching", 1000, 0, true, true, 1000),
752 ("timeout below flush interval", 4000, 5000, true, true, 6000),
753 (
754 "timeout equals flush interval",
755 5000,
756 5000,
757 true,
758 true,
759 6000,
760 ),
761 ("timeout below fallback", 5500, 5000, true, true, 6000),
762 ("timeout equals fallback", 6000, 5000, true, true, 6000),
763 ("timeout above fallback", 7000, 5000, true, true, 7000),
764 ];
765
766 for (name, timeout, flush_interval, enable, with_metric_engine, expected) in cases {
767 let mut opts = FrontendOptions::default();
768 opts.http.timeout = Duration::from_millis(timeout);
769 opts.prom_store.pending_rows_flush_interval = Duration::from_millis(flush_interval);
770 opts.prom_store.enable = enable;
771 opts.prom_store.with_metric_engine = with_metric_engine;
772
773 assert_eq!(
774 Duration::from_millis(expected),
775 effective_http_options_with_sync(&opts, true).timeout,
776 "{name}"
777 );
778 }
779 }
780
781 #[test]
782 fn test_effective_http_timeout_skips_fallback_in_async_batch_mode() {
783 let mut opts = FrontendOptions::default();
787 opts.http.timeout = Duration::from_millis(1000);
788 opts.prom_store.pending_rows_flush_interval = Duration::from_millis(5000);
789
790 assert_eq!(
791 Duration::from_millis(1000),
792 effective_http_options_with_sync(&opts, false).timeout,
793 );
794 assert_eq!(
795 Duration::from_millis(6000),
796 effective_http_options_with_sync(&opts, true).timeout,
797 );
798 }
799
800 #[test]
801 fn test_invalid_shared_batching_preserves_prom_store_options() {
802 type KnobMutator = fn(&mut FrontendOptions);
803 let cases: [KnobMutator; 5] = [
804 |opts| opts.pending_rows_batcher.table.max_concurrent_flushes = usize::MAX,
805 |opts| opts.pending_rows_batcher.table.worker_channel_capacity = usize::MAX,
806 |opts| opts.pending_rows_batcher.table.max_inflight_requests = usize::MAX,
807 |opts| {
808 opts.pending_rows_batcher
809 .table
810 .flow_notification_queue_capacity = NonZeroUsize::new(usize::MAX).unwrap()
811 },
812 |opts| opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::MAX,
813 ];
814 for invalidate in cases {
815 let mut opts = FrontendOptions::default();
816 opts.http.timeout = Duration::from_secs(1);
817 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(5);
818 opts.pending_rows_batcher.table.protocols =
819 vec![BatchingProtocol::Prom, BatchingProtocol::Influxdb];
820 opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::from_secs(10);
821 invalidate(&mut opts);
822 assert!(
823 !opts
824 .pending_rows_batcher
825 .table
826 .pending_rows_batching_enabled()
827 );
828 assert_eq!(opts.prom_store, effective_prom_store_options(&opts));
829 assert_eq!(
830 Duration::from_secs(6),
831 effective_http_options_with_sync(&opts, true).timeout
832 );
833 opts.prom_store.pending_rows_flush_interval = Duration::ZERO;
834 assert_eq!(
835 Duration::from_secs(1),
836 effective_http_options_with_sync(&opts, true).timeout
837 );
838 }
839 }
840
841 #[test]
842 fn test_effective_http_timeout_skips_fallback_when_batcher_disabled() {
843 type KnobMutator = fn(&mut FrontendOptions);
847 let cases: [(&str, KnobMutator); 9] = [
848 ("oversized flush concurrency", |opts| {
849 opts.prom_store.max_concurrent_flushes = usize::MAX
850 }),
851 ("oversized worker channel", |opts| {
852 opts.prom_store.worker_channel_capacity = usize::MAX
853 }),
854 ("oversized inflight limit", |opts| {
855 opts.prom_store.max_inflight_requests = usize::MAX
856 }),
857 ("oversized flow queue", |opts| {
858 opts.prom_store.flow_notification_queue_capacity =
859 std::num::NonZeroUsize::new(usize::MAX).unwrap()
860 }),
861 ("unrepresentable deadline", |opts| {
862 opts.prom_store.pending_rows_flush_interval = Duration::MAX
863 }),
864 ("zero max_batch_rows", |opts| {
865 opts.prom_store.max_batch_rows = 0
866 }),
867 ("zero max_concurrent_flushes", |opts| {
868 opts.prom_store.max_concurrent_flushes = 0
869 }),
870 ("zero worker_channel_capacity", |opts| {
871 opts.prom_store.worker_channel_capacity = 0
872 }),
873 ("zero max_inflight_requests", |opts| {
874 opts.prom_store.max_inflight_requests = 0
875 }),
876 ];
877
878 for (name, disable_batcher) in cases {
879 let mut opts = FrontendOptions::default();
880 opts.http.timeout = Duration::from_millis(1000);
881 opts.prom_store.pending_rows_flush_interval = Duration::from_millis(5000);
882 disable_batcher(&mut opts);
883
884 assert_eq!(
885 Duration::from_millis(1000),
886 effective_http_options_with_sync(&opts, true).timeout,
887 "{name}"
888 );
889 }
890 }
891
892 #[tokio::test]
893 async fn test_database_sql_authentication_differs_between_public_and_internal_grpc() {
894 let options = FrontendOptions {
895 http: HttpOptions {
896 addr: "127.0.0.1:0".to_string(),
897 ..Default::default()
898 },
899 grpc: GrpcOptions::default().with_bind_addr("127.0.0.1:0"),
900 internal_grpc: Some(GrpcOptions::default().with_bind_addr("127.0.0.1:0")),
901 mysql: crate::service_config::MysqlOptions {
902 enable: false,
903 ..Default::default()
904 },
905 postgres: crate::service_config::PostgresOptions {
906 enable: false,
907 ..Default::default()
908 },
909 ..Default::default()
910 };
911 let meta_client = Arc::new(
912 MetaClientBuilder::new(0, Role::Frontend)
913 .enable_procedure()
914 .build(),
915 );
916 let instance = Arc::new(
917 FrontendBuilder::new_test(&options, meta_client)
918 .try_build()
919 .await
920 .unwrap(),
921 );
922 let plugins = Plugins::new();
923 let provider =
924 static_user_provider_from_option("static_user_provider:cmd:greptime=greptime").unwrap();
925 plugins.insert::<UserProviderRef>(Arc::new(provider));
926 let public_flight_handler = Arc::new(GreptimeRequestHandler::new(
927 instance.clone(),
928 plugins.get::<UserProviderRef>(),
929 None,
930 options.grpc.flight_compression,
931 )) as FlightCraftRef;
932 let internal_flight_handler = Arc::new(CountingFlightCraft {
933 inner: Arc::new(GreptimeRequestHandler::new(
934 instance.clone(),
935 None,
936 None,
937 options.grpc.flight_compression,
938 )),
939 do_get_calls: AtomicUsize::new(0),
940 do_put_calls: AtomicUsize::new(0),
941 });
942 let internal_flight_handler_ref = internal_flight_handler.clone() as FlightCraftRef;
943 let mut services = Services::new(options, instance, plugins)
944 .with_flight_handler(public_flight_handler)
945 .with_internal_flight_handler(internal_flight_handler_ref)
946 .build()
947 .unwrap();
948
949 services.start_all().await.unwrap();
950 let public_addr = services.addr(GRPC_SERVER).unwrap();
951 let internal_addr = services.addr("INTERNAL_GRPC_SERVER").unwrap();
952 let public_database = Database::new(
953 "greptime",
954 "public",
955 Client::with_query_and_control_managers(
956 ChannelManager::new(),
957 ChannelManager::new(),
958 [public_addr.to_string()],
959 ),
960 );
961 let internal_database = Database::new(
962 "greptime",
963 "public",
964 Client::with_query_and_control_managers(
965 ChannelManager::new(),
966 ChannelManager::new(),
967 [internal_addr.to_string()],
968 ),
969 );
970
971 let internal_result = internal_database.sql("SELECT 1").await;
972 let put_result = internal_database
973 .do_put(Box::pin(futures::stream::empty()))
974 .await;
975 let public_result = public_database.sql("SELECT 1").await;
976
977 services.shutdown_all().await.unwrap();
978
979 assert!(internal_result.is_ok());
980 assert!(put_result.is_ok());
981 assert_eq!(
982 1,
983 internal_flight_handler.do_get_calls.load(Ordering::SeqCst)
984 );
985 assert_eq!(
986 1,
987 internal_flight_handler.do_put_calls.load(Ordering::SeqCst)
988 );
989 assert_eq!(
990 Some(Code::Unauthenticated),
991 public_result
992 .as_ref()
993 .err()
994 .and_then(|err| err.tonic_code())
995 );
996 }
997
998 #[tokio::test]
999 async fn test_services_builder_health_check_is_reachable() {
1000 let options = FrontendOptions {
1002 http: HttpOptions {
1003 addr: "127.0.0.1:0".to_string(),
1004 ..Default::default()
1005 },
1006 grpc: GrpcOptions::default().with_bind_addr("127.0.0.1:0"),
1007 mysql: crate::service_config::MysqlOptions {
1008 enable: false,
1009 ..Default::default()
1010 },
1011 postgres: crate::service_config::PostgresOptions {
1012 enable: false,
1013 ..Default::default()
1014 },
1015 ..Default::default()
1016 };
1017 let meta_client = Arc::new(
1018 MetaClientBuilder::new(0, Role::Frontend)
1019 .enable_procedure()
1020 .build(),
1021 );
1022 let instance = Arc::new(
1023 FrontendBuilder::new_test(&options, meta_client)
1024 .try_build()
1025 .await
1026 .unwrap(),
1027 );
1028 let mut services = Services::new(options, instance, Default::default())
1029 .build()
1030 .unwrap();
1031
1032 services.start_all().await.unwrap();
1034 let addr = services.addr(GRPC_SERVER).unwrap();
1035 let health_check = HealthCheckClient::connect(format!("http://{addr}"))
1036 .await
1037 .unwrap()
1038 .health_check(HealthCheckRequest {})
1039 .await;
1040 services.shutdown_all().await.unwrap();
1041
1042 assert!(health_check.is_ok());
1044 }
1045
1046 #[tokio::test]
1047 async fn test_internal_grpc_server_never_serves_cors() {
1048 let options = FrontendOptions {
1049 http: HttpOptions {
1050 addr: "127.0.0.1:0".to_string(),
1051 ..Default::default()
1052 },
1053 grpc: GrpcOptions {
1054 enable_cors: true,
1055 ..GrpcOptions::default().with_bind_addr("127.0.0.1:0")
1056 },
1057 internal_grpc: Some(GrpcOptions {
1058 enable_cors: true,
1059 ..GrpcOptions::default().with_bind_addr("127.0.0.1:0")
1060 }),
1061 mysql: crate::service_config::MysqlOptions {
1062 enable: false,
1063 ..Default::default()
1064 },
1065 postgres: crate::service_config::PostgresOptions {
1066 enable: false,
1067 ..Default::default()
1068 },
1069 ..Default::default()
1070 };
1071 let meta_client = Arc::new(
1072 MetaClientBuilder::new(0, Role::Frontend)
1073 .enable_procedure()
1074 .build(),
1075 );
1076 let instance = Arc::new(
1077 FrontendBuilder::new_test(&options, meta_client)
1078 .try_build()
1079 .await
1080 .unwrap(),
1081 );
1082 let mut services = Services::new(options, instance, Default::default())
1083 .build()
1084 .unwrap();
1085
1086 services.start_all().await.unwrap();
1087 let public_addr = services.addr(GRPC_SERVER).unwrap();
1088 let internal_addr = services.addr("INTERNAL_GRPC_SERVER").unwrap();
1089 let public = send_cors_preflight(public_addr).await;
1090 let internal = send_cors_preflight(internal_addr).await;
1091 services.shutdown_all().await.unwrap();
1092
1093 assert!(public.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN));
1094 assert!(!internal.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN));
1095 }
1096
1097 async fn send_cors_preflight(addr: std::net::SocketAddr) -> reqwest::Response {
1098 reqwest::Client::new()
1099 .request(
1100 reqwest::Method::OPTIONS,
1101 format!("http://{addr}/greptime.v1.HealthCheck/Check"),
1102 )
1103 .header(ORIGIN, "https://example.com")
1104 .header(ACCESS_CONTROL_REQUEST_METHOD, "POST")
1105 .send()
1106 .await
1107 .unwrap()
1108 }
1109}