panda-rpki-oss/src/repository/storage/memory.rs

//! RocksDB and process-memory observation.
//!
//! This module owns the serialization-friendly memory snapshot model and the
//! mapping from RocksDB property names to that model. Storage mutation stays
//! in the parent module.

use rocksdb::DB;
use serde::Serialize;

const ROCKSDB_MEMORY_PROPERTY_NAMES: &[(&str, &str)] = &[
    ("cur_size_all_mem_tables", "rocksdb.cur-size-all-mem-tables"),
    ("size_all_mem_tables", "rocksdb.size-all-mem-tables"),
    (
        "estimate_table_readers_mem",
        "rocksdb.estimate-table-readers-mem",
    ),
    ("block_cache_capacity", "rocksdb.block-cache-capacity"),
    ("block_cache_usage", "rocksdb.block-cache-usage"),
    (
        "block_cache_pinned_usage",
        "rocksdb.block-cache-pinned-usage",
    ),
    ("num_snapshots", "rocksdb.num-snapshots"),
    ("background_errors", "rocksdb.background-errors"),
];

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct RocksDbMemoryProperties {
    pub cur_size_all_mem_tables: Option<u64>,
    pub size_all_mem_tables: Option<u64>,
    pub estimate_table_readers_mem: Option<u64>,
    pub block_cache_capacity: Option<u64>,
    pub block_cache_usage: Option<u64>,
    pub block_cache_pinned_usage: Option<u64>,
    pub num_snapshots: Option<u64>,
    pub background_errors: Option<u64>,
    pub errors: Vec<String>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct RocksDbColumnFamilyMemoryProperties {
    pub name: String,
    pub properties: RocksDbMemoryProperties,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct RocksDbMemoryDbSnapshot {
    pub label: String,
    pub properties: RocksDbMemoryProperties,
    pub column_families: Vec<RocksDbColumnFamilyMemoryProperties>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct RocksDbMemoryTotals {
    pub cur_size_all_mem_tables: u64,
    pub size_all_mem_tables: u64,
    pub estimate_table_readers_mem: u64,
    pub block_cache_capacity: u64,
    pub block_cache_usage: u64,
    pub block_cache_pinned_usage: u64,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct RocksDbMemorySnapshot {
    pub databases: Vec<RocksDbMemoryDbSnapshot>,
    pub totals: RocksDbMemoryTotals,
}

impl RocksDbMemoryTotals {
    pub(super) fn add_properties(&mut self, properties: &RocksDbMemoryProperties) {
        self.cur_size_all_mem_tables += properties.cur_size_all_mem_tables.unwrap_or(0);
        self.size_all_mem_tables += properties.size_all_mem_tables.unwrap_or(0);
        self.estimate_table_readers_mem += properties.estimate_table_readers_mem.unwrap_or(0);
        self.block_cache_capacity += properties.block_cache_capacity.unwrap_or(0);
        self.block_cache_usage += properties.block_cache_usage.unwrap_or(0);
        self.block_cache_pinned_usage += properties.block_cache_pinned_usage.unwrap_or(0);
    }
}

fn set_memory_property(properties: &mut RocksDbMemoryProperties, name: &str, value: u64) {
    match name {
        "cur_size_all_mem_tables" => properties.cur_size_all_mem_tables = Some(value),
        "size_all_mem_tables" => properties.size_all_mem_tables = Some(value),
        "estimate_table_readers_mem" => properties.estimate_table_readers_mem = Some(value),
        "block_cache_capacity" => properties.block_cache_capacity = Some(value),
        "block_cache_usage" => properties.block_cache_usage = Some(value),
        "block_cache_pinned_usage" => properties.block_cache_pinned_usage = Some(value),
        "num_snapshots" => properties.num_snapshots = Some(value),
        "background_errors" => properties.background_errors = Some(value),
        _ => {}
    }
}

fn parse_rocksdb_property_int(raw: Option<String>) -> Option<u64> {
    raw.and_then(|value| value.trim().parse::<u64>().ok())
}

fn memory_properties_for_db(db: &DB) -> RocksDbMemoryProperties {
    let mut properties = RocksDbMemoryProperties::default();
    for (field_name, property_name) in ROCKSDB_MEMORY_PROPERTY_NAMES {
        match db.property_value(*property_name) {
            Ok(value) => {
                if let Some(parsed) = parse_rocksdb_property_int(value) {
                    set_memory_property(&mut properties, field_name, parsed);
                }
            }
            Err(err) => properties.errors.push(format!("{property_name}: {}", err)),
        }
    }
    properties
}

fn memory_properties_for_cf(db: &DB, cf_name: &'static str) -> RocksDbColumnFamilyMemoryProperties {
    let mut properties = RocksDbMemoryProperties::default();
    let Some(cf) = db.cf_handle(cf_name) else {
        properties
            .errors
            .push(format!("missing column family: {cf_name}"));
        return RocksDbColumnFamilyMemoryProperties {
            name: cf_name.to_string(),
            properties,
        };
    };
    for (field_name, property_name) in ROCKSDB_MEMORY_PROPERTY_NAMES {
        match db.property_value_cf(cf, *property_name) {
            Ok(value) => {
                if let Some(parsed) = parse_rocksdb_property_int(value) {
                    set_memory_property(&mut properties, field_name, parsed);
                }
            }
            Err(err) => properties.errors.push(format!("{property_name}: {}", err)),
        }
    }
    RocksDbColumnFamilyMemoryProperties {
        name: cf_name.to_string(),
        properties,
    }
}

pub(crate) fn memory_db_snapshot_for_column_families(
    label: impl Into<String>,
    db: &DB,
    column_families: Option<&[&'static str]>,
) -> RocksDbMemoryDbSnapshot {
    RocksDbMemoryDbSnapshot {
        label: label.into(),
        properties: memory_properties_for_db(db),
        column_families: column_families
            .map(|names| {
                names
                    .iter()
                    .map(|name| memory_properties_for_cf(db, name))
                    .collect()
            })
            .unwrap_or_default(),
    }
}

#[cfg(test)]
mod tests {
    use super::{
        RocksDbMemoryProperties, RocksDbMemoryTotals, parse_rocksdb_property_int,
        set_memory_property,
    };

    #[test]
    fn parses_rocksdb_integer_properties_without_panicking() {
        assert_eq!(
            parse_rocksdb_property_int(Some(" 42 ".to_string())),
            Some(42)
        );
        assert_eq!(
            parse_rocksdb_property_int(Some("not-a-number".to_string())),
            None
        );
        assert_eq!(parse_rocksdb_property_int(None), None);
    }

    #[test]
    fn aggregates_only_known_memory_properties() {
        let mut properties = RocksDbMemoryProperties::default();
        set_memory_property(&mut properties, "block_cache_usage", 11);
        set_memory_property(&mut properties, "unknown", 99);

        let mut totals = RocksDbMemoryTotals::default();
        totals.add_properties(&properties);

        assert_eq!(totals.block_cache_usage, 11);
        assert_eq!(totals.block_cache_capacity, 0);
    }
}