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 with_metric_engine = effective_prom_store_options(&opts).with_metric_engine;
413 let opts = &opts.postgres;
414 let pg_addr = parse_addr(&opts.addr)?;
415
416 let tls_server_config = Arc::new(
417 ReloadableTlsServerConfig::try_new(opts.tls.clone()).context(StartServerSnafu)?,
418 );
419
420 maybe_watch_server_tls_config(tls_server_config.clone()).context(StartServerSnafu)?;
421
422 let pg_server = Box::new(
423 PostgresServer::new(
424 instance.clone(),
425 instance.clone(),
426 opts.tls.should_force_tls(),
427 tls_server_config,
428 opts.keep_alive.as_secs(),
429 common_runtime::global_runtime(),
430 user_provider.clone(),
431 Some(self.instance.process_manager().clone()),
432 )
433 .with_influxdb_handler(instance.clone())
434 .with_otlp_handler(instance.clone(), with_metric_engine)
435 .with_batching_enabled(postgres_batching),
436 ) as Box<dyn Server>;
437
438 handlers.insert((pg_server, pg_addr));
439 }
440
441 Ok(handlers)
442 }
443}
444
445fn effective_prom_store_options(opts: &FrontendOptions) -> PromStoreOptions {
447 let mut prom_store = opts.prom_store.clone();
448 let shared = opts.logical_batcher_options();
449 if shared.protocols.contains(&BatchingProtocol::Prom) {
450 prom_store.pending_rows_flush_interval = shared.pending_rows_flush_interval;
451 prom_store.max_batch_rows = shared.max_batch_rows;
452 prom_store.max_concurrent_flushes = shared.max_concurrent_flushes;
453 prom_store.worker_channel_capacity = shared.worker_channel_capacity;
454 prom_store.max_inflight_requests = shared.max_inflight_requests;
455 prom_store.flow_notification_queue_capacity = shared.flow_notification_queue_capacity;
456 } else {
457 prom_store.pending_rows_flush_interval = Duration::ZERO;
458 }
459 prom_store
460}
461
462pub fn effective_http_options(opts: &FrontendOptions) -> HttpOptions {
464 effective_http_options_with_sync(opts, pending_rows_batch_sync_enabled())
465}
466
467fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> HttpOptions {
468 let mut http = opts.http.clone();
469 let prom_store = effective_prom_store_options(opts);
470 let shared = opts.table_batcher_options();
471 let common_enabled = batch_sync
472 && shared.pending_rows_batching_enabled()
473 && shared.protocols.iter().any(|protocol| {
474 !matches!(
475 protocol,
476 BatchingProtocol::Mysql | BatchingProtocol::Postgres
477 ) && (*protocol != BatchingProtocol::Prom
478 || (prom_store.enable && !prom_store.with_metric_engine))
479 });
480 let common_interval = common_enabled.then_some(shared.pending_rows_flush_interval);
481 let prom_interval = (prom_store.pending_rows_batching_enabled() && batch_sync)
482 .then_some(prom_store.pending_rows_flush_interval);
483 let logical = opts.logical_batcher_options();
484 let otlp_interval = (batch_sync
485 && opts.otlp.enable
486 && opts.prom_store.with_metric_engine
487 && logical.protocols.contains(&BatchingProtocol::Otlp)
488 && logical.pending_rows_batching_enabled())
489 .then_some(logical.pending_rows_flush_interval);
490 let Some(flush_interval) = common_interval
491 .into_iter()
492 .chain(prom_interval)
493 .chain(otlp_interval)
494 .max()
495 else {
496 return http;
497 };
498 let fallback_timeout = flush_interval.saturating_add(Duration::from_secs(1));
499 if http.timeout.is_zero() || http.timeout > fallback_timeout {
500 return http;
501 }
502
503 let configured_timeout = http.timeout;
504 http.timeout = fallback_timeout;
505 warn!(
506 ?configured_timeout,
507 ?flush_interval,
508 ?fallback_timeout,
509 "HTTP request timeout is not longer than the pending-row timeout fallback; using the fallback"
510 );
511 http
512}
513
514fn parse_addr(addr: &str) -> Result<SocketAddr> {
515 addr.parse().context(error::ParseAddrSnafu { addr })
516}
517
518#[cfg(test)]
519mod tests {
520 use std::num::NonZeroUsize;
521 use std::sync::Arc;
522 use std::sync::atomic::{AtomicUsize, Ordering};
523 use std::time::Duration;
524
525 use api::v1::HealthCheckRequest;
526 use api::v1::health_check_client::HealthCheckClient;
527 use api::v1::meta::Role;
528 use arrow_flight::{FlightData, PutResult, Ticket};
529 use async_trait::async_trait;
530 use auth::{UserProviderRef, static_user_provider_from_option};
531 use client::{Client, Database};
532 use common_grpc::channel_manager::ChannelManager;
533 use meta_client::client::MetaClientBuilder;
534 use reqwest::header::{ACCESS_CONTROL_ALLOW_ORIGIN, ACCESS_CONTROL_REQUEST_METHOD, ORIGIN};
535 use servers::grpc::GRPC_SERVER;
536 use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream};
537 use tonic::{Code, Request, Response, Status, Streaming};
538
539 use crate::instance::builder::FrontendBuilder;
540 use crate::server::*;
541
542 #[tokio::test]
543 async fn test_logical_batcher_shared_without_prom_endpoint() {
544 use crate::service_config::BatcherOptions;
545 let mut options = FrontendOptions::default();
546 options.prom_store.enable = false;
547 options.pending_rows_batcher.logical_table = Some(BatcherOptions {
548 protocols: vec![BatchingProtocol::Otlp],
549 pending_rows_flush_interval: Duration::from_millis(5),
550 ..Default::default()
551 });
552 let meta_client = Arc::new(
553 MetaClientBuilder::new(0, Role::Frontend)
554 .enable_procedure()
555 .build(),
556 );
557 let instance = Arc::new(
558 FrontendBuilder::new_test(&options, meta_client)
559 .try_build()
560 .await
561 .unwrap(),
562 );
563 let mut disabled_options = options.clone();
564 disabled_options.otlp.enable = false;
565 instance.init_logical_batcher(&disabled_options);
566 assert!(instance.logical_batcher().is_none());
567
568 let services = Services::new(options.clone(), instance.clone(), Plugins::default());
569 let batcher = instance.logical_batcher().unwrap().clone();
570 instance.init_logical_batcher(&options);
571 assert!(Arc::ptr_eq(&batcher, instance.logical_batcher().unwrap()));
572 let weak = Arc::downgrade(&instance);
573 drop(services);
574 drop(instance);
575 assert!(
576 weak.upgrade().is_none(),
577 "batcher must not retain its schema owner"
578 );
579 }
580
581 #[test]
582 fn test_logical_batcher_http_timeout_and_prom_disable() {
583 use crate::service_config::pending_rows_batcher::BatcherOptions;
584 let mut opts = FrontendOptions::default();
585 opts.http.timeout = Duration::from_secs(1);
586 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2);
587 opts.pending_rows_batcher.logical_table = Some(BatcherOptions {
588 protocols: vec![BatchingProtocol::Otlp],
589 pending_rows_flush_interval: Duration::from_secs(5),
590 ..Default::default()
591 });
592 assert!(!effective_prom_store_options(&opts).pending_rows_batching_enabled());
593 opts.prom_store.enable = false;
596 assert_eq!(
597 effective_http_options_with_sync(&opts, true).timeout,
598 Duration::from_secs(6)
599 );
600 assert_eq!(
601 effective_http_options_with_sync(&opts, false).timeout,
602 opts.http.timeout
603 );
604 opts.prom_store.with_metric_engine = false;
605 assert_eq!(
606 effective_http_options_with_sync(&opts, false).timeout,
607 opts.http.timeout
608 );
609 opts.prom_store.with_metric_engine = true;
610 opts.pending_rows_batcher
611 .logical_table
612 .as_mut()
613 .unwrap()
614 .protocols
615 .clear();
616 assert_eq!(
617 effective_http_options_with_sync(&opts, true).timeout,
618 opts.http.timeout
619 );
620 }
621
622 #[test]
623 fn test_effective_prom_batching_controls() {
624 for (selected, shared_enabled, metric_engine, prom_enabled) in [
626 (true, true, true, true),
627 (false, true, true, true),
628 (true, false, true, true),
629 (true, true, false, true),
630 (true, true, true, false),
631 ] {
632 let mut opts = FrontendOptions::default();
633 opts.http.timeout = Duration::from_millis(1);
634 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2);
635 opts.prom_store.with_metric_engine = metric_engine;
636 opts.prom_store.enable = prom_enabled;
637 let shared = &mut opts.pending_rows_batcher.table;
638 shared.protocols = vec![if selected {
639 BatchingProtocol::Prom
640 } else {
641 BatchingProtocol::Influxdb
642 }];
643 shared.pending_rows_flush_interval = if shared_enabled {
644 Duration::from_secs(5)
645 } else {
646 Duration::ZERO
647 };
648 shared.max_batch_rows = 7;
649 shared.max_concurrent_flushes = 3;
650 shared.worker_channel_capacity = 11;
651 shared.max_inflight_requests = 13;
652 shared.flow_notification_queue_capacity = NonZeroUsize::new(17).unwrap();
653
654 let mut expected = opts.prom_store.clone();
655 if selected && shared_enabled {
656 expected.pending_rows_flush_interval = shared.pending_rows_flush_interval;
657 expected.max_batch_rows = shared.max_batch_rows;
658 expected.max_concurrent_flushes = shared.max_concurrent_flushes;
659 expected.worker_channel_capacity = shared.worker_channel_capacity;
660 expected.max_inflight_requests = shared.max_inflight_requests;
661 expected.flow_notification_queue_capacity = shared.flow_notification_queue_capacity;
662 }
663 let actual = effective_prom_store_options(&opts);
664 assert_eq!(actual, expected);
665 assert_eq!(
666 actual.pending_rows_batching_enabled(),
667 metric_engine && prom_enabled
668 );
669 }
670 }
671
672 #[test]
673 fn test_http_timeout_covers_synchronous_batchers() {
674 for (
676 protocols,
677 metric_engine,
678 batch_sync,
679 shared_secs,
680 legacy_secs,
681 timeout_secs,
682 expected_secs,
683 ) in [
684 (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 1),
685 (vec![BatchingProtocol::Mysql], false, false, 5, 0, 1, 1),
686 (vec![BatchingProtocol::Postgres], false, false, 5, 0, 1, 1),
687 (vec![BatchingProtocol::Mysql], false, true, 5, 0, 1, 1),
688 (vec![BatchingProtocol::Postgres], false, true, 5, 0, 1, 1),
689 (vec![BatchingProtocol::Prom], true, false, 5, 2, 1, 1),
690 (vec![BatchingProtocol::Prom], true, true, 5, 2, 1, 6),
691 (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 1),
692 (vec![BatchingProtocol::Influxdb], true, true, 5, 8, 1, 9),
693 (vec![BatchingProtocol::Influxdb], true, true, 8, 5, 1, 9),
694 (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 0, 0),
695 (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 10, 10),
696 (vec![BatchingProtocol::Prom], false, false, 0, 2, 1, 1),
697 (vec![], true, false, 5, 2, 1, 1),
698 (vec![], true, true, 5, 2, 1, 3),
699 ] {
700 let mut opts = FrontendOptions::default();
701 opts.http.timeout = Duration::from_secs(timeout_secs);
702 opts.prom_store.with_metric_engine = metric_engine;
703 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(legacy_secs);
704 opts.pending_rows_batcher.table.protocols = protocols;
705 opts.pending_rows_batcher.table.pending_rows_flush_interval =
706 Duration::from_secs(shared_secs);
707 assert_eq!(
708 effective_http_options_with_sync(&opts, batch_sync).timeout,
709 Duration::from_secs(expected_secs)
710 );
711 }
712 }
713
714 struct CountingFlightCraft {
715 inner: FlightCraftRef,
716 do_get_calls: AtomicUsize,
717 do_put_calls: AtomicUsize,
718 }
719
720 #[async_trait]
721 impl FlightCraft for CountingFlightCraft {
722 async fn do_get(
723 &self,
724 request: Request<Ticket>,
725 ) -> std::result::Result<Response<TonicStream<FlightData>>, Status> {
726 assert_eq!(
727 request.extensions().get::<Channel>(),
728 Some(&Channel::Internal)
729 );
730 self.do_get_calls.fetch_add(1, Ordering::SeqCst);
731 self.inner.do_get(request).await
732 }
733
734 async fn do_put(
735 &self,
736 request: Request<Streaming<FlightData>>,
737 ) -> std::result::Result<Response<TonicStream<PutResult>>, Status> {
738 assert_eq!(
739 request.extensions().get::<Channel>(),
740 Some(&Channel::Internal)
741 );
742 self.do_put_calls.fetch_add(1, Ordering::SeqCst);
743 self.inner.do_put(request).await
744 }
745 }
746
747 #[test]
748 fn test_effective_http_timeout_for_pending_rows() {
749 let cases = [
750 ("disabled timeout", 0, 5000, true, true, 0),
751 ("disabled prom store", 1000, 5000, false, true, 1000),
752 ("disabled metric engine", 1000, 5000, true, false, 1000),
753 ("disabled batching", 1000, 0, true, true, 1000),
754 ("timeout below flush interval", 4000, 5000, true, true, 6000),
755 (
756 "timeout equals flush interval",
757 5000,
758 5000,
759 true,
760 true,
761 6000,
762 ),
763 ("timeout below fallback", 5500, 5000, true, true, 6000),
764 ("timeout equals fallback", 6000, 5000, true, true, 6000),
765 ("timeout above fallback", 7000, 5000, true, true, 7000),
766 ];
767
768 for (name, timeout, flush_interval, enable, with_metric_engine, expected) in cases {
769 let mut opts = FrontendOptions::default();
770 opts.http.timeout = Duration::from_millis(timeout);
771 opts.prom_store.pending_rows_flush_interval = Duration::from_millis(flush_interval);
772 opts.prom_store.enable = enable;
773 opts.prom_store.with_metric_engine = with_metric_engine;
774
775 assert_eq!(
776 Duration::from_millis(expected),
777 effective_http_options_with_sync(&opts, true).timeout,
778 "{name}"
779 );
780 }
781 }
782
783 #[test]
784 fn test_effective_http_timeout_skips_fallback_in_async_batch_mode() {
785 let mut opts = FrontendOptions::default();
789 opts.http.timeout = Duration::from_millis(1000);
790 opts.prom_store.pending_rows_flush_interval = Duration::from_millis(5000);
791
792 assert_eq!(
793 Duration::from_millis(1000),
794 effective_http_options_with_sync(&opts, false).timeout,
795 );
796 assert_eq!(
797 Duration::from_millis(6000),
798 effective_http_options_with_sync(&opts, true).timeout,
799 );
800 }
801
802 #[test]
803 fn test_invalid_shared_batching_preserves_prom_store_options() {
804 type KnobMutator = fn(&mut FrontendOptions);
805 let cases: [KnobMutator; 5] = [
806 |opts| opts.pending_rows_batcher.table.max_concurrent_flushes = usize::MAX,
807 |opts| opts.pending_rows_batcher.table.worker_channel_capacity = usize::MAX,
808 |opts| opts.pending_rows_batcher.table.max_inflight_requests = usize::MAX,
809 |opts| {
810 opts.pending_rows_batcher
811 .table
812 .flow_notification_queue_capacity = NonZeroUsize::new(usize::MAX).unwrap()
813 },
814 |opts| opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::MAX,
815 ];
816 for invalidate in cases {
817 let mut opts = FrontendOptions::default();
818 opts.http.timeout = Duration::from_secs(1);
819 opts.prom_store.pending_rows_flush_interval = Duration::from_secs(5);
820 opts.pending_rows_batcher.table.protocols =
821 vec![BatchingProtocol::Prom, BatchingProtocol::Influxdb];
822 opts.pending_rows_batcher.table.pending_rows_flush_interval = Duration::from_secs(10);
823 invalidate(&mut opts);
824 assert!(
825 !opts
826 .pending_rows_batcher
827 .table
828 .pending_rows_batching_enabled()
829 );
830 assert_eq!(opts.prom_store, effective_prom_store_options(&opts));
831 assert_eq!(
832 Duration::from_secs(6),
833 effective_http_options_with_sync(&opts, true).timeout
834 );
835 opts.prom_store.pending_rows_flush_interval = Duration::ZERO;
836 assert_eq!(
837 Duration::from_secs(1),
838 effective_http_options_with_sync(&opts, true).timeout
839 );
840 }
841 }
842
843 #[test]
844 fn test_effective_http_timeout_skips_fallback_when_batcher_disabled() {
845 type KnobMutator = fn(&mut FrontendOptions);
849 let cases: [(&str, KnobMutator); 9] = [
850 ("oversized flush concurrency", |opts| {
851 opts.prom_store.max_concurrent_flushes = usize::MAX
852 }),
853 ("oversized worker channel", |opts| {
854 opts.prom_store.worker_channel_capacity = usize::MAX
855 }),
856 ("oversized inflight limit", |opts| {
857 opts.prom_store.max_inflight_requests = usize::MAX
858 }),
859 ("oversized flow queue", |opts| {
860 opts.prom_store.flow_notification_queue_capacity =
861 std::num::NonZeroUsize::new(usize::MAX).unwrap()
862 }),
863 ("unrepresentable deadline", |opts| {
864 opts.prom_store.pending_rows_flush_interval = Duration::MAX
865 }),
866 ("zero max_batch_rows", |opts| {
867 opts.prom_store.max_batch_rows = 0
868 }),
869 ("zero max_concurrent_flushes", |opts| {
870 opts.prom_store.max_concurrent_flushes = 0
871 }),
872 ("zero worker_channel_capacity", |opts| {
873 opts.prom_store.worker_channel_capacity = 0
874 }),
875 ("zero max_inflight_requests", |opts| {
876 opts.prom_store.max_inflight_requests = 0
877 }),
878 ];
879
880 for (name, disable_batcher) in cases {
881 let mut opts = FrontendOptions::default();
882 opts.http.timeout = Duration::from_millis(1000);
883 opts.prom_store.pending_rows_flush_interval = Duration::from_millis(5000);
884 disable_batcher(&mut opts);
885
886 assert_eq!(
887 Duration::from_millis(1000),
888 effective_http_options_with_sync(&opts, true).timeout,
889 "{name}"
890 );
891 }
892 }
893
894 #[tokio::test]
895 async fn test_database_sql_authentication_differs_between_public_and_internal_grpc() {
896 let options = FrontendOptions {
897 http: HttpOptions {
898 addr: "127.0.0.1:0".to_string(),
899 ..Default::default()
900 },
901 grpc: GrpcOptions::default().with_bind_addr("127.0.0.1:0"),
902 internal_grpc: Some(GrpcOptions::default().with_bind_addr("127.0.0.1:0")),
903 mysql: crate::service_config::MysqlOptions {
904 enable: false,
905 ..Default::default()
906 },
907 postgres: crate::service_config::PostgresOptions {
908 enable: false,
909 ..Default::default()
910 },
911 ..Default::default()
912 };
913 let meta_client = Arc::new(
914 MetaClientBuilder::new(0, Role::Frontend)
915 .enable_procedure()
916 .build(),
917 );
918 let instance = Arc::new(
919 FrontendBuilder::new_test(&options, meta_client)
920 .try_build()
921 .await
922 .unwrap(),
923 );
924 let plugins = Plugins::new();
925 let provider =
926 static_user_provider_from_option("static_user_provider:cmd:greptime=greptime").unwrap();
927 plugins.insert::<UserProviderRef>(Arc::new(provider));
928 let public_flight_handler = Arc::new(GreptimeRequestHandler::new(
929 instance.clone(),
930 plugins.get::<UserProviderRef>(),
931 None,
932 options.grpc.flight_compression,
933 )) as FlightCraftRef;
934 let internal_flight_handler = Arc::new(CountingFlightCraft {
935 inner: Arc::new(GreptimeRequestHandler::new(
936 instance.clone(),
937 None,
938 None,
939 options.grpc.flight_compression,
940 )),
941 do_get_calls: AtomicUsize::new(0),
942 do_put_calls: AtomicUsize::new(0),
943 });
944 let internal_flight_handler_ref = internal_flight_handler.clone() as FlightCraftRef;
945 let mut services = Services::new(options, instance, plugins)
946 .with_flight_handler(public_flight_handler)
947 .with_internal_flight_handler(internal_flight_handler_ref)
948 .build()
949 .unwrap();
950
951 services.start_all().await.unwrap();
952 let public_addr = services.addr(GRPC_SERVER).unwrap();
953 let internal_addr = services.addr("INTERNAL_GRPC_SERVER").unwrap();
954 let public_database = Database::new(
955 "greptime",
956 "public",
957 Client::with_query_and_control_managers(
958 ChannelManager::new(),
959 ChannelManager::new(),
960 [public_addr.to_string()],
961 ),
962 );
963 let internal_database = Database::new(
964 "greptime",
965 "public",
966 Client::with_query_and_control_managers(
967 ChannelManager::new(),
968 ChannelManager::new(),
969 [internal_addr.to_string()],
970 ),
971 );
972
973 let internal_result = internal_database.sql("SELECT 1").await;
974 let put_result = internal_database
975 .do_put(Box::pin(futures::stream::empty()))
976 .await;
977 let public_result = public_database.sql("SELECT 1").await;
978
979 services.shutdown_all().await.unwrap();
980
981 assert!(internal_result.is_ok());
982 assert!(put_result.is_ok());
983 assert_eq!(
984 1,
985 internal_flight_handler.do_get_calls.load(Ordering::SeqCst)
986 );
987 assert_eq!(
988 1,
989 internal_flight_handler.do_put_calls.load(Ordering::SeqCst)
990 );
991 assert_eq!(
992 Some(Code::Unauthenticated),
993 public_result
994 .as_ref()
995 .err()
996 .and_then(|err| err.tonic_code())
997 );
998 }
999
1000 #[tokio::test]
1001 async fn test_services_builder_health_check_is_reachable() {
1002 let options = FrontendOptions {
1004 http: HttpOptions {
1005 addr: "127.0.0.1:0".to_string(),
1006 ..Default::default()
1007 },
1008 grpc: GrpcOptions::default().with_bind_addr("127.0.0.1:0"),
1009 mysql: crate::service_config::MysqlOptions {
1010 enable: false,
1011 ..Default::default()
1012 },
1013 postgres: crate::service_config::PostgresOptions {
1014 enable: false,
1015 ..Default::default()
1016 },
1017 ..Default::default()
1018 };
1019 let meta_client = Arc::new(
1020 MetaClientBuilder::new(0, Role::Frontend)
1021 .enable_procedure()
1022 .build(),
1023 );
1024 let instance = Arc::new(
1025 FrontendBuilder::new_test(&options, meta_client)
1026 .try_build()
1027 .await
1028 .unwrap(),
1029 );
1030 let mut services = Services::new(options, instance, Default::default())
1031 .build()
1032 .unwrap();
1033
1034 services.start_all().await.unwrap();
1036 let addr = services.addr(GRPC_SERVER).unwrap();
1037 let health_check = HealthCheckClient::connect(format!("http://{addr}"))
1038 .await
1039 .unwrap()
1040 .health_check(HealthCheckRequest {})
1041 .await;
1042 services.shutdown_all().await.unwrap();
1043
1044 assert!(health_check.is_ok());
1046 }
1047
1048 #[tokio::test]
1049 async fn test_internal_grpc_server_never_serves_cors() {
1050 let options = FrontendOptions {
1051 http: HttpOptions {
1052 addr: "127.0.0.1:0".to_string(),
1053 ..Default::default()
1054 },
1055 grpc: GrpcOptions {
1056 enable_cors: true,
1057 ..GrpcOptions::default().with_bind_addr("127.0.0.1:0")
1058 },
1059 internal_grpc: Some(GrpcOptions {
1060 enable_cors: true,
1061 ..GrpcOptions::default().with_bind_addr("127.0.0.1:0")
1062 }),
1063 mysql: crate::service_config::MysqlOptions {
1064 enable: false,
1065 ..Default::default()
1066 },
1067 postgres: crate::service_config::PostgresOptions {
1068 enable: false,
1069 ..Default::default()
1070 },
1071 ..Default::default()
1072 };
1073 let meta_client = Arc::new(
1074 MetaClientBuilder::new(0, Role::Frontend)
1075 .enable_procedure()
1076 .build(),
1077 );
1078 let instance = Arc::new(
1079 FrontendBuilder::new_test(&options, meta_client)
1080 .try_build()
1081 .await
1082 .unwrap(),
1083 );
1084 let mut services = Services::new(options, instance, Default::default())
1085 .build()
1086 .unwrap();
1087
1088 services.start_all().await.unwrap();
1089 let public_addr = services.addr(GRPC_SERVER).unwrap();
1090 let internal_addr = services.addr("INTERNAL_GRPC_SERVER").unwrap();
1091 let public = send_cors_preflight(public_addr).await;
1092 let internal = send_cors_preflight(internal_addr).await;
1093 services.shutdown_all().await.unwrap();
1094
1095 assert!(public.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN));
1096 assert!(!internal.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN));
1097 }
1098
1099 async fn send_cors_preflight(addr: std::net::SocketAddr) -> reqwest::Response {
1100 reqwest::Client::new()
1101 .request(
1102 reqwest::Method::OPTIONS,
1103 format!("http://{addr}/greptime.v1.HealthCheck/Check"),
1104 )
1105 .header(ORIGIN, "https://example.com")
1106 .header(ACCESS_CONTROL_REQUEST_METHOD, "POST")
1107 .send()
1108 .await
1109 .unwrap()
1110 }
1111}