Skip to main content

frontend/
server.rs

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