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

rw_resource_util/
lib.rs
