rtsc/
system.rs

1use crate::{Error, Result};
2
3/// Get absolute number of CPUs, including isolated (Linux only)
4#[cfg(target_os = "linux")]
5pub fn num_cpus() -> Result<usize> {
6    use std::io::BufRead;
7
8    let f = std::fs::File::open("/proc/cpuinfo")?;
9    let reader = std::io::BufReader::new(f);
10    let lines = reader.lines();
11    let mut count = 0;
12    for line in lines {
13        let line = line?;
14        if line
15            .split(':')
16            .next()
17            .ok_or_else(|| Error::Failed("invalid line".into()))?
18            .trim_end()
19            == "processor"
20        {
21            count += 1;
22        }
23    }
24    Ok(count)
25}
26
27/// Get absolute number of CPUs, including isolated (Linux only)
28#[cfg(not(target_os = "linux"))]
29pub fn num_cpus() -> Result<usize> {
30    Err(Error::Unimplemented)
31}
32
33/// Linux-specific system tools
34pub mod linux {
35    use crate::Result;
36
37    use core::fmt;
38    use std::{collections::BTreeMap, fs};
39    use tracing::warn;
40
41    /// Configure system parameters (global) while the process is running. Does nothing in simulated
42    /// mode
43    ///
44    /// Example:
45    ///
46    /// ```rust,no_run
47    /// use rtsc::system::linux::SystemConfig;
48    ///
49    /// let _sys = SystemConfig::new().set("kernel/sched_rt_runtime_us", -1)
50    ///     .apply()
51    ///     .expect("Unable to set system config");
52    /// // some code
53    /// // system config is restored at the end of the scope
54    /// ```
55    #[derive(Default)]
56    pub struct SystemConfig {
57        values: BTreeMap<&'static str, String>,
58        prev_values: BTreeMap<&'static str, String>,
59    }
60
61    impl SystemConfig {
62        /// Creates a new system config object
63        #[must_use]
64        pub fn new() -> Self {
65            Self::default()
66        }
67        /// Set a parameter to configure
68        pub fn set<V: fmt::Display>(mut self, key: &'static str, value: V) -> Self {
69            self.values.insert(key, value.to_string());
70            self
71        }
72        /// Apply values to /proc/sys keys
73        pub fn apply(mut self) -> Result<SystemConfigGuard> {
74            for (key, value) in &self.values {
75                let fname = format!("/proc/sys/{}", key);
76                let prev_value = fs::read_to_string(&fname)?;
77                self.prev_values.insert(key, prev_value);
78                fs::write(fname, value)?;
79            }
80            Ok(SystemConfigGuard { config: self })
81        }
82    }
83
84    /// A guard object to restore system parameters when dropped
85    #[derive(Default)]
86    pub struct SystemConfigGuard {
87        config: SystemConfig,
88    }
89
90    impl Drop for SystemConfigGuard {
91        fn drop(&mut self) {
92            for (key, value) in &self.config.prev_values {
93                if let Err(error) = fs::write(format!("/proc/sys/{}", key), value) {
94                    warn!(key, value, %error, "Failed to restore system config");
95                }
96            }
97        }
98    }
99
100    /// Configure CPU governors for the given CPUs
101    #[derive(Default)]
102    pub struct CpuGovernor {
103        prev_governor: BTreeMap<usize, String>,
104    }
105
106    impl CpuGovernor {
107        /// Set performance governor for the given CPUs. This sets the maximum frequency for the
108        /// CPUs, increasing the power consumption but lowering their latency. It is enough to
109        /// specify a single logical core number per physical core. The governor is restored when
110        /// the returned guard object is dropped.
111        pub fn performance<I>(performance_cpus: I) -> Result<CpuGovernor>
112        where
113            I: IntoIterator<Item = usize>,
114        {
115            let mut prev_governor = BTreeMap::new();
116            for cpu in performance_cpus {
117                let fname = format!(
118                    "/sys/devices/system/cpu/cpu{}/cpufreq/scaling_governor",
119                    cpu
120                );
121                let prev_value = fs::read_to_string(fname)?;
122                prev_governor.insert(cpu, prev_value.trim().to_string());
123            }
124            for cpu in prev_governor.keys() {
125                let fname = format!(
126                    "/sys/devices/system/cpu/cpu{}/cpufreq/scaling_governor",
127                    cpu
128                );
129                fs::write(fname, "performance")?;
130            }
131            Ok(CpuGovernor { prev_governor })
132        }
133    }
134
135    impl Drop for CpuGovernor {
136        fn drop(&mut self) {
137            for (cpu, governor) in &self.prev_governor {
138                if let Err(error) = fs::write(
139                    format!(
140                        "/sys/devices/system/cpu/cpu{}/cpufreq/scaling_governor",
141                        cpu
142                    ),
143                    governor,
144                ) {
145                    warn!(cpu, %error, "Failed to restore governor");
146                }
147            }
148        }
149    }
150}
rtsc/system.rs

rtsc/
system.rs
