cpubind_bash_builtin/
lib.rs

1#![doc = include_str!("../README.md")]
2use bash_builtins::variables::find_as_string;
3use bash_builtins::{builtin_metadata, Args, Builtin, BuiltinOptions, Result};
4use core_affinity::get_core_ids;
5use std::io::{self, BufWriter, Write};
6
7builtin_metadata!(
8    name = "cpubind",
9    create = CpuBind::default,
10    short_doc = "cpubind [-i identifier]",
11    long_doc = "
12    Prints information about the task and it's cpu affinity
13
14    Options:
15        -i identifier used to identify that particular task
16    ",
17);
18
19#[derive(BuiltinOptions)]
20enum Opt {
21    #[opt = 'i']
22    Identifier(String),
23}
24
25#[derive(Default)]
26struct CpuBind;
27
28/// Gets the variable named var if it exists and return its value
29/// or an empty string if it does not exists.
30fn get_variable_string(var: &str) -> String {
31    let var_ostring = find_as_string(var);
32    if let Some(var_str) = var_ostring.as_ref().and_then(|v| v.to_str().ok()) {
33        var_str.to_string()
34    } else {
35        "".to_string()
36    }
37}
38
39impl Builtin for CpuBind {
40    fn call(&mut self, args: &mut Args) -> Result<()> {
41        let mut identifier: String = "".to_string();
42
43        // managing options argument if any - none is ok
44        if !args.is_empty() {
45            for opt in args.options() {
46                match opt? {
47                    Opt::Identifier(s) => identifier = s,
48                };
49            }
50        }
51
52        // It is an error if we receive free arguments.
53        args.finished()?;
54
55        let core_ids = match get_core_ids() {
56            Some(core_ids) => core_ids,
57            None => Vec::new(),
58        };
59
60        let slurm_job_id = get_variable_string("SLURM_JOB_ID");
61        let slurm_procid = get_variable_string("SLURM_PROCID");
62        let slurm_localid = get_variable_string("SLURM_LOCALID");
63
64        let slurm_str = format!("{identifier} - {slurm_job_id} - {slurm_procid} - {slurm_localid}");
65
66        let hostname = match hostname::get()?.into_string() {
67            Ok(string) => string,
68            Err(_) => "".to_string(),
69        };
70
71        // using write!() and writeln!() instead of println!() to avoid
72        // panicking if stdout is closed.
73        let stdout_handle = io::stdout();
74        let mut output = BufWriter::new(stdout_handle.lock());
75
76        write!(&mut output, "{hostname} - {slurm_str} - cpu affinity:")?;
77        for core in core_ids.into_iter() {
78            write!(&mut output, " {}", core.id)?
79        }
80        writeln!(&mut output)?;
81
82        Ok(())
83    }
84}
cpubind_bash_builtin/lib.rs

cpubind_bash_builtin/
lib.rs
