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#[cfg(test)]
43mod tests {
44    use super::*;
45
46    #[test]
47    fn test_get_cpus() {
48        assert!(get_cpus() > 0);
49    }
50
51    #[test]
52    fn test_get_sys_total_memory() {
53        assert!(get_sys_total_memory().unwrap() > ReadableSize::mb(0));
54    }
55}
common_config/utils.rs

common_config/
utils.rs
