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        String::new()
36    }
37}
38
39impl Builtin for CpuBind {
40    fn call(&mut self, args: &mut Args) -> Result<()> {
41        let mut identifier: String = String::new();
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 = get_core_ids().unwrap_or_default();
56
57        let slurm_job_id = get_variable_string("SLURM_JOB_ID");
58        let slurm_procid = get_variable_string("SLURM_PROCID");
59        let slurm_localid = get_variable_string("SLURM_LOCALID");
60
61        let slurm_str = format!("{identifier} - {slurm_job_id} - {slurm_procid} - {slurm_localid}");
62
63        let hostname = hostname::get()?.into_string().unwrap_or_default();
64
65        // using write!() and writeln!() instead of println!() to avoid
66        // panicking if stdout is closed.
67        let stdout_handle = io::stdout();
68        let mut output = BufWriter::new(stdout_handle.lock());
69
70        write!(&mut output, "{hostname} - {slurm_str} - cpu affinity:")?;
71        for core in core_ids {
72            write!(&mut output, " {}", core.id)?;
73        }
74        writeln!(&mut output)?;
75
76        Ok(())
77    }
78}
cpubind_bash_builtin/lib.rs

cpubind_bash_builtin/
lib.rs
