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

frontend/
frontend.rs
