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