rw_resource_util/
lib.rs

1// Copyright 2025 RisingWave Labs
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/// Returns the hostname of the current machine, or an empty string if failed,
16/// which rarely happens.
17pub fn hostname() -> String {
18    hostname::get()
19        .unwrap_or_default()
20        .into_string()
21        .unwrap_or_default()
22}
23
24#[derive(Debug, Clone, Copy)]
25pub enum CgroupVersion {
26    V1,
27    V2,
28}
29
30/// Current controllers available in implementation.
31pub enum Controller {
32    Cpu,
33    Memory,
34}
35
36/// Default constant Cgroup paths and hierarchy.
37const DEFAULT_CGROUP_ROOT_HIERARCYHY: &str = "/sys/fs/cgroup";
38const DEFAULT_CGROUP_V2_CONTROLLER_LIST_PATH: &str = "/sys/fs/cgroup/cgroup.controllers";
39const DEFAULT_CGROUP_MAX_INDICATOR: &str = "max";
40
41mod runtime {
42    use std::env;
43    use std::path::Path;
44
45    use thiserror_ext::AsReport;
46
47    use super::CgroupVersion;
48    use super::util::parse_controller_enable_file_for_cgroup_v2;
49    const DEFAULT_DOCKER_ENV_PATH: &str = "/.dockerenv";
50    const DEFAULT_LINUX_IDENTIFIER: &str = "linux";
51    const DEFAULT_IN_CONTAINER_ENV_VARIABLE: &str = "IN_CONTAINER";
52    const DEFAULT_KUBERNETES_SECRETS_PATH: &str = "/var/run/secrets/kubernetes.io";
53
54    fn is_linux_machine() -> bool {
55        env::consts::OS.eq(DEFAULT_LINUX_IDENTIFIER)
56    }
57
58    /// checks if is running in a docker container by checking for docker env file, or if it is
59    /// running in a kubernetes pod.
60    fn is_running_in_container() -> bool {
61        return env_var_check_if_running_in_container()
62            || docker_env_exists()
63            || is_running_in_kubernetes_pod();
64
65        /// checks for existence of docker env file
66        fn docker_env_exists() -> bool {
67            Path::new(DEFAULT_DOCKER_ENV_PATH).exists()
68        }
69
70        /// checks for environment
71        fn env_var_check_if_running_in_container() -> bool {
72            env::var(DEFAULT_IN_CONTAINER_ENV_VARIABLE).is_ok()
73        }
74
75        /// checks if it is running in a kubernetes pod
76        fn is_running_in_kubernetes_pod() -> bool {
77            Path::new(DEFAULT_KUBERNETES_SECRETS_PATH).exists()
78        }
79    }
80
81    /// Given a certain controller, checks if it is enabled.
82    /// For cgroup v1, existence of directory with controller name is checked in cgroup default root
83    /// hierarchy. e.g if directory "/sys/fs/cgroup"/cpu" exists then CPU controller is enabled.
84    /// For cgroup v2, check the controller list path for the controller name.
85    pub fn is_controller_activated(
86        controller_type: super::Controller,
87        cgroup_version: CgroupVersion,
88    ) -> bool {
89        let controller_name: &str = match controller_type {
90            super::Controller::Cpu => "cpu",
91            super::Controller::Memory => "memory",
92        };
93        match cgroup_version {
94            super::CgroupVersion::V1 => Path::new(super::DEFAULT_CGROUP_ROOT_HIERARCYHY)
95                .join(controller_name)
96                .is_dir(),
97            super::CgroupVersion::V2 => parse_controller_enable_file_for_cgroup_v2(
98                super::DEFAULT_CGROUP_V2_CONTROLLER_LIST_PATH,
99                controller_name,
100            ),
101        }
102    }
103
104    /// If cgroup exists or is enabled in kernel, returnb true, else false.
105    fn cgroup_exists() -> bool {
106        Path::new(super::DEFAULT_CGROUP_ROOT_HIERARCYHY).is_dir()
107    }
108
109    pub fn get_resource<T>(
110        desc: &str,
111        controller_type: super::Controller,
112        get_system: fn() -> T,
113        get_container: fn(CgroupVersion) -> Result<T, std::io::Error>,
114    ) -> T {
115        if !is_linux_machine() || !is_running_in_container() || !cgroup_exists() {
116            return get_system();
117        };
118
119        // if cgroup.controllers exist, v2 is used.
120        let cgroup_version = if Path::new(super::DEFAULT_CGROUP_V2_CONTROLLER_LIST_PATH).exists() {
121            super::CgroupVersion::V2
122        } else {
123            super::CgroupVersion::V1
124        };
125        if !is_controller_activated(controller_type, cgroup_version) {
126            return get_system();
127        }
128
129        match get_container(cgroup_version) {
130            Ok(value) => value,
131            Err(err) => {
132                tracing::warn!(
133                    error = %err.as_report(),
134                    cgroup_version = ?cgroup_version,
135                    "failed to get {desc} in container, use system value instead"
136                );
137                get_system()
138            }
139        }
140    }
141}
142
143pub mod memory {
144    use sysinfo::System;
145
146    use super::runtime::get_resource;
147
148    /// Default paths for memory limitations and usage for cgroup v1 and cgroup v2.
149    const V1_MEMORY_LIMIT_PATH: &str = "/sys/fs/cgroup/memory/memory.limit_in_bytes";
150    const V1_MEMORY_CURRENT_PATH: &str = "/sys/fs/cgroup/memory/memory.usage_in_bytes";
151    const V2_MEMORY_LIMIT_PATH: &str = "/sys/fs/cgroup/memory.max";
152    const V2_MEMORY_CURRENT_PATH: &str = "/sys/fs/cgroup/memory.current";
153
154    /// Returns the system memory.
155    pub fn get_system_memory() -> usize {
156        let mut sys = System::new();
157        sys.refresh_memory();
158        sys.total_memory() as usize
159    }
160
161    /// Returns the used memory of the system.
162    pub fn get_system_memory_used() -> usize {
163        let mut sys = System::new();
164        sys.refresh_memory();
165        sys.used_memory() as usize
166    }
167
168    /// Returns the total memory used by the system in bytes.
169    ///
170    /// If running in container, this function will read the cgroup interface files for the
171    /// memory used, if interface files are not found, will return the memory used in the system
172    /// as default. The cgroup mount point is assumed to be at /sys/fs/cgroup by default.
173    ///
174    ///
175    /// # Examples
176    ///
177    /// Basic usage:
178    /// ``` ignore
179    /// let mem_used = memory::total_memory_used_bytes();
180    /// ```
181    pub fn total_memory_used_bytes() -> usize {
182        get_resource(
183            "memory used",
184            super::Controller::Memory,
185            get_system_memory_used,
186            get_container_memory_used,
187        )
188    }
189
190    /// Returns the total memory available by the system in bytes.
191    ///
192    /// If running in container, this function will read the cgroup interface files for the
193    /// memory available/limit, if interface files are not found, will return the system memory
194    /// volume by default. The cgroup mount point is assumed to be at /sys/fs/cgroup by default.
195    ///
196    ///
197    /// # Examples
198    ///
199    /// Basic usage:
200    /// ``` ignore
201    /// let mem_available = memory::system_memory_available_bytes();
202    /// ```
203    pub fn system_memory_available_bytes() -> usize {
204        get_resource(
205            "memory available",
206            super::Controller::Memory,
207            get_system_memory,
208            get_container_memory_limit,
209        )
210    }
211
212    /// Returns the memory limit of a container if running in a container else returns the system
213    /// memory available.
214    /// When the limit is set to max, [`system_memory_available_bytes()`] will return default system
215    /// memory.
216    fn get_container_memory_limit(
217        cgroup_version: super::CgroupVersion,
218    ) -> Result<usize, std::io::Error> {
219        let limit_path = match cgroup_version {
220            super::CgroupVersion::V1 => V1_MEMORY_LIMIT_PATH,
221            super::CgroupVersion::V2 => V2_MEMORY_LIMIT_PATH,
222        };
223        let system = get_system_memory();
224        let value = super::util::read_usize_or_max(limit_path, system)?;
225        Ok(std::cmp::min(value, system))
226    }
227
228    /// Returns the memory used in a container if running in a container else returns the system
229    /// memory used.
230    fn get_container_memory_used(
231        cgroup_version: super::CgroupVersion,
232    ) -> Result<usize, std::io::Error> {
233        let usage_path = match cgroup_version {
234            super::CgroupVersion::V1 => V1_MEMORY_CURRENT_PATH,
235            super::CgroupVersion::V2 => V2_MEMORY_CURRENT_PATH,
236        };
237        let system = get_system_memory_used();
238        let value = super::util::read_usize_or_max(usage_path, system)?;
239        Ok(std::cmp::min(value, system))
240    }
241}
242
243pub mod cpu {
244    use std::thread;
245
246    use thiserror_ext::AsReport;
247
248    use super::runtime::get_resource;
249    use super::util::parse_error;
250
251    /// Default constant Cgroup paths and hierarchy.
252    const V1_CPU_QUOTA_PATH: &str = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us";
253    const V1_CPU_PERIOD_PATH: &str = "/sys/fs/cgroup/cpu/cpu.cfs_period_us";
254    const V2_CPU_LIMIT_PATH: &str = "/sys/fs/cgroup/cpu.max";
255
256    /// Returns the total number of cpu available as a float.
257    ///
258    /// If running in container, this function will return the cpu limit by the container. If not,
259    /// it will return the ```available_parallelism``` by the system. A panic will be invoked if
260    /// invoking process does not have permission to read appropriate values by
261    /// ```std::thread::available_parallelism``` or if the platform is not supported. The cgroup
262    /// mount point is assumed to be at /sys/fs/cgroup by default.
263    ///
264    ///
265    /// # Examples
266    ///
267    /// Basic usage:
268    /// ``` ignore
269    /// let cpu_available = cpu::total_cpu_available();
270    /// ```
271    pub fn total_cpu_available() -> f32 {
272        get_resource(
273            "cpu quota",
274            super::Controller::Cpu,
275            get_system_cpu,
276            get_container_cpu_limit,
277        )
278    }
279
280    /// Returns the CPU limit of the container.
281    fn get_container_cpu_limit(
282        cgroup_version: super::CgroupVersion,
283    ) -> Result<f32, std::io::Error> {
284        let max_cpu = get_system_cpu();
285        match cgroup_version {
286            super::CgroupVersion::V1 => {
287                get_cpu_limit_v1(V1_CPU_QUOTA_PATH, V1_CPU_PERIOD_PATH, max_cpu)
288            }
289            super::CgroupVersion::V2 => get_cpu_limit_v2(V2_CPU_LIMIT_PATH, max_cpu),
290        }
291    }
292
293    /// Returns the total system cpu.
294    pub fn get_system_cpu() -> f32 {
295        match thread::available_parallelism() {
296            Ok(available_parallelism) => available_parallelism.get() as f32,
297            Err(e) => panic!(
298                "Failed to get available parallelism, error: {}",
299                e.as_report()
300            ),
301        }
302    }
303
304    /// Returns the CPU limit when cgroup v1 is utilised.
305    pub fn get_cpu_limit_v1(
306        quota_path: &str,
307        period_path: &str,
308        max_value: f32,
309    ) -> Result<f32, std::io::Error> {
310        let content = std::fs::read_to_string(quota_path)?;
311        let cpu_quota = content
312            .trim()
313            .parse::<i64>()
314            .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "not a number"))?;
315        // According to the kernel documentation, if the value is negative, it means no limit.
316        // https://docs.kernel.org/scheduler/sched-bwc.html#management
317        if cpu_quota < 0 {
318            return Ok(max_value);
319        }
320
321        let cpu_period = super::util::read_usize(period_path)?;
322
323        Ok((cpu_quota as f32) / (cpu_period as f32))
324    }
325
326    /// Returns the CPU limit when cgroup v2 is utilised.
327    pub fn get_cpu_limit_v2(limit_path: &str, max_value: f32) -> Result<f32, std::io::Error> {
328        let cpu_limit_string = fs_err::read_to_string(limit_path)?;
329
330        let cpu_data: Vec<&str> = cpu_limit_string.split_whitespace().collect();
331        match cpu_data.get(0..2) {
332            Some(cpu_data_values) => {
333                if cpu_data_values[0] == super::DEFAULT_CGROUP_MAX_INDICATOR {
334                    return Ok(max_value);
335                }
336                let cpu_quota = cpu_data_values[0]
337                    .parse::<usize>()
338                    .map_err(|e| parse_error(limit_path, &cpu_limit_string, e))?;
339                let cpu_period = cpu_data_values[1]
340                    .parse::<usize>()
341                    .map_err(|e| parse_error(limit_path, &cpu_limit_string, e))?;
342                Ok((cpu_quota as f32) / (cpu_period as f32))
343            }
344            None => Err(std::io::Error::new(
345                std::io::ErrorKind::InvalidData,
346                format!(
347                    "Invalid format in Cgroup CPU interface file, path: {limit_path}, content: {cpu_limit_string}"
348                ),
349            )),
350        }
351    }
352}
353
354mod util {
355    /// Parses the filepath and checks for the existence of `controller_name` in the file.
356    pub fn parse_controller_enable_file_for_cgroup_v2(
357        file_path: &str,
358        controller_name: &str,
359    ) -> bool {
360        match fs_err::read_to_string(file_path) {
361            Ok(controller_string) => {
362                for controller in controller_string.split_whitespace() {
363                    if controller.eq(controller_name) {
364                        return true;
365                    };
366                }
367                false
368            }
369            Err(_) => false,
370        }
371    }
372
373    pub fn parse_error(
374        file_path: &str,
375        content: &str,
376        e: impl std::fmt::Display,
377    ) -> std::io::Error {
378        std::io::Error::new(
379            std::io::ErrorKind::InvalidData,
380            format!("failed to parse, path: {file_path}, content: {content}, error: {e}"),
381        )
382    }
383
384    /// Reads an integer value from a file path.
385    pub fn read_usize(file_path: &str) -> Result<usize, std::io::Error> {
386        let content = fs_err::read_to_string(file_path)?;
387        let limit_val = content
388            .trim()
389            .parse::<usize>()
390            .map_err(|e| parse_error(file_path, &content, e))?;
391        Ok(limit_val)
392    }
393
394    /// Helper function that helps to retrieve value in file, if value is "max", `max_value` will be
395    /// returned instead.
396    pub fn read_usize_or_max(file_path: &str, max_value: usize) -> Result<usize, std::io::Error> {
397        let content = fs_err::read_to_string(file_path)?;
398        if content.trim() == super::DEFAULT_CGROUP_MAX_INDICATOR {
399            return Ok(max_value);
400        }
401        let limit_val = content
402            .trim()
403            .parse::<usize>()
404            .map_err(|e| parse_error(file_path, &content, e))?;
405        Ok(limit_val)
406    }
407
408    #[cfg(test)]
409    mod tests {
410        use std::collections::HashMap;
411        use std::io::prelude::*;
412        use std::thread;
413
414        use super::*;
415        use crate::cpu::{self, get_system_cpu};
416        use crate::memory::get_system_memory;
417        use crate::{Controller, DEFAULT_CGROUP_MAX_INDICATOR};
418        const DEFAULT_NON_EXISTENT_PATH: &str = "default-non-existent-path";
419
420        #[test]
421        fn test_read_integer_from_file_path() {
422            struct TestCase {
423                file_exists: bool,
424                value_in_file: String,
425                expected: Result<usize, std::io::Error>,
426            }
427
428            let test_cases = HashMap::from([
429                (
430                    "valid-integer-value-in-file",
431                    TestCase {
432                        file_exists: true,
433                        value_in_file: String::from("10000"),
434                        expected: Ok(10000),
435                    },
436                ),
437                (
438                    "valid-integer-value-in-file-with-spaces-after",
439                    TestCase {
440                        file_exists: true,
441                        value_in_file: String::from("10000   "),
442                        expected: Ok(10000),
443                    },
444                ),
445                (
446                    "valid-integer-value-in-file-with-spaces-before",
447                    TestCase {
448                        file_exists: true,
449                        value_in_file: String::from("   10000"),
450                        expected: Ok(10000),
451                    },
452                ),
453                (
454                    "invalid-integer-value-in-file",
455                    TestCase {
456                        file_exists: true,
457                        value_in_file: String::from("test-string"),
458                        expected: Err(std::io::Error::new(
459                            std::io::ErrorKind::InvalidData,
460                            "not a number",
461                        )),
462                    },
463                ),
464                (
465                    "file-not-exist",
466                    TestCase {
467                        file_exists: false,
468                        value_in_file: String::from(""),
469                        expected: Err(std::io::Error::new(
470                            std::io::ErrorKind::NotFound,
471                            "File not found",
472                        )),
473                    },
474                ),
475                (
476                    "max-value-in-file",
477                    TestCase {
478                        file_exists: true,
479                        value_in_file: String::from(DEFAULT_CGROUP_MAX_INDICATOR),
480                        expected: Err(std::io::Error::new(
481                            std::io::ErrorKind::InvalidData,
482                            "not a number",
483                        )),
484                    },
485                ),
486            ]);
487
488            for tc in test_cases {
489                let curr_test_case = &tc.1;
490                let mut file: tempfile::NamedTempFile;
491                let mut test_file_path = String::from(DEFAULT_NON_EXISTENT_PATH);
492                if curr_test_case.file_exists {
493                    file = tempfile::NamedTempFile::new()
494                        .expect("Error encountered while creating file!");
495                    file.as_file_mut()
496                        .write_all(curr_test_case.value_in_file.as_bytes())
497                        .expect("Error while writing to file");
498                    test_file_path = String::from(file.path().to_str().unwrap())
499                }
500                match read_usize(&test_file_path) {
501                    Ok(int_val) => assert_eq!(&int_val, curr_test_case.expected.as_ref().unwrap()),
502                    Err(e) => assert_eq!(
503                        e.kind(),
504                        curr_test_case.expected.as_ref().unwrap_err().kind()
505                    ),
506                }
507            }
508        }
509
510        #[test]
511        fn test_get_value_from_file() {
512            struct TestCase {
513                file_exists: bool,
514                value_in_file: String,
515                expected: Result<usize, std::io::Error>,
516            }
517
518            let test_cases = HashMap::from([
519                (
520                    "valid-integer-value-in-file",
521                    TestCase {
522                        file_exists: true,
523                        value_in_file: String::from("10000"),
524                        expected: Ok(10000),
525                    },
526                ),
527                (
528                    "valid-integer-value-in-file-with-spaces-after",
529                    TestCase {
530                        file_exists: true,
531                        value_in_file: String::from("10000   "),
532                        expected: Ok(10000),
533                    },
534                ),
535                (
536                    "valid-integer-value-in-file-with-spaces-before",
537                    TestCase {
538                        file_exists: true,
539                        value_in_file: String::from("   10000"),
540                        expected: Ok(10000),
541                    },
542                ),
543                (
544                    "invalid-integer-value-in-file",
545                    TestCase {
546                        file_exists: true,
547                        value_in_file: String::from("test-string"),
548                        expected: Err(std::io::Error::new(
549                            std::io::ErrorKind::InvalidData,
550                            "not a number",
551                        )),
552                    },
553                ),
554                (
555                    "file-not-exist",
556                    TestCase {
557                        file_exists: false,
558                        value_in_file: String::from(""),
559                        expected: Err(std::io::Error::new(
560                            std::io::ErrorKind::NotFound,
561                            "File not found",
562                        )),
563                    },
564                ),
565                (
566                    "max-value-in-file",
567                    TestCase {
568                        file_exists: true,
569                        value_in_file: String::from(DEFAULT_CGROUP_MAX_INDICATOR),
570                        expected: Ok(get_system_memory()),
571                    },
572                ),
573            ]);
574
575            for tc in test_cases {
576                let curr_test_case = &tc.1;
577                let mut file: tempfile::NamedTempFile;
578                let mut test_file_path = String::from(DEFAULT_NON_EXISTENT_PATH);
579                if curr_test_case.file_exists {
580                    file = tempfile::NamedTempFile::new()
581                        .expect("Error encountered while creating file!");
582                    file.as_file_mut()
583                        .write_all(curr_test_case.value_in_file.as_bytes())
584                        .expect("Error while writing to file");
585                    test_file_path = String::from(file.path().to_str().unwrap())
586                }
587                match read_usize_or_max(&test_file_path, get_system_memory()) {
588                    Ok(int_val) => assert_eq!(&int_val, curr_test_case.expected.as_ref().unwrap()),
589                    Err(e) => assert_eq!(
590                        e.kind(),
591                        curr_test_case.expected.as_ref().unwrap_err().kind()
592                    ),
593                }
594            }
595        }
596
597        #[test]
598        fn test_get_cpu_limit_v1() {
599            #[derive(Debug)]
600            struct TestCase {
601                file_exists: bool,
602                value_in_quota_file: String,
603                value_in_period_file: String,
604                expected: Result<f32, std::io::Error>,
605            }
606
607            let test_cases = HashMap::from([
608                (
609                    "default-values",
610                    TestCase {
611                        file_exists: true,
612                        value_in_quota_file: String::from("-1"),
613                        value_in_period_file: String::from("10000"),
614                        expected: Ok(thread::available_parallelism().unwrap().get() as f32),
615                    },
616                ),
617                (
618                    "valid-values-in-file",
619                    TestCase {
620                        file_exists: true,
621                        value_in_quota_file: String::from("10000"),
622                        value_in_period_file: String::from("20000"),
623                        expected: Ok(10000.0 / 20000.0),
624                    },
625                ),
626                (
627                    "empty-value-in-files",
628                    TestCase {
629                        file_exists: true,
630                        value_in_quota_file: String::from(""),
631                        value_in_period_file: String::from(""),
632                        expected: Err(std::io::Error::new(
633                            std::io::ErrorKind::InvalidData,
634                            "Invalid format in Cgroup CPU interface file",
635                        )),
636                    },
637                ),
638                (
639                    "Invalid-string-value-in-file",
640                    TestCase {
641                        file_exists: true,
642                        value_in_quota_file: String::from("10000"),
643                        value_in_period_file: String::from("test-string "),
644                        expected: Err(std::io::Error::new(
645                            std::io::ErrorKind::InvalidData,
646                            "not a number",
647                        )),
648                    },
649                ),
650                (
651                    "negative-value-in-file",
652                    TestCase {
653                        file_exists: true,
654                        value_in_quota_file: String::from("-2"),
655                        value_in_period_file: String::from("20000"),
656                        expected: Ok(thread::available_parallelism().unwrap().get() as f32),
657                    },
658                ),
659                (
660                    "file-not-exist",
661                    TestCase {
662                        file_exists: false,
663                        value_in_quota_file: String::from("10000 20000"),
664                        value_in_period_file: String::from("10000 20000"),
665                        expected: Err(std::io::Error::new(
666                            std::io::ErrorKind::NotFound,
667                            "File not found",
668                        )),
669                    },
670                ),
671            ]);
672            for tc in test_cases {
673                let curr_test_case = &tc.1;
674                let mut quota_file: tempfile::NamedTempFile;
675                let mut period_file: tempfile::NamedTempFile;
676                let mut test_quota_file_path = String::from(DEFAULT_NON_EXISTENT_PATH);
677                let mut test_period_file_path = String::from(DEFAULT_NON_EXISTENT_PATH);
678                if curr_test_case.file_exists {
679                    quota_file = tempfile::NamedTempFile::new()
680                        .expect("Error encountered while creating file!");
681                    quota_file
682                        .as_file_mut()
683                        .write_all(curr_test_case.value_in_quota_file.as_bytes())
684                        .expect("Error while writing to file");
685                    test_quota_file_path = String::from(quota_file.path().to_str().unwrap());
686
687                    period_file = tempfile::NamedTempFile::new()
688                        .expect("Error encountered while creating file!");
689                    period_file
690                        .as_file_mut()
691                        .write_all(curr_test_case.value_in_period_file.as_bytes())
692                        .expect("Error while writing to file");
693                    test_period_file_path = String::from(period_file.path().to_str().unwrap());
694                }
695                match cpu::get_cpu_limit_v1(
696                    &test_quota_file_path,
697                    &test_period_file_path,
698                    get_system_cpu(),
699                ) {
700                    Ok(int_val) => assert_eq!(
701                        &int_val,
702                        curr_test_case.expected.as_ref().unwrap(),
703                        "{:?}",
704                        tc
705                    ),
706                    Err(e) => assert_eq!(
707                        e.kind(),
708                        curr_test_case.expected.as_ref().unwrap_err().kind()
709                    ),
710                }
711            }
712        }
713
714        #[test]
715        fn test_get_cpu_limit_v2() {
716            struct TestCase {
717                file_exists: bool,
718                value_in_file: String,
719                expected: Result<f32, std::io::Error>,
720            }
721
722            let test_cases = HashMap::from([
723                (
724                    "valid-values-in-file",
725                    TestCase {
726                        file_exists: true,
727                        value_in_file: String::from("10000 20000"),
728                        expected: Ok(10000.0 / 20000.0),
729                    },
730                ),
731                (
732                    "Invalid-single-value-in-file",
733                    TestCase {
734                        file_exists: true,
735                        value_in_file: String::from("10000"),
736                        expected: Err(std::io::Error::new(
737                            std::io::ErrorKind::InvalidData,
738                            "Invalid format in Cgroup CPU interface file",
739                        )),
740                    },
741                ),
742                (
743                    "Invalid-string-value-in-file",
744                    TestCase {
745                        file_exists: true,
746                        value_in_file: String::from("10000 test-string "),
747                        expected: Err(std::io::Error::new(
748                            std::io::ErrorKind::InvalidData,
749                            "not a number",
750                        )),
751                    },
752                ),
753                (
754                    "max-value-in-file",
755                    TestCase {
756                        file_exists: true,
757                        value_in_file: String::from("max 20000"),
758                        expected: Ok(thread::available_parallelism().unwrap().get() as f32),
759                    },
760                ),
761                (
762                    "file-not-exist",
763                    TestCase {
764                        file_exists: false,
765                        value_in_file: String::from(""),
766                        expected: Err(std::io::Error::new(
767                            std::io::ErrorKind::NotFound,
768                            "File not found",
769                        )),
770                    },
771                ),
772            ]);
773            for tc in test_cases {
774                let curr_test_case = &tc.1;
775                let mut file: tempfile::NamedTempFile;
776                let mut test_file_path = String::from(DEFAULT_NON_EXISTENT_PATH);
777                if curr_test_case.file_exists {
778                    file = tempfile::NamedTempFile::new()
779                        .expect("Error encountered while creating file!");
780                    file.as_file_mut()
781                        .write_all(curr_test_case.value_in_file.as_bytes())
782                        .expect("Error while writing to file");
783                    test_file_path = String::from(file.path().to_str().unwrap())
784                }
785                match cpu::get_cpu_limit_v2(&test_file_path, get_system_cpu()) {
786                    Ok(int_val) => assert_eq!(&int_val, curr_test_case.expected.as_ref().unwrap()),
787                    Err(e) => assert_eq!(
788                        e.kind(),
789                        curr_test_case.expected.as_ref().unwrap_err().kind()
790                    ),
791                }
792            }
793        }
794
795        #[test]
796        fn test_parse_controller_enable_file_for_cgroup_v2() {
797            struct TestCase {
798                file_exists: bool,
799                value_in_file: String,
800                controller_type: Controller,
801                expected: bool,
802            }
803
804            let test_cases = HashMap::from([
805                (
806                    "cpu-enabled",
807                    TestCase {
808                        file_exists: true,
809                        value_in_file: String::from("cpu memory IO"),
810                        controller_type: Controller::Cpu,
811                        expected: true,
812                    },
813                ),
814                (
815                    "memory-enabled",
816                    TestCase {
817                        file_exists: true,
818                        value_in_file: String::from("cpu memory IO"),
819                        controller_type: Controller::Memory,
820                        expected: true,
821                    },
822                ),
823                (
824                    "memory-disabled",
825                    TestCase {
826                        file_exists: true,
827                        value_in_file: String::from("cpu IO"),
828                        controller_type: Controller::Memory,
829                        expected: false,
830                    },
831                ),
832                (
833                    "cpu-disabled",
834                    TestCase {
835                        file_exists: true,
836                        value_in_file: String::from("memory IO"),
837                        controller_type: Controller::Cpu,
838                        expected: false,
839                    },
840                ),
841                (
842                    "Invalid-value-in-file",
843                    TestCase {
844                        file_exists: true,
845                        value_in_file: String::from("test-string test-string"),
846                        controller_type: Controller::Cpu,
847                        expected: false,
848                    },
849                ),
850                (
851                    "controller-file-not-exist",
852                    TestCase {
853                        file_exists: false,
854                        value_in_file: String::from(""),
855                        controller_type: Controller::Memory,
856                        expected: false,
857                    },
858                ),
859            ]);
860
861            for tc in test_cases {
862                let curr_test_case = &tc.1;
863                let controller_name: &str = match curr_test_case.controller_type {
864                    Controller::Cpu => "cpu",
865                    Controller::Memory => "memory",
866                };
867                let mut file: tempfile::NamedTempFile;
868                let mut test_file_path = String::from(DEFAULT_NON_EXISTENT_PATH);
869                if curr_test_case.file_exists {
870                    file = tempfile::NamedTempFile::new()
871                        .expect("Error encountered while creating file!");
872                    file.as_file_mut()
873                        .write_all(curr_test_case.value_in_file.as_bytes())
874                        .expect("Error while writing to file");
875                    test_file_path = String::from(file.path().to_str().unwrap())
876                }
877                assert_eq!(
878                    parse_controller_enable_file_for_cgroup_v2(&test_file_path, controller_name),
879                    curr_test_case.expected
880                );
881            }
882        }
883    }
884}
rw_resource_util/lib.rs

rw_resource_util/
lib.rs
