libamdgpu_top 0.11.4

use std::collections::{HashMap, HashSet};
use std::fs;
use std::io::Read;
use std::time::Duration;
use std::path::PathBuf;

use super::XdnaFdInfoUsage;
use crate::stat::ProcInfo;

#[derive(Clone, Debug, Default, Eq, PartialEq, PartialOrd)]
pub struct XdnaProcUsage {
    pub pid: i32,
    pub name: String,
    pub ids_count: usize,
    pub usage: XdnaFdInfoUsage,
    // pub cpu_usage: i64, // %
}

#[derive(Clone, Default)]
pub struct XdnaFdInfoStat {
    pub pre_proc_usage_map: HashMap<i32, XdnaFdInfoUsage>,
    pub drm_client_ids: HashSet<usize>,
    pub proc_usage: Vec<XdnaProcUsage>,
    pub interval: Duration,
}

impl XdnaFdInfoStat {
    pub fn get_proc_usage(&mut self, proc_info: &ProcInfo) {
        let pid = proc_info.pid;
        let name = &proc_info.name;
        let mut stat = XdnaFdInfoUsage::default();
        let mut buf = String::new();
        let mut ids_count = 0usize;
        let mut path = PathBuf::with_capacity(24);

        path.push("/proc");
        path.push(pid.to_string());

        if !path.exists() {
            self.pre_proc_usage_map.remove(&pid);
            return;
        }

        path.push("fdinfo");

        for fd in &proc_info.fds {
            buf.clear();

            {
                path.push(fd.to_string());
                let Ok(mut f) = fs::File::open(&path) else { continue };
                if f.read_to_string(&mut buf).is_err() { continue }
                path.pop(); // fd
            }

            let mut lines = buf.lines().skip_while(|l| !l.starts_with("drm-client-id"));
            if let Some(id) = lines.next().and_then(XdnaFdInfoUsage::id_parse) {
                ids_count += 1;
                if !self.drm_client_ids.insert(id) { continue }
            } else {
                continue;
            }

            for l in lines {
                let Some(s) = l.get(0..13) else { continue };

                match s {
                    "drm-total-mem" => stat.total_memory_usage_parse(l),
                    "drm-shared-me" => stat.shared_memory_usage_parse(l),
                    "drm-active-me" => stat.active_memory_usage_parse(l),
                    "drm-engine-np" => stat.engine_usage_parse(l),
                    _ => {},
                }
            }
        }

        let diff = if let Some(pre_stat) = self.pre_proc_usage_map.get_mut(&pid) {
            let tmp = stat.calc_usage(pre_stat, &self.interval);
            *pre_stat = stat;

            tmp
        } else {
            let usage = XdnaFdInfoUsage {
                total_memory: stat.total_memory,
                shared_memory: stat.shared_memory,
                active_memory: stat.active_memory,
                ..Default::default()
            };

            self.pre_proc_usage_map.insert(pid, stat);

            usage
        };

        self.proc_usage.push(XdnaProcUsage {
            pid,
            name: name.to_string(),
            ids_count,
            usage: diff,
        });
    }

    pub fn update_proc_usage(&mut self, proc_index: &[ProcInfo]) {
        self.proc_usage.clear();
        self.drm_client_ids.clear();

        // If a process disappears from proc_index while its usage remains in pre_proc_usage_map,
        // it will persist in pre_proc_usage_map and needs to be removed
        if proc_index.len() < self.pre_proc_usage_map.len() {
            let mut path = PathBuf::with_capacity(16);
            path.push("/proc");
            self.pre_proc_usage_map.retain(|&pid, _| {
                path.push(pid.to_string());

                let b = path.exists();

                path.pop(); // pid

                b
            });
        }

        for pu in proc_index {
            self.get_proc_usage(pu);
        }
    }

    pub fn fold_fdinfo_usage(&self) -> XdnaFdInfoUsage {
        self.proc_usage.iter().fold(XdnaFdInfoUsage::default(), |acc, pu| acc + pu.usage)
    }
}