procutils_hugetop/

lib.rs

use clap::Parser;
use crossterm::event::{self, Event, KeyCode, KeyEventKind};
use procfs::{prelude::*, process};
use procutils_common::{MAX_TERM_WIDTH, man::ManContent};
use ratatui::{
    DefaultTerminal, Frame,
    layout::{Constraint, Layout},
    style::{Color, Modifier, Style},
    text::{Line, Span},
    widgets::{Block, Borders, Cell, Paragraph, Row, Table},
};
use std::{process::ExitCode, time::Duration};

pub const MAN: ManContent = ManContent {
    description: Some(include_str!("../man/description.man")),
    extra_sections: &[
        ("KEY COMMANDS", include_str!("../man/key_commands.man")),
        (
            "FIELD DESCRIPTIONS",
            include_str!("../man/field_descriptions.man"),
        ),
        ("EXAMPLES", include_str!("../man/examples.man")),
        ("NOTES", include_str!("../man/notes.man")),
        ("SEE ALSO", include_str!("../man/see_also.man")),
    ],
};

/// Display hugepage usage by process.
#[derive(Parser)]
#[command(name = "hugetop", version, about, max_term_width = MAX_TERM_WIDTH)]
pub struct Args {
    /// Delay between updates in seconds.
    #[arg(short, long, default_value = "2")]
    delay: f64,

    /// Batch mode (non-interactive, for piping).
    #[arg(short, long)]
    batch: bool,

    /// Exit after this many display updates (only honoured with --batch).
    #[arg(short = 'n', long)]
    iterations: Option<u64>,
}

struct ProcessHuge {
    pid: i32,
    user: String,
    comm: String,
    hugetlb_kb: u64,
    anon_huge_kb: u64,
    total_huge_kb: u64,
}

struct SystemHuge {
    anon_hugepages_kb: u64,
    shmem_hugepages_kb: u64,
    hugepages_total: u64,
    hugepages_free: u64,
    hugepages_rsvd: u64,
    hugepages_surp: u64,
    hugepagesize_kb: u64,
    hugetlb_kb: u64,
}

impl SystemHuge {
    fn read() -> Self {
        let m = procfs::Meminfo::current().ok();
        Self {
            anon_hugepages_kb: m
                .as_ref()
                .and_then(|m| m.anon_hugepages)
                .unwrap_or(0)
                / 1024,
            shmem_hugepages_kb: m
                .as_ref()
                .and_then(|m| m.shmem_hugepages)
                .unwrap_or(0)
                / 1024,
            hugepages_total: m
                .as_ref()
                .and_then(|m| m.hugepages_total)
                .unwrap_or(0),
            hugepages_free: m
                .as_ref()
                .and_then(|m| m.hugepages_free)
                .unwrap_or(0),
            hugepages_rsvd: m
                .as_ref()
                .and_then(|m| m.hugepages_rsvd)
                .unwrap_or(0),
            hugepages_surp: m
                .as_ref()
                .and_then(|m| m.hugepages_surp)
                .unwrap_or(0),
            hugepagesize_kb: m
                .as_ref()
                .and_then(|m| m.hugepagesize)
                .unwrap_or(0)
                / 1024,
            hugetlb_kb: m.as_ref().and_then(|m| m.hugetlb).unwrap_or(0) / 1024,
        }
    }
}

fn format_kb(kb: u64) -> String {
    if kb >= 1_048_576 {
        format!("{:.1}G", kb as f64 / 1_048_576.0)
    } else if kb >= 1024 {
        format!("{:.1}M", kb as f64 / 1024.0)
    } else {
        format!("{kb}K")
    }
}

struct App {
    processes: Vec<ProcessHuge>,
    system: SystemHuge,
    delay: Duration,
    uid_cache: procutils_common::uid::UidCache,
}

impl App {
    fn new(args: &Args) -> Self {
        Self {
            processes: Vec::new(),
            system: SystemHuge::read(),
            delay: Duration::from_secs_f64(args.delay.max(0.5)),
            uid_cache: procutils_common::uid::UidCache::new(),
        }
    }

    fn refresh(&mut self) {
        self.system = SystemHuge::read();
        self.processes.clear();

        let all_procs = match process::all_processes() {
            Ok(iter) => iter,
            Err(_) => return,
        };

        for proc_result in all_procs {
            let proc = match proc_result {
                Ok(p) => p,
                Err(_) => continue,
            };

            let status = match proc.status() {
                Ok(s) => s,
                Err(_) => continue,
            };

            // /proc/[pid]/status reports HugetlbPages in kB; procfs strips
            // " kB" and returns the integer as-is (no /1024 needed). Other
            // procfs fields (Meminfo, smaps) return bytes — those need /1024.
            let hugetlb_kb = status.hugetlbpages.unwrap_or(0);

            // Read AnonHugePages from smaps_rollup if available, fall back to status
            let anon_huge_kb = proc
                .smaps_rollup()
                .ok()
                .and_then(|r| {
                    r.memory_map_rollup.0.first().and_then(|m| {
                        m.extension.map.get("AnonHugePages").copied()
                    })
                })
                .unwrap_or(0)
                / 1024;

            let total = hugetlb_kb + anon_huge_kb;
            if total == 0 {
                continue;
            }

            let stat = match proc.stat() {
                Ok(s) => s,
                Err(_) => continue,
            };

            self.processes.push(ProcessHuge {
                pid: stat.pid,
                user: self.uid_cache.get(status.euid).to_string(),
                comm: stat.comm.clone(),
                hugetlb_kb,
                anon_huge_kb,
                total_huge_kb: total,
            });
        }

        self.processes
            .sort_by(|a, b| b.total_huge_kb.cmp(&a.total_huge_kb));
    }

    fn draw(&self, frame: &mut Frame) {
        let area = frame.area();

        let chunks =
            Layout::vertical([Constraint::Length(4), Constraint::Min(5)])
                .split(area);

        // System summary
        let s = &self.system;
        let summary = vec![
            Line::from(vec![
                Span::styled(" HugePages: ", Style::default().fg(Color::Cyan)),
                Span::raw(format!(
                    "total={}, free={}, rsvd={}, surp={}, size={}K",
                    s.hugepages_total,
                    s.hugepages_free,
                    s.hugepages_rsvd,
                    s.hugepages_surp,
                    s.hugepagesize_kb,
                )),
            ]),
            Line::from(vec![
                Span::styled(
                    " Transparent: ",
                    Style::default().fg(Color::Cyan),
                ),
                Span::raw(format!(
                    "anon={}, shmem={}",
                    format_kb(s.anon_hugepages_kb),
                    format_kb(s.shmem_hugepages_kb),
                )),
                Span::raw(format!("  HugeTLB: {}", format_kb(s.hugetlb_kb))),
            ]),
            Line::from(vec![Span::styled(
                format!(" {} processes using hugepages", self.processes.len()),
                Style::default().fg(Color::White),
            )]),
        ];

        frame.render_widget(
            Paragraph::new(summary)
                .block(Block::default().borders(Borders::BOTTOM)),
            chunks[0],
        );

        // Process table
        let header = Row::new(vec![
            Cell::from("PID")
                .style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("USER")
                .style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("HUGETLB")
                .style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("ANON_HUGE")
                .style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("TOTAL")
                .style(Style::default().add_modifier(Modifier::BOLD)),
            Cell::from("COMMAND")
                .style(Style::default().add_modifier(Modifier::BOLD)),
        ]);

        let rows: Vec<Row> = self
            .processes
            .iter()
            .map(|p| {
                let color = if p.total_huge_kb >= 1_048_576 {
                    Color::Red
                } else if p.total_huge_kb >= 1024 {
                    Color::Yellow
                } else {
                    Color::default()
                };

                Row::new(vec![
                    Cell::from(p.pid.to_string()),
                    Cell::from(p.user.clone()),
                    Cell::from(format_kb(p.hugetlb_kb))
                        .style(Style::default().fg(color)),
                    Cell::from(format_kb(p.anon_huge_kb))
                        .style(Style::default().fg(color)),
                    Cell::from(format_kb(p.total_huge_kb)).style(
                        Style::default().fg(color).add_modifier(Modifier::BOLD),
                    ),
                    Cell::from(p.comm.clone()),
                ])
            })
            .collect();

        let widths = [
            Constraint::Length(8),
            Constraint::Length(12),
            Constraint::Length(10),
            Constraint::Length(10),
            Constraint::Length(10),
            Constraint::Fill(1),
        ];

        let table = Table::new(rows, widths)
            .header(header)
            .block(Block::default().borders(Borders::NONE));

        frame.render_widget(table, chunks[1]);
    }
}

pub fn run(args: Args) -> ExitCode {
    if args.batch {
        return run_batch(&args);
    }

    let mut terminal = match ratatui::try_init() {
        Ok(t) => t,
        Err(e) => {
            eprintln!("hugetop: failed to initialize terminal: {e}");
            return ExitCode::FAILURE;
        }
    };

    let result = run_app(&mut terminal, &args);

    ratatui::restore();

    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(e) => {
            eprintln!("hugetop: {e}");
            ExitCode::FAILURE
        }
    }
}

fn run_batch(args: &Args) -> ExitCode {
    let mut app = App::new(args);
    let mut iteration = 0u64;
    let max = args.iterations.unwrap_or(1);

    loop {
        if iteration >= max {
            break;
        }
        if iteration > 0 {
            std::thread::sleep(app.delay);
        }

        app.refresh();
        print_batch(&app);
        iteration += 1;
    }

    ExitCode::SUCCESS
}

fn print_batch(app: &App) {
    let s = &app.system;
    println!(
        "HugePages: total={}, free={}, rsvd={}, surp={}, size={}K",
        s.hugepages_total,
        s.hugepages_free,
        s.hugepages_rsvd,
        s.hugepages_surp,
        s.hugepagesize_kb,
    );
    println!(
        "Transparent: anon={}, shmem={}",
        format_kb(s.anon_hugepages_kb),
        format_kb(s.shmem_hugepages_kb),
    );
    println!("HugeTLB: {}", format_kb(s.hugetlb_kb));
    println!("{} processes using hugepages", app.processes.len());
    println!();

    println!(
        "{:>7} {:<10} {:>10} {:>10} {:>10} COMMAND",
        "PID", "USER", "HUGETLB", "ANON_HUGE", "TOTAL",
    );
    for p in &app.processes {
        println!(
            "{:>7} {:<10} {:>10} {:>10} {:>10} {}",
            p.pid,
            p.user,
            format_kb(p.hugetlb_kb),
            format_kb(p.anon_huge_kb),
            format_kb(p.total_huge_kb),
            p.comm,
        );
    }
}

fn run_app(
    terminal: &mut DefaultTerminal,
    args: &Args,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut app = App::new(args);

    loop {
        app.refresh();

        terminal.draw(|frame| app.draw(frame))?;

        if event::poll(app.delay)?
            && let Event::Key(key) = event::read()?
            && key.kind == KeyEventKind::Press
            && key.code == KeyCode::Char('q')
        {
            return Ok(());
        }
    }
}