diff --git a/Cargo.lock b/Cargo.lock index 0c701f61..184138d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1023,6 +1023,7 @@ dependencies = [ "tempfile", "textwrap", "uu_blockdev", + "uu_chcpu", "uu_ctrlaltdel", "uu_dmesg", "uu_fsfreeze", @@ -1051,6 +1052,14 @@ dependencies = [ "uucore", ] +[[package]] +name = "uu_chcpu" +version = "0.0.1" +dependencies = [ + "clap", + "uucore", +] + [[package]] name = "uu_ctrlaltdel" version = "0.0.1" diff --git a/Cargo.toml b/Cargo.toml index 5bad5d97..ffda8ae4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ uudoc = [] feat_common_core = [ "blockdev", + "chcpu", "ctrlaltdel", "dmesg", "fsfreeze", @@ -76,6 +77,7 @@ uucore = { workspace = true } # blockdev = { optional = true, version = "0.0.1", package = "uu_blockdev", path = "src/uu/blockdev" } +chcpu = { optional = true, version = "0.0.1", package = "uu_chcpu", path = "src/uu/chcpu" } ctrlaltdel = { optional = true, version = "0.0.1", package = "uu_ctrlaltdel", path = "src/uu/ctrlaltdel" } dmesg = { optional = true, version = "0.0.1", package = "uu_dmesg", path = "src/uu/dmesg" } fsfreeze = { optional = true, version = "0.0.1", package = "uu_fsfreeze", path = "src/uu/fsfreeze" } diff --git a/src/uu/chcpu/Cargo.toml b/src/uu/chcpu/Cargo.toml new file mode 100644 index 00000000..f62c1843 --- /dev/null +++ b/src/uu/chcpu/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "uu_chcpu" +version = "0.0.1" +edition = "2021" + +[lib] +path = "src/chcpu.rs" + +[[bin]] +name = "chcpu" +path = "src/main.rs" + +[dependencies] +uucore = { workspace = true } +clap = { workspace = true } diff --git a/src/uu/chcpu/chcpu.md b/src/uu/chcpu/chcpu.md new file mode 100644 index 00000000..6ed48e84 --- /dev/null +++ b/src/uu/chcpu/chcpu.md @@ -0,0 +1,7 @@ +# chcpu + +``` +chcpu [OPTION]... +``` + +Configure CPUs in a multi-processor system. diff --git a/src/uu/chcpu/src/chcpu.rs b/src/uu/chcpu/src/chcpu.rs new file mode 100644 index 00000000..8b65df8b --- /dev/null +++ b/src/uu/chcpu/src/chcpu.rs @@ -0,0 +1,437 @@ +// This file is part of the uutils util-linux package. +// +// For the full copyright and license information, please view the LICENSE +// file that was distributed with this source code. + +use std::{ + fs::{self, File}, + io::Write, + path::PathBuf, +}; + +use clap::{crate_version, Arg, ArgAction, ArgGroup, Command}; +use uucore::{ + error::{UResult, USimpleError}, + format_usage, help_about, help_usage, +}; + +// Takes in a human-readable list of CPUs, and returns a list of indices parsed from that list +// These can come in the form of a plain range like `X-Y`, or a comma-separated ranges and indices ie. `1,3-4,7-8,10` +// Kernel docs with examples: https://www.kernel.org/doc/html/latest/admin-guide/cputopology.html +// TODO: Move function to uucore +fn parse_cpu_list(list: &str) -> UResult> { + let mut out: Vec = vec![]; + + if list.is_empty() { + return Ok(out); + } + + for part in list.trim().split(",") { + if part.contains("-") { + let bounds: Vec<_> = part.split("-").flat_map(|x| x.parse::()).collect(); + + if bounds.len() != 2 { + return Err(USimpleError::new(1, format!("Invalid CPU range: {}", part))); + } + + for idx in bounds[0]..bounds[1] + 1 { + out.push(idx) + } + } else { + let idx = part.parse::().expect("Invalid CPU index value"); + out.push(idx); + } + } + Ok(out) +} + +#[derive(Debug)] +struct Cpu(usize); + +impl Cpu { + fn get_path(&self) -> PathBuf { + PathBuf::from(format!("/sys/devices/system/cpu/cpu{}", self.0)) + } + + fn write_to_cpu_file(&self, file: &str, value: &[u8]) -> std::io::Result<()> { + File::create(self.get_path().join(file)).and_then(|mut f| f.write_all(value)) + } + + fn read_cpu_file(&self, file: &str) -> String { + fs::read_to_string(self.get_path().join(file)).unwrap() + } + + fn ensure_exists(&self) -> UResult { + if !self.get_path().exists() { + return Err(USimpleError::new( + 1, + format!("CPU {} does not exist", self.0), + )); + }; + Ok(true) + } + + // CPUs which are not hot-pluggable will not have the `/online` file in their directory + fn is_hotpluggable(&self) -> bool { + self.get_path().join("online").exists() + } + + fn is_online(&self) -> bool { + fs::read_to_string(self.get_path().join("online")) + .map(|content| match content.trim() { + "0" => false, + "1" => true, + other => panic!("Unrecognized CPU online state: {}", other), + }) + // Just in case the caller forgot to check `is_hotpluggable` first, + // instead of panicing that the file doesn't exist, return true + // This is because a non-hotpluggable CPU is assumed to be always online + .unwrap_or(true) + } + + fn enable(&self) -> UResult<()> { + if !self.is_hotpluggable() { + return Err(USimpleError::new( + 1, + format!("CPU {} is not hot-pluggable", self.0), + )); + } + + if self.is_online() { + return Err(USimpleError::new( + 1, + format!("CPU {} is already enabled", self.0), + )); + } + + match self.write_to_cpu_file("online", b"1") { + Ok(_) => println!("CPU {} enabled", self.0), + Err(e) => { + return Err(USimpleError::new( + 1, + format!("CPU {} enable failed: {}", self.0, e.kind()), + )) + } + }; + Ok(()) + } + + fn disable(&self) -> UResult<()> { + if !self.is_hotpluggable() { + return Err(USimpleError::new( + 1, + format!("CPU {} is not hot-pluggable", self.0), + )); + } + + if !self.is_online() { + return Err(USimpleError::new( + 1, + format!("CPU {} is already disabled", self.0), + )); + } + + if get_online_cpus()?.len() == 1 { + return Err(USimpleError::new( + 1, + format!("CPU {} disable failed (last enabled CPU)", self.0), + )); + } + + match self.write_to_cpu_file("online", b"0") { + Ok(_) => println!("CPU {} disabled", self.0), + Err(e) => { + return Err(USimpleError::new( + 1, + format!("CPU {} disable failed: {}", self.0, e.kind()), + )) + } + }; + + Ok(()) + } + + fn is_configurable(&self) -> bool { + self.get_path().join("configure").exists() + } + + fn configure(&self) -> UResult<()> { + if !self.is_configurable() { + return Err(USimpleError::new( + 1, + format!("CPU {} is not configurable", self.0), + )); + } + + let configured = self.read_cpu_file("configure"); + if configured.trim() == "1" { + return Err(USimpleError::new( + 1, + format!("CPU {} is already configured", self.0), + )); + }; + + match self.write_to_cpu_file("configure", b"1") { + Ok(_) => println!("CPU {} configured", self.0), + Err(e) => { + return Err(USimpleError::new( + 1, + format!("CPU {} configure failed: {}", self.0, e.kind()), + )) + } + }; + Ok(()) + } + + fn deconfigure(&self) -> UResult<()> { + if !self.is_configurable() { + return Err(USimpleError::new( + 1, + format!("CPU {} is not configurable", self.0), + )); + } + + let configured = self.read_cpu_file("configure"); + if configured.trim() == "0" { + return Err(USimpleError::new( + 1, + format!("CPU {} is already deconfigured", self.0), + )); + }; + + if self.is_online() { + return Err(USimpleError::new( + 1, + format!("CPU {} deconfigure failed (CPU is enabled)", self.0), + )); + } + + match self.write_to_cpu_file("configure", b"0") { + Ok(_) => println!("CPU {} deconfigured", self.0), + Err(e) => { + return Err(USimpleError::new( + 1, + format!("CPU {} deconfigure failed: {}", self.0, e.kind()), + )) + } + }; + + Ok(()) + } +} + +fn get_online_cpus() -> UResult> { + let cpu_list = fs::read_to_string("/sys/devices/system/cpu/online")?; + let cpus = parse_cpu_list(&cpu_list)?.iter().map(|n| Cpu(*n)).collect(); + Ok(cpus) +} + +fn trigger_rescan() -> UResult<()> { + let path = PathBuf::from("/sys/devices/system/cpu/rescan"); + + if !path.exists() { + return Err(USimpleError::new( + 1, + "This system does not support rescanning of CPUs", + )); + } + + let result = File::create(path).and_then(|mut f| f.write_all(b"1")); + match result { + Ok(_) => println!("Triggered rescan of CPUs"), + Err(e) => { + return Err(USimpleError::new( + 1, + format!("Failed to trigger rescan of CPUs: {}", e.kind()), + )) + } + }; + + Ok(()) +} + +fn process_cpus(cpu_list: &str, action: impl Fn(&Cpu) -> UResult<()>) -> UResult<()> { + parse_cpu_list(cpu_list)? + .into_iter() + .map(Cpu) + .try_for_each(|cpu| cpu.ensure_exists().and_then(|_| action(&cpu)))?; + Ok(()) +} + +fn enable_cpus(cpu_list: &str) -> UResult<()> { + process_cpus(cpu_list, Cpu::enable) +} + +fn disable_cpus(cpu_list: &str) -> UResult<()> { + process_cpus(cpu_list, Cpu::disable) +} + +fn configure_cpus(cpu_list: &str) -> UResult<()> { + process_cpus(cpu_list, Cpu::configure) +} + +fn deconfigure_cpus(cpu_list: &str) -> UResult<()> { + process_cpus(cpu_list, Cpu::deconfigure) +} + +fn set_dispatch_mode(mode: &str) -> UResult<()> { + let mode_num: u8 = match mode { + "horizontal" => 0, + "vertical" => 1, + _ => { + return Err(USimpleError::new( + 1, + format!( + "Unsupported dispatching mode: '{}'. Must be either 'horizontal' or 'vertical'", + mode + ), + )) + } + }; + + let path = PathBuf::from("/sys/devices/system/cpu/dispatching"); + + if !path.exists() { + return Err(USimpleError::new( + 1, + "This system does not support setting the dispatching mode of CPUs", + )); + } + + let result = File::create(path).and_then(|mut f| f.write_all(&[mode_num])); + match result { + Ok(_) => println!("Successfully set {} dispatching mode", mode), + Err(e) => { + return Err(USimpleError::new( + 1, + format!("Failed to set {} dispatching mode: {}", mode, e.kind()), + )) + } + }; + + Ok(()) +} + +#[uucore::main] +pub fn uumain(args: impl uucore::Args) -> UResult<()> { + let matches = uu_app().try_get_matches_from(args)?; + + if matches.get_flag(options::RESCAN) { + trigger_rescan()?; + } + + if let Some(cpu_list) = matches.get_one::(options::ENABLE) { + enable_cpus(cpu_list)?; + } + + if let Some(cpu_list) = matches.get_one::(options::DISABLE) { + disable_cpus(cpu_list)?; + } + + if let Some(cpu_list) = matches.get_one::(options::CONFIGURE) { + configure_cpus(cpu_list)?; + } + + if let Some(cpu_list) = matches.get_one::(options::DECONFIGURE) { + deconfigure_cpus(cpu_list)?; + } + + if let Some(mode) = matches.get_one::(options::DISPATCH) { + set_dispatch_mode(mode)?; + } + + Ok(()) +} + +mod options { + pub const ENABLE: &str = "enable"; + pub const DISABLE: &str = "disable"; + pub const CONFIGURE: &str = "configure"; + pub const DECONFIGURE: &str = "deconfigure"; + pub const DISPATCH: &str = "dispatch"; + pub const RESCAN: &str = "rescan"; +} + +const ABOUT: &str = help_about!("chcpu.md"); +const USAGE: &str = help_usage!("chcpu.md"); + +pub fn uu_app() -> Command { + Command::new(uucore::util_name()) + .version(crate_version!()) + .about(ABOUT) + .override_usage(format_usage(USAGE)) + .infer_long_args(true) + .arg( + Arg::new(options::ENABLE) + .short('e') + .long("enable") + .value_name("cpu-list") + .action(ArgAction::Set), + ) + .arg( + Arg::new(options::DISABLE) + .short('d') + .long("disable") + .value_name("cpu-list") + .action(ArgAction::Set), + ) + .arg( + Arg::new(options::CONFIGURE) + .short('c') + .long("configure") + .value_name("cpu-list") + .action(ArgAction::Set), + ) + .arg( + Arg::new(options::DECONFIGURE) + .short('g') + .long("deconfigure") + .value_name("cpu-list") + .action(ArgAction::Set), + ) + .arg( + Arg::new(options::DISPATCH) + .short('p') + .long("dispatch") + .value_name("mode") + .action(ArgAction::Set), + ) + .group( + ArgGroup::new("action") + .args(vec![ + options::ENABLE, + options::DISABLE, + options::CONFIGURE, + options::DECONFIGURE, + options::DISPATCH, + ]) + .multiple(false), // These 5 are mutually exclusive + ) + .arg( + Arg::new(options::RESCAN) + .short('r') + .long("rescan") + .action(ArgAction::SetTrue), + ) +} + +#[test] +fn test_parse_cpu_list() { + assert_eq!(parse_cpu_list("").unwrap(), Vec::::new()); + assert_eq!( + parse_cpu_list("1-3").unwrap(), + Vec::::from([1, 2, 3]) + ); + assert_eq!( + parse_cpu_list("1,2,3").unwrap(), + Vec::::from([1, 2, 3]) + ); + assert_eq!( + parse_cpu_list("1,3-6,8").unwrap(), + Vec::::from([1, 3, 4, 5, 6, 8]) + ); + assert_eq!( + parse_cpu_list("1-2,3-5,7").unwrap(), + Vec::::from([1, 2, 3, 4, 5, 7]) + ); +} diff --git a/src/uu/chcpu/src/main.rs b/src/uu/chcpu/src/main.rs new file mode 100644 index 00000000..fe72ddc7 --- /dev/null +++ b/src/uu/chcpu/src/main.rs @@ -0,0 +1 @@ +uucore::bin!(uu_chcpu); diff --git a/tests/by-util/test_chcpu.rs b/tests/by-util/test_chcpu.rs new file mode 100644 index 00000000..603728dc --- /dev/null +++ b/tests/by-util/test_chcpu.rs @@ -0,0 +1,42 @@ +// This file is part of the uutils util-linux package. +// +// For the full copyright and license information, please view the LICENSE +// file that was distributed with this source code. + +use crate::common::util::TestScenario; + +#[test] +fn test_invalid_arg() { + new_ucmd!().arg("--definitely-invalid").fails().code_is(1); +} + +#[test] +fn test_invalid_cpu_range() { + new_ucmd!() + .arg("-e") + .arg("non-numeric-range") + .fails() + .code_is(1); +} + +#[test] +fn test_invalid_cpu_index() { + new_ucmd!() + .arg("-e") + .arg("10000") // Assuming no test environment will ever have 10000 CPUs + .fails() + .code_is(1) + .stderr_contains("CPU 10000 does not exist"); +} + +#[test] +fn test_invalid_dispatch_mode() { + new_ucmd!() + .arg("-p") + .arg("not-horizontal-or-vertical") + .fails() + .code_is(1) + .stderr_contains("Unsupported dispatching mode"); +} + +// TODO: Find a way to implement "happy-case" tests that doesn't rely on the host `/sys/` filesystem diff --git a/tests/tests.rs b/tests/tests.rs index 0fb3ed14..59e1bbf3 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -57,6 +57,10 @@ mod test_dmesg; #[path = "by-util/test_fsfreeze.rs"] mod test_fsfreeze; +#[cfg(feature = "chcpu")] +#[path = "by-util/test_chcpu.rs"] +mod test_chcpu; + #[cfg(feature = "mcookie")] #[path = "by-util/test_mcookie.rs"] mod test_mcookie;