rpki/src/memory_telemetry.rs

use serde::Serialize;

use crate::storage::RocksDbMemorySnapshot;

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct MallocTrimProbe {
    pub supported: bool,
    pub return_value: Option<i32>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct ProcessMemorySnapshot {
    pub label: String,
    pub vm_rss_kb: Option<u64>,
    pub vm_size_kb: Option<u64>,
    pub vm_data_kb: Option<u64>,
    pub vm_swap_kb: Option<u64>,
    pub rss_anon_kb: Option<u64>,
    pub rss_file_kb: Option<u64>,
    pub rss_shmem_kb: Option<u64>,
    pub threads: Option<u64>,
    pub fd_count: Option<u64>,
    pub smaps_rollup: Option<SmapsRollupSnapshot>,
    pub smaps_mapping_summary: Option<SmapsMappingSummary>,
    pub errors: Vec<String>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct SmapsRollupSnapshot {
    pub rss_kb: Option<u64>,
    pub pss_kb: Option<u64>,
    pub shared_clean_kb: Option<u64>,
    pub shared_dirty_kb: Option<u64>,
    pub private_clean_kb: Option<u64>,
    pub private_dirty_kb: Option<u64>,
    pub anonymous_kb: Option<u64>,
    pub swap_kb: Option<u64>,
    pub swap_pss_kb: Option<u64>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct SmapsMappingSummary {
    pub heap: SmapsMappingCategory,
    pub anonymous_mmap: SmapsMappingCategory,
    pub file_backed: SmapsMappingCategory,
    pub stack: SmapsMappingCategory,
    pub special: SmapsMappingCategory,
    pub total: SmapsMappingCategory,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct SmapsMappingCategory {
    pub mappings: u64,
    pub size_kb: u64,
    pub rss_kb: u64,
    pub pss_kb: u64,
    pub private_clean_kb: u64,
    pub private_dirty_kb: u64,
    pub anonymous_kb: u64,
    pub largest_mapping_rss_kb: u64,
    pub large_mapping_count_64m: u64,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct MemoryTelemetryCheckpoint {
    pub label: String,
    pub elapsed_ms: u64,
    pub process: ProcessMemorySnapshot,
    pub rocksdb: RocksDbMemorySnapshot,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct MemoryTelemetrySummary {
    pub checkpoints: Vec<MemoryTelemetryCheckpoint>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub object_graph: Option<ObjectGraphMemorySummary>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub malloc_trim_probes: Vec<MallocTrimProbe>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct ObjectGraphMemorySummary {
    pub captured_at_label: String,
    pub total_estimated_bytes: u64,
    pub sections: Vec<ObjectGraphMemorySection>,
    pub notes: Vec<String>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct ObjectGraphMemorySection {
    pub name: String,
    pub item_count: u64,
    pub shallow_bytes: u64,
    pub heap_bytes: u64,
    pub estimated_bytes: u64,
    pub string_count: u64,
    pub string_bytes: u64,
    pub string_capacity_bytes: u64,
    pub vec_count: u64,
    pub vec_heap_bytes: u64,
    pub vec_capacity_bytes: u64,
    pub details: Vec<ObjectGraphMemoryMetric>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct ObjectGraphMemoryMetric {
    pub name: String,
    pub value: u64,
}

pub fn process_memory_snapshot(label: impl Into<String>) -> ProcessMemorySnapshot {
    let label = label.into();
    let mut snapshot = ProcessMemorySnapshot {
        label,
        vm_rss_kb: None,
        vm_size_kb: None,
        vm_data_kb: None,
        vm_swap_kb: None,
        rss_anon_kb: None,
        rss_file_kb: None,
        rss_shmem_kb: None,
        threads: None,
        fd_count: current_fd_count(),
        smaps_rollup: None,
        smaps_mapping_summary: None,
        errors: Vec::new(),
    };

    match std::fs::read_to_string("/proc/self/status") {
        Ok(status) => parse_status(&status, &mut snapshot),
        Err(err) => snapshot
            .errors
            .push(format!("read /proc/self/status failed: {err}")),
    }

    match std::fs::read_to_string("/proc/self/smaps_rollup") {
        Ok(smaps) => snapshot.smaps_rollup = Some(parse_smaps_rollup(&smaps)),
        Err(err) => snapshot
            .errors
            .push(format!("read /proc/self/smaps_rollup failed: {err}")),
    }

    match std::fs::read_to_string("/proc/self/smaps") {
        Ok(smaps) => snapshot.smaps_mapping_summary = Some(parse_smaps_mapping_summary(&smaps)),
        Err(err) => snapshot
            .errors
            .push(format!("read /proc/self/smaps failed: {err}")),
    }

    snapshot
}

pub fn malloc_trim_probe() -> MallocTrimProbe {
    #[cfg(all(target_os = "linux", target_env = "gnu"))]
    {
        MallocTrimProbe {
            supported: true,
            return_value: Some(unsafe { malloc_trim(0) }),
        }
    }
    #[cfg(not(all(target_os = "linux", target_env = "gnu")))]
    {
        MallocTrimProbe {
            supported: false,
            return_value: None,
        }
    }
}

#[cfg(all(target_os = "linux", target_env = "gnu"))]
unsafe extern "C" {
    fn malloc_trim(pad: usize) -> i32;
}

fn current_fd_count() -> Option<u64> {
    std::fs::read_dir("/proc/self/fd")
        .ok()
        .map(|entries| entries.filter_map(Result::ok).count() as u64)
}

fn parse_status(status: &str, snapshot: &mut ProcessMemorySnapshot) {
    for line in status.lines() {
        let Some((key, value)) = line.split_once(':') else {
            continue;
        };
        let parsed = parse_kb_or_plain_u64(value);
        match key {
            "VmRSS" => snapshot.vm_rss_kb = parsed,
            "VmSize" => snapshot.vm_size_kb = parsed,
            "VmData" => snapshot.vm_data_kb = parsed,
            "VmSwap" => snapshot.vm_swap_kb = parsed,
            "RssAnon" => snapshot.rss_anon_kb = parsed,
            "RssFile" => snapshot.rss_file_kb = parsed,
            "RssShmem" => snapshot.rss_shmem_kb = parsed,
            "Threads" => snapshot.threads = parsed,
            _ => {}
        }
    }
}

fn parse_smaps_rollup(smaps: &str) -> SmapsRollupSnapshot {
    let mut snapshot = SmapsRollupSnapshot::default();
    for line in smaps.lines() {
        let Some((key, value)) = line.split_once(':') else {
            continue;
        };
        let parsed = parse_kb_or_plain_u64(value);
        match key {
            "Rss" => snapshot.rss_kb = parsed,
            "Pss" => snapshot.pss_kb = parsed,
            "Shared_Clean" => snapshot.shared_clean_kb = parsed,
            "Shared_Dirty" => snapshot.shared_dirty_kb = parsed,
            "Private_Clean" => snapshot.private_clean_kb = parsed,
            "Private_Dirty" => snapshot.private_dirty_kb = parsed,
            "Anonymous" => snapshot.anonymous_kb = parsed,
            "Swap" => snapshot.swap_kb = parsed,
            "SwapPss" => snapshot.swap_pss_kb = parsed,
            _ => {}
        }
    }
    snapshot
}

fn parse_smaps_mapping_summary(smaps: &str) -> SmapsMappingSummary {
    let mut summary = SmapsMappingSummary::default();
    let mut current_path = String::new();
    let mut current = SmapsMappingCategory::default();
    let mut have_mapping = false;

    for line in smaps.lines() {
        if is_smaps_mapping_header(line) {
            if have_mapping {
                add_mapping(&mut summary, &current_path, &current);
            }
            current_path = smaps_header_path(line);
            current = SmapsMappingCategory {
                mappings: 1,
                ..SmapsMappingCategory::default()
            };
            have_mapping = true;
            continue;
        }

        if !have_mapping {
            continue;
        }

        let Some((key, value)) = line.split_once(':') else {
            continue;
        };
        let parsed = parse_kb_or_plain_u64(value).unwrap_or(0);
        match key {
            "Size" => current.size_kb = parsed,
            "Rss" => current.rss_kb = parsed,
            "Pss" => current.pss_kb = parsed,
            "Private_Clean" => current.private_clean_kb = parsed,
            "Private_Dirty" => current.private_dirty_kb = parsed,
            "Anonymous" => current.anonymous_kb = parsed,
            _ => {}
        }
    }

    if have_mapping {
        add_mapping(&mut summary, &current_path, &current);
    }

    summary
}

fn is_smaps_mapping_header(line: &str) -> bool {
    let mut parts = line.split_whitespace();
    let Some(range) = parts.next() else {
        return false;
    };
    let Some(perms) = parts.next() else {
        return false;
    };
    let Some((start, end)) = range.split_once('-') else {
        return false;
    };
    !start.is_empty()
        && !end.is_empty()
        && start.as_bytes().iter().all(u8::is_ascii_hexdigit)
        && end.as_bytes().iter().all(u8::is_ascii_hexdigit)
        && perms.len() == 4
        && perms
            .as_bytes()
            .iter()
            .all(|b| matches!(b, b'r' | b'w' | b'x' | b's' | b'p' | b'-'))
}

fn smaps_header_path(line: &str) -> String {
    line.split_whitespace()
        .skip(5)
        .collect::<Vec<_>>()
        .join(" ")
}

fn add_mapping(summary: &mut SmapsMappingSummary, path: &str, mapping: &SmapsMappingCategory) {
    add_category(&mut summary.total, mapping);
    match mapping_category(path) {
        MappingCategory::Heap => add_category(&mut summary.heap, mapping),
        MappingCategory::AnonymousMmap => add_category(&mut summary.anonymous_mmap, mapping),
        MappingCategory::FileBacked => add_category(&mut summary.file_backed, mapping),
        MappingCategory::Stack => add_category(&mut summary.stack, mapping),
        MappingCategory::Special => add_category(&mut summary.special, mapping),
    }
}

fn add_category(target: &mut SmapsMappingCategory, source: &SmapsMappingCategory) {
    target.mappings += source.mappings;
    target.size_kb += source.size_kb;
    target.rss_kb += source.rss_kb;
    target.pss_kb += source.pss_kb;
    target.private_clean_kb += source.private_clean_kb;
    target.private_dirty_kb += source.private_dirty_kb;
    target.anonymous_kb += source.anonymous_kb;
    target.largest_mapping_rss_kb = target.largest_mapping_rss_kb.max(source.rss_kb);
    if source.rss_kb >= 64 * 1024 {
        target.large_mapping_count_64m += source.mappings;
    }
}

enum MappingCategory {
    Heap,
    AnonymousMmap,
    FileBacked,
    Stack,
    Special,
}

fn mapping_category(path: &str) -> MappingCategory {
    if path == "[heap]" {
        MappingCategory::Heap
    } else if path.starts_with("[stack") {
        MappingCategory::Stack
    } else if path.is_empty() {
        MappingCategory::AnonymousMmap
    } else if path.starts_with('/') {
        MappingCategory::FileBacked
    } else {
        MappingCategory::Special
    }
}

fn parse_kb_or_plain_u64(value: &str) -> Option<u64> {
    value.split_whitespace().next()?.parse::<u64>().ok()
}