Skip to main content

linuxutils_system/
choom.rs

1use linuxutils_common::man::ManContent;
2
3pub const MAN: ManContent = ManContent::empty();
4
5use clap::Parser;
6use std::{fs, os::unix::process::CommandExt, process, process::ExitCode};
7
8#[derive(Parser)]
9#[command(
10    name = "choom",
11    about = "Display and adjust OOM-killer score",
12    override_usage = "choom [options] -p pid\n       \
13                      choom [options] -n number -p pid\n       \
14                      choom [options] -n number [--] command [args...]"
15)]
16pub struct Args {
17    /// OOM score adjust value to set (-1000 to 1000)
18    #[arg(short = 'n', long = "adjust", value_name = "num")]
19    adjust: Option<i16>,
20
21    /// Process ID to display or adjust
22    #[arg(short = 'p', long, value_name = "num")]
23    pid: Option<u32>,
24
25    /// Command and arguments to run with adjusted OOM score
26    #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
27    command: Vec<String>,
28}
29
30fn read_oom(pid: u32) -> std::io::Result<(i32, i16)> {
31    let score: i32 = fs::read_to_string(format!("/proc/{pid}/oom_score"))?
32        .trim()
33        .parse()
34        .map_err(std::io::Error::other)?;
35    let adj: i16 = fs::read_to_string(format!("/proc/{pid}/oom_score_adj"))?
36        .trim()
37        .parse()
38        .map_err(std::io::Error::other)?;
39    Ok((score, adj))
40}
41
42fn write_oom_adj(pid: u32, adj: i16) -> std::io::Result<()> {
43    fs::write(format!("/proc/{pid}/oom_score_adj"), format!("{adj}\n"))
44}
45
46pub fn run(args: Args) -> ExitCode {
47    // Validate: needs either -p or a command (but not both).
48    let has_command = !args.command.is_empty();
49
50    if args.pid.is_none() && !has_command {
51        eprintln!("choom: no PID specified and no command given");
52        eprintln!("Try 'choom --help' for more information.");
53        return ExitCode::FAILURE;
54    }
55
56    if args.pid.is_some() && has_command {
57        eprintln!("choom: --pid and command are mutually exclusive");
58        return ExitCode::FAILURE;
59    }
60
61    if has_command {
62        // Set our own OOM adjust value, then exec the command.
63        if let Some(adj) = args.adjust
64            && let Err(e) = write_oom_adj(std::process::id(), adj)
65        {
66            eprintln!("choom: failed to set OOM score: {e}");
67            return ExitCode::FAILURE;
68        }
69        let (prog, prog_args) = args.command.split_first().unwrap();
70        let err = process::Command::new(prog).args(prog_args).exec();
71        eprintln!("choom: {prog}: {err}");
72        return ExitCode::FAILURE;
73    }
74
75    let pid = args.pid.unwrap();
76
77    if let Some(adj) = args.adjust
78        && let Err(e) = write_oom_adj(pid, adj)
79    {
80        eprintln!("choom: failed to set OOM score for pid {pid}: {e}");
81        return ExitCode::FAILURE;
82    }
83
84    match read_oom(pid) {
85        Ok((score, adj)) => {
86            println!("pid {pid}'s current OOM score: {score}");
87            println!("pid {pid}'s current OOM score adjust value: {adj}");
88            ExitCode::SUCCESS
89        }
90        Err(e) => {
91            eprintln!("choom: failed to read OOM score for pid {pid}: {e}");
92            ExitCode::FAILURE
93        }
94    }
95}