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