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