frontend/
frontend.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::sync::Arc;
16
17use common_base::readable_size::ReadableSize;
18use common_config::config::Configurable;
19use common_event_recorder::EventRecorderOptions;
20use common_memory_manager::OnExhaustedPolicy;
21use common_options::datanode::DatanodeClientOptions;
22use common_options::memory::MemoryOptions;
23use common_telemetry::logging::{LoggingOptions, SlowQueryOptions, TracingOptions};
24use meta_client::MetaClientOptions;
25use query::options::QueryOptions;
26use serde::{Deserialize, Serialize};
27use servers::grpc::GrpcOptions;
28use servers::http::HttpOptions;
29use servers::server::ServerHandlers;
30use snafu::ResultExt;
31
32use crate::error;
33use crate::error::Result;
34use crate::heartbeat::HeartbeatTask;
35use crate::instance::Instance;
36use crate::service_config::{
37    InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions,
38    PromStoreOptions,
39};
40
41#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
42#[serde(default)]
43pub struct FrontendOptions {
44    pub node_id: Option<String>,
45    pub default_timezone: Option<String>,
46    pub default_column_prefix: Option<String>,
47    /// Server-side global switch for auto table creation on write.
48    /// Acts as an upper bound: when `false`, missing tables are never auto-created
49    /// even if a request sets the `auto_create_table` hint to `true`. When `true`
50    /// (default), the per-request hint still applies. Default: `true`.
51    pub auto_create_table: bool,
52    /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
53    /// Set to 0 to disable the limit. Default: "0" (unlimited)
54    pub max_in_flight_write_bytes: ReadableSize,
55    /// Policy when write bytes quota is exhausted.
56    /// Options: "wait" (default, 10s), "wait(<duration>)", "fail"
57    pub write_bytes_exhausted_policy: OnExhaustedPolicy,
58    pub http: HttpOptions,
59    pub grpc: GrpcOptions,
60    /// The internal gRPC options for the frontend service.
61    /// it provide the same service as the public gRPC service, just only for internal use.
62    pub internal_grpc: Option<GrpcOptions>,
63    pub mysql: MysqlOptions,
64    pub postgres: PostgresOptions,
65    pub opentsdb: OpentsdbOptions,
66    pub influxdb: InfluxdbOptions,
67    pub prom_store: PromStoreOptions,
68    pub jaeger: JaegerOptions,
69    pub otlp: OtlpOptions,
70    pub meta_client: Option<MetaClientOptions>,
71    pub logging: LoggingOptions,
72    pub datanode: DatanodeClientOptions,
73    pub user_provider: Option<String>,
74    pub tracing: TracingOptions,
75    pub query: QueryOptions,
76    pub slow_query: SlowQueryOptions,
77    pub memory: MemoryOptions,
78    /// The event recorder options.
79    pub event_recorder: EventRecorderOptions,
80    /// Environment variable keys to read and report in heartbeat messages.
81    pub heartbeat_env_vars: Vec<String>,
82}
83
84impl Default for FrontendOptions {
85    fn default() -> Self {
86        Self {
87            node_id: None,
88            default_timezone: None,
89            default_column_prefix: None,
90            auto_create_table: true,
91            max_in_flight_write_bytes: ReadableSize(0),
92            write_bytes_exhausted_policy: OnExhaustedPolicy::default(),
93            http: HttpOptions::default(),
94            grpc: GrpcOptions::default(),
95            internal_grpc: None,
96            mysql: MysqlOptions::default(),
97            postgres: PostgresOptions::default(),
98            opentsdb: OpentsdbOptions::default(),
99            influxdb: InfluxdbOptions::default(),
100            jaeger: JaegerOptions::default(),
101            prom_store: PromStoreOptions::default(),
102            otlp: OtlpOptions::default(),
103            meta_client: None,
104            logging: LoggingOptions::default(),
105            datanode: DatanodeClientOptions::default(),
106            user_provider: None,
107            tracing: TracingOptions::default(),
108            query: QueryOptions::default(),
109            slow_query: SlowQueryOptions::default(),
110            memory: MemoryOptions::default(),
111            event_recorder: EventRecorderOptions::default(),
112            heartbeat_env_vars: vec![],
113        }
114    }
115}
116
117impl Configurable for FrontendOptions {
118    fn env_list_keys() -> Option<&'static [&'static str]> {
119        Some(&["heartbeat_env_vars", "meta_client.metasrv_addrs"])
120    }
121}
122
123/// The [`Frontend`] struct is the main entry point for the frontend service
124/// which contains server handlers, frontend instance and some background tasks.
125pub struct Frontend {
126    pub instance: Arc<Instance>,
127    pub servers: ServerHandlers,
128    pub heartbeat_task: Option<HeartbeatTask>,
129}
130
131impl Frontend {
132    pub async fn start(&mut self) -> Result<()> {
133        if let Some(t) = &self.heartbeat_task {
134            t.start().await?;
135        }
136
137        self.servers
138            .start_all()
139            .await
140            .context(error::StartServerSnafu)
141    }
142
143    pub async fn shutdown(&mut self) -> Result<()> {
144        self.servers
145            .shutdown_all()
146            .await
147            .context(error::ShutdownServerSnafu)
148    }
149
150    pub fn server_handlers(&self) -> &ServerHandlers {
151        &self.servers
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use std::sync::atomic::{AtomicBool, Ordering};
158    use std::time::Duration;
159
160    use api::v1::meta::heartbeat_server::HeartbeatServer;
161    use api::v1::meta::mailbox_message::Payload;
162    use api::v1::meta::{
163        AskLeaderRequest, AskLeaderResponse, HeartbeatRequest, HeartbeatResponse, MailboxMessage,
164        Peer, ResponseHeader, Role, heartbeat_server,
165    };
166    use async_trait::async_trait;
167    use client::{Client, Database};
168    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
169    use common_error::ext::ErrorExt;
170    use common_error::from_header_to_err_code_msg;
171    use common_error::status_code::StatusCode;
172    use common_grpc::channel_manager::ChannelManager;
173    use common_meta::heartbeat::handler::HandlerGroupExecutor;
174    use common_meta::heartbeat::handler::parse_mailbox_message::ParseMailboxMessageHandler;
175    use common_meta::heartbeat::handler::suspend::SuspendHandler;
176    use common_meta::instruction::Instruction;
177    use common_stat::ResourceStatImpl;
178    use meta_client::MetaClientRef;
179    use meta_client::client::MetaClientBuilder;
180    use meta_srv::service::GrpcStream;
181    use servers::grpc::{FlightCompression, GRPC_SERVER};
182    use servers::http::HTTP_SERVER;
183    use servers::http::result::greptime_result_v1::GreptimedbV1Response;
184    use tokio::sync::mpsc;
185    use tonic::codec::CompressionEncoding;
186    use tonic::codegen::tokio_stream::StreamExt;
187    use tonic::codegen::tokio_stream::wrappers::ReceiverStream;
188    use tonic::{Request, Response, Status, Streaming};
189
190    use super::*;
191    use crate::instance::builder::FrontendBuilder;
192    use crate::server::Services;
193
194    #[test]
195    fn test_toml() {
196        let opts = FrontendOptions::default();
197        let toml_string = toml::to_string(&opts).unwrap();
198        let _parsed: FrontendOptions = toml::from_str(&toml_string).unwrap();
199    }
200
201    struct SuspendableHeartbeatServer {
202        suspend: Arc<AtomicBool>,
203    }
204
205    #[async_trait]
206    impl heartbeat_server::Heartbeat for SuspendableHeartbeatServer {
207        type HeartbeatStream = GrpcStream<HeartbeatResponse>;
208
209        async fn heartbeat(
210            &self,
211            request: Request<Streaming<HeartbeatRequest>>,
212        ) -> std::result::Result<Response<Self::HeartbeatStream>, Status> {
213            let (tx, rx) = mpsc::channel(4);
214
215            common_runtime::spawn_global({
216                let mut requests = request.into_inner();
217                let suspend = self.suspend.clone();
218                async move {
219                    // Make the heartbeat interval short in unit tests to reduce the waiting time.
220                    // Only the handshake response needs to populate it (as metasrv does).
221                    let heartbeat_interval_ms = Duration::from_millis(200).as_millis() as u64;
222                    let mut is_handshake = true;
223                    while let Some(request) = requests.next().await {
224                        if let Err(e) = request {
225                            let _ = tx.send(Err(e)).await;
226                            return;
227                        }
228
229                        let mailbox_message =
230                            suspend.load(Ordering::Relaxed).then(|| MailboxMessage {
231                                payload: Some(Payload::Json(
232                                    serde_json::to_string(&Instruction::Suspend).unwrap(),
233                                )),
234                                ..Default::default()
235                            });
236                        let heartbeat_config =
237                            is_handshake.then_some(api::v1::meta::HeartbeatConfig {
238                                heartbeat_interval_ms,
239                                retry_interval_ms: heartbeat_interval_ms,
240                            });
241                        is_handshake = false;
242                        let response = HeartbeatResponse {
243                            header: Some(ResponseHeader::success()),
244                            mailbox_message,
245                            heartbeat_config,
246                            ..Default::default()
247                        };
248
249                        let _ = tx.send(Ok(response)).await;
250                    }
251                }
252            });
253
254            Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
255        }
256
257        async fn ask_leader(
258            &self,
259            _: Request<AskLeaderRequest>,
260        ) -> std::result::Result<Response<AskLeaderResponse>, Status> {
261            Ok(Response::new(AskLeaderResponse {
262                header: Some(ResponseHeader::success()),
263                leader: Some(Peer {
264                    addr: "localhost:0".to_string(),
265                    ..Default::default()
266                }),
267            }))
268        }
269    }
270
271    async fn create_meta_client(
272        options: &MetaClientOptions,
273        heartbeat_server: Arc<SuspendableHeartbeatServer>,
274    ) -> MetaClientRef {
275        let (client, server) = tokio::io::duplex(1024);
276
277        // create the heartbeat server:
278        common_runtime::spawn_global(async move {
279            let mut router = tonic::transport::Server::builder();
280            let router = router.add_service(
281                HeartbeatServer::from_arc(heartbeat_server)
282                    .accept_compressed(CompressionEncoding::Zstd)
283                    .send_compressed(CompressionEncoding::Zstd),
284            );
285            router
286                .serve_with_incoming(futures::stream::iter([Ok::<_, std::io::Error>(server)]))
287                .await
288        });
289
290        // Move client to an option so we can _move_ the inner value
291        // on the first attempt to connect. All other attempts will fail.
292        let mut client = Some(client);
293        let connector = tower::service_fn(move |_| {
294            let client = client.take();
295            async move {
296                if let Some(client) = client {
297                    Ok(hyper_util::rt::TokioIo::new(client))
298                } else {
299                    Err(std::io::Error::other("client already taken"))
300                }
301            }
302        });
303        let manager = ChannelManager::new();
304        manager
305            .reset_with_connector("localhost:0", connector)
306            .unwrap();
307
308        // create the heartbeat client:
309        let mut client = MetaClientBuilder::new(0, Role::Frontend)
310            .enable_heartbeat()
311            .heartbeat_channel_manager(manager)
312            .build();
313        client.start(&options.metasrv_addrs).await.unwrap();
314        Arc::new(client)
315    }
316
317    async fn create_frontend(
318        options: &FrontendOptions,
319        meta_client: MetaClientRef,
320    ) -> Result<Frontend> {
321        let instance = Arc::new(
322            FrontendBuilder::new_test(options, meta_client.clone())
323                .try_build()
324                .await?,
325        );
326
327        let servers =
328            Services::new(options.clone(), instance.clone(), Default::default()).build()?;
329
330        let executor = Arc::new(HandlerGroupExecutor::new(vec![
331            Arc::new(ParseMailboxMessageHandler),
332            Arc::new(SuspendHandler::new(instance.suspend_state())),
333        ]));
334        let heartbeat_task = Some(HeartbeatTask::new(
335            instance.frontend_peer_addr().to_string(),
336            options,
337            meta_client,
338            executor,
339            Arc::new(ResourceStatImpl::default()),
340        ));
341
342        let mut frontend = Frontend {
343            instance,
344            servers,
345            heartbeat_task,
346        };
347        frontend.start().await?;
348        Ok(frontend)
349    }
350
351    async fn verify_suspend_state_by_http(
352        frontend: &Frontend,
353        expected: std::result::Result<&str, (StatusCode, &str)>,
354    ) {
355        let addr = frontend.server_handlers().addr(HTTP_SERVER).unwrap();
356        let response = reqwest::get(format!("http://{}/v1/sql?sql=SELECT 1", addr))
357            .await
358            .unwrap();
359
360        let headers = response.headers();
361        let response = if let Some((code, error)) = from_header_to_err_code_msg(headers) {
362            Err((code, error))
363        } else {
364            Ok(response.text().await.unwrap())
365        };
366
367        match (response, expected) {
368            (Ok(response), Ok(expected)) => {
369                let response: GreptimedbV1Response = serde_json::from_str(&response).unwrap();
370                let response = serde_json::to_string(response.output()).unwrap();
371                assert_eq!(&response, expected);
372            }
373            (Err(actual), Err(expected)) => assert_eq!(actual, expected),
374            _ => unreachable!(),
375        }
376    }
377
378    async fn verify_suspend_state_by_grpc(
379        frontend: &Frontend,
380        expected: std::result::Result<&str, (StatusCode, &str)>,
381    ) {
382        let addr = frontend.server_handlers().addr(GRPC_SERVER).unwrap();
383        let client = Client::with_urls([addr.to_string()]);
384        let client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client);
385        let response = client.sql("SELECT 1").await;
386
387        match (response, expected) {
388            (Ok(response), Ok(expected)) => {
389                let response = response.data.pretty_print().await;
390                assert_eq!(&response, expected.trim());
391            }
392            (Err(actual), Err(expected)) => {
393                assert_eq!(actual.status_code(), expected.0);
394                assert_eq!(actual.output_msg(), expected.1);
395            }
396            _ => unreachable!(),
397        }
398    }
399
400    async fn wait_for_suspend_state(frontend: &Frontend, expected: bool) {
401        let check = || frontend.instance.is_suspended() == expected;
402        if check() {
403            return;
404        }
405
406        tokio::time::timeout(Duration::from_secs(5), async move {
407            while !check() {
408                tokio::time::sleep(Duration::from_millis(20)).await;
409            }
410        })
411        .await
412        .unwrap();
413    }
414
415    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
416    async fn test_suspend_frontend() -> Result<()> {
417        common_telemetry::init_default_ut_logging();
418
419        let meta_client_options = MetaClientOptions {
420            metasrv_addrs: vec!["localhost:0".to_string()],
421            ..Default::default()
422        };
423        let options = FrontendOptions {
424            http: HttpOptions {
425                addr: "127.0.0.1:0".to_string(),
426                ..Default::default()
427            },
428            grpc: GrpcOptions {
429                bind_addr: "127.0.0.1:0".to_string(),
430                flight_compression: FlightCompression::None,
431                ..Default::default()
432            },
433            mysql: MysqlOptions {
434                enable: false,
435                ..Default::default()
436            },
437            postgres: PostgresOptions {
438                enable: false,
439                ..Default::default()
440            },
441            meta_client: Some(meta_client_options.clone()),
442            ..Default::default()
443        };
444
445        let server = Arc::new(SuspendableHeartbeatServer {
446            suspend: Arc::new(AtomicBool::new(false)),
447        });
448        let meta_client = create_meta_client(&meta_client_options, server.clone()).await;
449        let frontend = create_frontend(&options, meta_client).await?;
450
451        // initial state: not suspend:
452        assert!(!frontend.instance.is_suspended());
453        verify_suspend_state_by_http(&frontend, Ok(r#"[{"records":{"schema":{"column_schemas":[{"name":"Int64(1)","data_type":"Int64"}]},"rows":[[1]],"total_rows":1}}]"#)).await;
454        verify_suspend_state_by_grpc(
455            &frontend,
456            Ok(r#"
457+----------+
458| Int64(1) |
459+----------+
460| 1        |
461+----------+"#),
462        )
463        .await;
464
465        // make heartbeat server returned "suspend" instruction,
466        server.suspend.store(true, Ordering::Relaxed);
467        wait_for_suspend_state(&frontend, true).await;
468        // ... then the frontend is suspended:
469        assert!(frontend.instance.is_suspended());
470        verify_suspend_state_by_http(
471            &frontend,
472            Err((
473                StatusCode::Suspended,
474                "error: Service suspended, execution_time_ms: 0",
475            )),
476        )
477        .await;
478        verify_suspend_state_by_grpc(&frontend, Err((StatusCode::Suspended, "Service suspended")))
479            .await;
480
481        // make heartbeat server NOT returned "suspend" instruction,
482        server.suspend.store(false, Ordering::Relaxed);
483        wait_for_suspend_state(&frontend, false).await;
484        // ... then frontend's suspend state is cleared:
485        assert!(!frontend.instance.is_suspended());
486        verify_suspend_state_by_http(&frontend, Ok(r#"[{"records":{"schema":{"column_schemas":[{"name":"Int64(1)","data_type":"Int64"}]},"rows":[[1]],"total_rows":1}}]"#)).await;
487        verify_suspend_state_by_grpc(
488            &frontend,
489            Ok(r#"
490+----------+
491| Int64(1) |
492+----------+
493| 1        |
494+----------+"#),
495        )
496        .await;
497        Ok(())
498    }
499}
frontend/frontend.rs

frontend/
frontend.rs
