kumo_server_common/
start.rs

1use crate::diagnostic_logging::LoggingConfig;
2use crate::nodeid::NodeId;
3use anyhow::Context;
4use chrono::{DateTime, Utc};
5use config::RegisterFunc;
6use kumo_machine_info::MachineInfo;
7use kumo_prometheus::declare_metric;
8use kumo_server_lifecycle::LifeCycle;
9use kumo_server_runtime::rt_spawn;
10use parking_lot::Mutex;
11use std::future::Future;
12use std::path::Path;
13use std::sync::LazyLock;
14
15pub static ONLINE_SINCE: LazyLock<DateTime<Utc>> = LazyLock::new(Utc::now);
16pub static MACHINE_INFO: LazyLock<Mutex<Option<MachineInfo>>> = LazyLock::new(|| Mutex::new(None));
17
18pub struct StartConfig<'a> {
19    pub logging: LoggingConfig<'a>,
20    pub lua_funcs: &'a [RegisterFunc],
21    pub policy: &'a Path,
22}
23
24impl StartConfig<'_> {
25    pub async fn run<INIT, FINI>(
26        self,
27        init_future: INIT,
28        shutdown_future: FINI,
29    ) -> anyhow::Result<()>
30    where
31        INIT: Future<Output = anyhow::Result<()>> + Send + 'static,
32        FINI: Future<Output = ()> + Send + 'static,
33    {
34        LazyLock::force(&ONLINE_SINCE);
35        self.logging.init()?;
36
37        start_cpu_usage_monitor();
38        crate::log_backlog::start_monitor();
39
40        rustls::crypto::aws_lc_rs::default_provider()
41            .install_default()
42            .map_err(|_| anyhow::anyhow!("failed to install default crypto provider"))?;
43
44        kumo_server_memory::setup_memory_limit().context("setup_memory_limit")?;
45
46        prometheus::register(Box::new(
47            tokio_metrics_collector::default_runtime_collector(),
48        ))
49        .context("failed to configure tokio-metrics-collector")?;
50
51        for &func in self.lua_funcs {
52            config::register(func);
53        }
54
55        config::set_policy_path(self.policy.to_path_buf())
56            .await
57            .with_context(|| format!("Error evaluating policy file {:?}", self.policy))?;
58
59        tokio::spawn(async move {
60            let mut info = MachineInfo::new();
61            info.node_id.replace(NodeId::get().uuid.to_string());
62            info.query_cloud_provider().await;
63            *MACHINE_INFO.lock() = Some(info);
64        });
65
66        let mut life_cycle = LifeCycle::new();
67
68        let init_handle = rt_spawn("initialize", async move {
69            let mut error = None;
70            if let Err(err) = init_future.await {
71                let err = format!("{err:#}");
72                tracing::error!("problem initializing: {err}");
73                LifeCycle::request_shutdown(Ok(())).await;
74                error.replace(err);
75            }
76            // This log line is depended upon by the integration
77            // test harness. Do not change or remove it without
78            // making appropriate adjustments over there!
79            tracing::info!("initialization complete");
80            error
81        })?;
82
83        let final_result = life_cycle.wait_for_shutdown().await;
84
85        // after waiting for those to idle out, shut down logging
86        shutdown_future.await;
87
88        tracing::info!("Shutdown completed OK!");
89
90        if let Some(error) = init_handle.await? {
91            anyhow::bail!("Initialization raised an error: {error}");
92        }
93        if let Some(Err(error)) = final_result {
94            anyhow::bail!("Shutdown due to error: {error}");
95        }
96        Ok(())
97    }
98}
99
100declare_metric! {
101/// The sum of the system-wide CPU usage for each CPU in the system, can add up to more than 100%.
102///
103/// Each CPU has a value from 0-100% busy; a value of 100% in this metric
104/// indicates that the load is equivalent to one fully utilized CPU.
105///
106/// A multi-CPU system can report more than 100% in this metric; a dual-CPU
107/// system reporting 200% indicates that both CPUs are fully utilized.
108///
109/// See system_cpu_usage_normalized for a version of this metric that scales from
110/// 0% (totally idle) to 100% (totally saturated).
111///
112/// This metric is scoped to the system, reflecting the total load on the
113/// system, not just from the kumo related process(es).
114static SYS_CPU_USAGE: IntGauge("system_cpu_usage_sum");
115}
116
117declare_metric! {
118/// The sum of the system-wide CPU usage for each CPU in the system, divided by the number of CPUs.
119///
120/// 100% in this metric indicates that all CPU cores are 100% busy.
121///
122/// This metric is scoped to the system, reflecting the total load on the
123/// system, not just from the kumo related process(es).
124static SYS_CPU_USAGE_NORM: IntGauge("system_cpu_usage_normalized");
125}
126
127declare_metric! {
128/// The sum of the process CPU usage for each CPU in the system, can add up to more than 100%.
129///
130/// Each CPU has a value from 0-100% busy; a value of 100% in this metric
131/// indicates that the load is equivalent to one fully utilized CPU.
132///
133/// A multi-CPU system can report more than 100% in this metric; a dual-CPU
134/// system reporting 200% indicates that both CPUs are fully utilized.
135///
136/// See process_cpu_usage_normalized for a version of this metric that scales from
137/// 0% (totally idle) to 100% (totally saturated).
138///
139/// This metric is scoped to the service process, reflecting the CPU used only
140/// by the process and not the system as a whole.
141static PROC_CPU_USAGE: IntGauge("process_cpu_usage_sum");
142}
143
144declare_metric! {
145/// The sum of the process CPU usage for each CPU in the system, divided by the number of CPUs.
146///
147/// 100% in this metric indicates that all CPU cores are 100% busy.
148///
149/// This metric is scoped to the service process, reflecting the CPU used only
150/// by the process and not the system as a whole.
151static PROC_CPU_USAGE_NORM: IntGauge("process_cpu_usage_normalized");
152}
153
154fn start_cpu_usage_monitor() {
155    std::thread::spawn(|| {
156        use std::time::Duration;
157        use sysinfo::{
158            get_current_pid, CpuRefreshKind, ProcessRefreshKind, ProcessesToUpdate, RefreshKind,
159            System,
160        };
161
162        let mut sys = System::new_with_specifics(
163            RefreshKind::nothing().with_cpu(CpuRefreshKind::everything()),
164        );
165
166        let my_pid = get_current_pid().expect("failed to get own pid!?");
167
168        let update_interval = Duration::from_secs(3).max(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
169        // Do two initial updates so that we have sufficient data for
170        // computing a delta; one here outside the loop, and another
171        // at the top of the loop.
172        sys.refresh_cpu_usage();
173        std::thread::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
174
175        loop {
176            sys.refresh_cpu_usage();
177            sys.refresh_processes_specifics(
178                ProcessesToUpdate::Some(&[my_pid]),
179                true,
180                ProcessRefreshKind::everything().without_tasks(),
181            );
182
183            let num_cpus = sys.cpus().len() as i64;
184
185            let sys_usage = sys.cpus().iter().map(|cpu| cpu.cpu_usage()).sum::<f32>() as i64;
186            SYS_CPU_USAGE.set(sys_usage);
187            SYS_CPU_USAGE_NORM.set(sys_usage / num_cpus);
188
189            if let Some(p) = sys.process(my_pid) {
190                let proc_usage = p.cpu_usage() as i64;
191                PROC_CPU_USAGE.set(proc_usage);
192                PROC_CPU_USAGE_NORM.set(proc_usage / num_cpus);
193            }
194
195            std::thread::sleep(update_interval);
196        }
197    });
198}