rs_malloc_tracker 1.0.0

use std::{
    fs,
    sync::{Arc, RwLock},
    time::{Duration, Instant},
};

use nodit::{Interval, NoditMap, interval::ie};
use rustc_hash::FxHashMap;

#[cfg(feature = "mmap")]
use regex::Regex;

use super::meta_instrumentation::MetaInstrumentationState;

#[derive(Clone)]
struct Map {
    // Represents a map in /proc/self/maps
    // Each Map contains _resolved_ symbols for that region

    // name as reported by /proc/self/maps (sometimes empty)
    name: String,

    // resolved symbol cache
    symbols: FxHashMap<usize, Option<String>>,
}

pub struct SymbolCache {
    meta_instrumentation: Arc<RwLock<MetaInstrumentationState>>,
    last_proc_maps_read: Instant,
    cache: NoditMap<usize, Interval<usize>, Map>,
}

#[derive(Default)]
pub struct Symbol<'a> {
    pub region: Option<&'a String>,
    pub function: Option<&'a String>,
}

impl SymbolCache {
    pub fn new(module: &super::PrometheusModule) -> Result<Self, String> {
        let mut c = SymbolCache {
            meta_instrumentation: Arc::clone(module.meta_instrumentation.as_ref().unwrap()),
            last_proc_maps_read: Instant::now(),
            cache: NoditMap::new(),
        };

        c.read_proc_maps()?;

        Ok(c)
    }

    /// Read and parse /proc/self/maps so a symbol cache is created for each region
    pub fn read_proc_maps(&mut self) -> Result<(), String> {
        self.meta_instrumentation.read().unwrap().meta_instrument(
            "symbol-cache-read-proc-maps",
            || -> Result<_, String> {
                read_self_maps(&mut |start, end, name| {
                    if !self.cache.contains_interval(ie(start, end)) {
                        let _ = self
                            .cache
                            .insert_overwrite(
                                ie(start, end),
                                Map {
                                    name: name.to_string(),
                                    symbols: FxHashMap::default(),
                                },
                            )
                            .collect::<Vec<_>>();
                    }
                })?;

                Ok(())
            },
        )
    }

    fn maybe_read_proc_maps(&mut self) -> Result<(), String> {
        if self.last_proc_maps_read.elapsed() > Duration::from_secs(120) {
            self.read_proc_maps()?;
        }

        self.last_proc_maps_read = Instant::now();

        Ok(())
    }

    pub fn resolve_symbol_at<'a>(
        &'a mut self,
        resolve_function: bool,
        ip: usize,
    ) -> Result<Symbol<'a>, String> {
        self.maybe_read_proc_maps()?;

        // get zone from /proc/self/maps
        if let Some(v) = self.cache.get_at_point_mut(ip) {
            Ok(Symbol {
                region: Some(&v.name),

                // try to resolve symbol
                function: v
                    .symbols
                    .entry(ip)
                    .or_insert_with(|| {
                        let mut function: Option<String> = None;

                        if resolve_function {
                            backtrace::resolve(ip as *mut _, |sym| {
                                function = sym
                                    .name()
                                    .map(|s| s.as_str().unwrap_or("<error>").to_string());
                            })
                        }

                        function
                    })
                    .as_ref(),
            })
        } else {
            Ok(Symbol::default())
        }
    }
}

pub fn read_self_maps<F>(cb: &mut F) -> Result<(), String>
where
    F: for<'a> FnMut(usize, usize, &'a str),
{
    let _ = fs::read_to_string("/proc/self/maps")
        .map_err(|e| format!("Failed to read /proc/self/maps: {}", e))?
        .as_str()
        .split_terminator("\n")
        .map(|line| {
            let s: Vec<&str> = line.split(" ").collect();

            let (start, end) = s[0].split_once("-").unwrap();
            let start = usize::from_str_radix(start, 16).unwrap();
            let end = usize::from_str_radix(end, 16).unwrap();

            let name = *s.last().unwrap();

            cb(start, end, name);
        })
        .collect::<Vec<_>>();
    Ok(())
}

#[cfg(feature = "mmap")]
#[derive(Default, Debug)]
pub struct Maps {
    pub owner: Option<String>,
    pub size: usize,
    pub rss: usize,
}

#[cfg(feature = "mmap")]
pub fn read_self_smaps() -> Result<NoditMap<usize, Interval<usize>, Maps>, String> {
    let head_re = Regex::new(
        r"^([a-f0-9]+)\-([a-f0-9]+)\s+(\S+)\s+[a-z0-9]+\s+\S+[a-z0-9]+\s+[a-z0-9]+\s*(\S*)$",
    )
    .unwrap();
    let val_re = Regex::new(r"^(\S+):\s*(\S+) kB$").unwrap();

    let mut maps = NoditMap::new();

    let mut cur_interval = ie(0, 1);
    let mut cur_map = Maps::default();

    for line in fs::read_to_string("/proc/self/smaps")
        .map_err(|e| format!("Failed to read /proc/self/smaps: {}", e))?
        .split_terminator("\n")
    {
        if let Some((_, [start, end, _perms, owner])) = head_re.captures(line).map(|c| c.extract())
        {
            maps.insert_strict(cur_interval, cur_map).unwrap();
            cur_interval = ie(
                usize::from_str_radix(start, 16).unwrap(),
                usize::from_str_radix(end, 16).unwrap(),
            );
            cur_map = Maps::default();
            if !owner.is_empty() {
                cur_map.owner = Some(owner.to_string());
            }
        } else if let Some((_, [name, size])) = val_re.captures(line).map(|c| c.extract()) {
            match name {
                "Size" => {
                    cur_map.size = size.parse::<usize>().unwrap() * 1024;
                }
                "Rss" => {
                    cur_map.rss = size.parse::<usize>().unwrap() * 1024;
                }
                _ => (),
            }
        }
    }
    maps.insert_strict(cur_interval, cur_map).unwrap();

    Ok(maps)
}