Skip to main content

cmd/
lib.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
15#![recursion_limit = "256"]
16
17use async_trait::async_trait;
18use common_error::ext::ErrorExt;
19use common_error::status_code::StatusCode;
20use common_mem_prof::activate_heap_profile;
21use common_stat::{get_total_cpu_millicores, get_total_memory_bytes};
22use common_telemetry::{error, info, warn};
23
24use crate::error::Result;
25
26pub mod cli;
27pub mod datanode;
28pub mod error;
29pub mod flownode;
30pub mod frontend;
31pub mod metasrv;
32pub mod options;
33pub mod standalone;
34pub mod user;
35
36lazy_static::lazy_static! {
37    static ref APP_VERSION: prometheus::IntGaugeVec =
38        prometheus::register_int_gauge_vec!("greptime_app_version", "app version", &["version", "short_version", "app"]).unwrap();
39
40    static ref CPU_LIMIT: prometheus::IntGaugeVec =
41        prometheus::register_int_gauge_vec!("greptime_cpu_limit_in_millicores", "cpu limit in millicores", &["app"]).unwrap();
42
43    static ref MEMORY_LIMIT: prometheus::IntGaugeVec =
44        prometheus::register_int_gauge_vec!("greptime_memory_limit_in_bytes", "memory limit in bytes", &["app"]).unwrap();
45}
46
47/// wait for the close signal, for unix platform it's SIGINT or SIGTERM
48#[cfg(unix)]
49async fn start_wait_for_close_signal() -> std::io::Result<()> {
50    use tokio::signal::unix::{SignalKind, signal};
51    let mut sigint = signal(SignalKind::interrupt())?;
52    let mut sigterm = signal(SignalKind::terminate())?;
53
54    tokio::select! {
55        _ = sigint.recv() => {
56            info!("Received SIGINT, shutting down");
57        }
58        _ = sigterm.recv() => {
59            info!("Received SIGTERM, shutting down");
60        }
61    }
62
63    Ok(())
64}
65
66/// wait for the close signal, for non-unix platform it's ctrl-c
67#[cfg(not(unix))]
68async fn start_wait_for_close_signal() -> std::io::Result<()> {
69    tokio::signal::ctrl_c().await
70}
71
72#[async_trait]
73pub trait App: Send {
74    fn name(&self) -> &str;
75
76    /// A hook for implementor to make something happened before actual startup. Defaults to no-op.
77    async fn pre_start(&mut self) -> Result<()> {
78        Ok(())
79    }
80
81    async fn start(&mut self) -> Result<()>;
82
83    /// Waits the quit signal by default.
84    fn wait_signal(&self) -> bool {
85        true
86    }
87
88    async fn stop(&mut self) -> Result<()>;
89
90    async fn run(&mut self) -> Result<()> {
91        info!("Starting app: {}", self.name());
92
93        self.pre_start().await?;
94
95        self.start().await?;
96
97        if self.wait_signal()
98            && let Err(e) = start_wait_for_close_signal().await
99        {
100            error!(e; "Failed to listen for close signal");
101            // It's unusual to fail to listen for close signal, maybe there's something unexpected in
102            // the underlying system. So we stop the app instead of running nonetheless to let people
103            // investigate the issue.
104        }
105
106        self.stop().await?;
107        info!("Goodbye!");
108        Ok(())
109    }
110}
111
112/// Log the versions of the application.
113///
114/// `version` should be the same as the output of cli "--version";
115/// and the `short_version` is the short version of the codes, often consist of git branch and commit.
116pub fn log_versions(version: &str, short_version: &str, app: &str) {
117    // Report app version as gauge.
118    APP_VERSION
119        .with_label_values(&[common_version::version(), short_version, app])
120        .inc();
121
122    info!("GreptimeDB version: {}", version);
123}
124
125pub fn create_resource_limit_metrics(app: &str) {
126    let cpu_limit = get_total_cpu_millicores();
127    if cpu_limit > 0 {
128        info!(
129            "GreptimeDB start with cpu limit in millicores: {}",
130            cpu_limit
131        );
132        CPU_LIMIT.with_label_values(&[app]).set(cpu_limit);
133    }
134
135    let memory_limit = get_total_memory_bytes();
136    if memory_limit > 0 {
137        info!(
138            "GreptimeDB start with memory limit in bytes: {}",
139            memory_limit
140        );
141        MEMORY_LIMIT.with_label_values(&[app]).set(memory_limit);
142    }
143}
144
145pub fn maybe_activate_heap_profile(memory_options: &common_options::memory::MemoryOptions) {
146    if memory_options.enable_heap_profiling {
147        match activate_heap_profile() {
148            Ok(()) => {
149                info!("Heap profile is active");
150            }
151            Err(err) => {
152                if err.status_code() == StatusCode::Unsupported {
153                    info!("Heap profile is not supported");
154                } else {
155                    warn!(err; "Failed to activate heap profile");
156                }
157            }
158        }
159    }
160}