common_config/
utils.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 common_base::readable_size::ReadableSize;
16use sysinfo::System;
17
18/// Get the CPU core number of system, aware of cgroups.
19pub fn get_cpus() -> usize {
20    // This function will check cgroups
21    num_cpus::get()
22}
23
24/// Get the total memory of the system.
25/// If `cgroup_limits` is enabled, it will also check it.
26pub fn get_sys_total_memory() -> Option<ReadableSize> {
27    if sysinfo::IS_SUPPORTED_SYSTEM {
28        let mut sys_info = System::new();
29        sys_info.refresh_memory();
30        let mut total_memory = sys_info.total_memory();
31        // Compare with cgroups memory limit, use smaller values
32        // This method is only implemented for Linux. It always returns None for all other systems.
33        if let Some(cgroup_limits) = sys_info.cgroup_limits() {
34            total_memory = total_memory.min(cgroup_limits.total_memory)
35        }
36        Some(ReadableSize(total_memory))
37    } else {
38        None
39    }
40}
41
42/// `ResourceSpec` holds the static resource specifications of a node,
43/// such as CPU cores and memory capacity. These values are fixed
44/// at startup and do not change dynamically during runtime.
45#[derive(Debug, Clone, Copy)]
46pub struct ResourceSpec {
47    pub cpus: usize,
48    pub memory: Option<ReadableSize>,
49}
50
51impl Default for ResourceSpec {
52    fn default() -> Self {
53        Self {
54            cpus: get_cpus(),
55            memory: get_sys_total_memory(),
56        }
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn test_get_cpus() {
66        assert!(get_cpus() > 0);
67    }
68
69    #[test]
70    fn test_get_sys_total_memory() {
71        assert!(get_sys_total_memory().unwrap() > ReadableSize::mb(0));
72    }
73}
common_config/utils.rs

common_config/
utils.rs
