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