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                    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
441/// Selected shared controls override legacy Prom batching knobs, not protocol behavior.
442fn 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
458/// Resolves HTTP options, allowing synchronous batches enough time to flush.
459pub 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        // Both logical protocols follow the global acknowledgement policy,
590        // independently of enabling the Prom HTTP endpoint.
591        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        // Only an enabled shared Prom selection replaces the legacy controls.
621        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        // Only synchronous batchers extend the HTTP timeout.
671        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        // With `PENDING_ROWS_BATCH_SYNC=false`, submissions return right after
782        // enqueue and no request waits for a pending-row flush, so the
783        // timeout must not be raised.
784        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        // Mirrors the conditions under which `LogicalTablePendingRowsBatcher::try_new`
842        // returns `None`; in these cases no request can wait for a pending-row
843        // flush, so the timeout must not be raised.
844        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        // Arrange
999        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        // Act
1031        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
1041        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}