rpki/src/storage/config.rs

use rocksdb::{ColumnFamilyDescriptor, DBCompressionType, Options};

pub const CF_REPOSITORY_VIEW: &str = "repository_view";
pub const CF_RAW_BY_HASH: &str = "raw_by_hash";
pub const CF_RAW_BLOB: &str = "raw_blob";
pub const CF_VCIR: &str = "vcir";
pub const CF_VCIR_FAILED_FETCH_REUSE_IDENTITY: &str = "vcir_failed_fetch_reuse_identity";
pub const CF_MANIFEST_REPLAY_META: &str = "manifest_replay_meta";
pub const CF_ROA_CACHE_PROJECTION: &str = "roa_cache_projection";
pub const CF_PUBLICATION_POINT_CACHE_PROJECTION: &str = "publication_point_cache_projection";
pub const CF_CHILD_CERTIFICATE_CACHE_PROJECTION: &str = "child_certificate_cache_projection";
pub const CF_RRDP_SOURCE: &str = "rrdp_source";
pub const CF_RRDP_SOURCE_MEMBER: &str = "rrdp_source_member";
pub const CF_RRDP_URI_OWNER: &str = "rrdp_uri_owner";
pub const CF_TRANSPORT_PREFETCH: &str = "transport_prefetch";

pub const ALL_COLUMN_FAMILY_NAMES: &[&str] = &[
    CF_REPOSITORY_VIEW,
    CF_RAW_BY_HASH,
    CF_RAW_BLOB,
    CF_VCIR,
    CF_VCIR_FAILED_FETCH_REUSE_IDENTITY,
    CF_MANIFEST_REPLAY_META,
    CF_ROA_CACHE_PROJECTION,
    CF_PUBLICATION_POINT_CACHE_PROJECTION,
    CF_CHILD_CERTIFICATE_CACHE_PROJECTION,
    CF_RRDP_SOURCE,
    CF_RRDP_SOURCE_MEMBER,
    CF_RRDP_URI_OWNER,
    CF_TRANSPORT_PREFETCH,
];

pub(super) const REPOSITORY_VIEW_KEY_PREFIX: &str = "repo_view:";
pub(super) const RAW_BY_HASH_KEY_PREFIX: &str = "rawbyhash:";
pub(super) const RAW_BLOB_KEY_PREFIX: &str = "rawblob:";
pub(super) const VCIR_KEY_PREFIX: &str = "vcir:";
pub(super) const VCIR_FAILED_FETCH_REUSE_IDENTITY_KEY_PREFIX: &str =
    "vcir_failed_fetch_reuse_identity:";
pub(super) const MANIFEST_REPLAY_META_KEY_PREFIX: &str = "manifest_replay_meta:";
pub(super) const ROA_CACHE_PROJECTION_KEY_PREFIX: &str = "roa_cache_projection:";
pub(super) const PUBLICATION_POINT_CACHE_PROJECTION_KEY_PREFIX: &str =
    "publication_point_cache_projection:";
pub(super) const CHILD_CERTIFICATE_CACHE_PROJECTION_KEY_PREFIX: &str =
    "child_certificate_cache_projection:";
pub(super) const RRDP_SOURCE_KEY_PREFIX: &str = "rrdp_source:";
pub(super) const RRDP_SOURCE_MEMBER_KEY_PREFIX: &str = "rrdp_source_member:";
pub(super) const RRDP_URI_OWNER_KEY_PREFIX: &str = "rrdp_uri_owner:";
pub(super) const TRANSPORT_PREFETCH_LAST_SNAPSHOT_KEY: &str = "transport_prefetch:last_snapshot";

const WORK_DB_BLOB_MODE_ENV: &str = "RPKI_WORK_DB_BLOB_MODE";
const WORK_DB_MEMORY_PROFILE_ENV: &str = "RPKI_WORK_DB_MEMORY_PROFILE";

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum WorkDbBlobMode {
    Current,
    Disabled,
    Lz4,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum WorkDbMemoryProfile {
    Default,
    Compact,
}

pub(super) fn parse_work_db_blob_mode(raw: &str) -> Option<WorkDbBlobMode> {
    match raw.trim().to_ascii_lowercase().as_str() {
        "" | "default" => Some(default_work_db_blob_mode()),
        "current" | "legacy" => Some(WorkDbBlobMode::Current),
        "disabled" | "disable" | "off" | "none" | "no_blob" | "no-blob" => {
            Some(WorkDbBlobMode::Disabled)
        }
        "lz4" | "blob_lz4" | "blob-lz4" => Some(WorkDbBlobMode::Lz4),
        _ => None,
    }
}

pub(super) fn default_work_db_blob_mode() -> WorkDbBlobMode {
    WorkDbBlobMode::Disabled
}

pub(super) fn work_db_blob_mode_from_env() -> WorkDbBlobMode {
    let Ok(raw) = std::env::var(WORK_DB_BLOB_MODE_ENV) else {
        return default_work_db_blob_mode();
    };
    match parse_work_db_blob_mode(&raw) {
        Some(mode) => mode,
        None => {
            eprintln!(
                "warning: unsupported {WORK_DB_BLOB_MODE_ENV}={raw:?}; using default work-db blobdb mode"
            );
            default_work_db_blob_mode()
        }
    }
}

pub(super) fn parse_work_db_memory_profile(raw: &str) -> Option<WorkDbMemoryProfile> {
    match raw.trim().to_ascii_lowercase().as_str() {
        "" | "default" | "off" | "none" => Some(WorkDbMemoryProfile::Default),
        "compact" | "low" | "low_memory" | "low-memory" => Some(WorkDbMemoryProfile::Compact),
        _ => None,
    }
}

pub(super) fn work_db_memory_profile_from_env() -> WorkDbMemoryProfile {
    let Ok(raw) = std::env::var(WORK_DB_MEMORY_PROFILE_ENV) else {
        return WorkDbMemoryProfile::Default;
    };
    match parse_work_db_memory_profile(&raw) {
        Some(profile) => profile,
        None => {
            eprintln!(
                "warning: unsupported {WORK_DB_MEMORY_PROFILE_ENV}={raw:?}; using default work-db memory profile"
            );
            WorkDbMemoryProfile::Default
        }
    }
}

pub(super) fn configure_work_db_options(
    opts: &mut Options,
    blob_mode: WorkDbBlobMode,
    memory_profile: WorkDbMemoryProfile,
) {
    opts.set_compression_type(DBCompressionType::Lz4);
    apply_work_db_memory_profile(opts, memory_profile);
    match blob_mode {
        WorkDbBlobMode::Current => enable_blobdb_current(opts),
        WorkDbBlobMode::Disabled => {}
        WorkDbBlobMode::Lz4 => {
            enable_blobdb_current(opts);
            opts.set_blob_compression_type(DBCompressionType::Lz4);
        }
    }
}

pub(super) fn cf_opts(blob_mode: WorkDbBlobMode, memory_profile: WorkDbMemoryProfile) -> Options {
    let mut opts = Options::default();
    configure_work_db_options(&mut opts, blob_mode, memory_profile);
    opts
}

pub fn column_family_descriptors() -> Vec<ColumnFamilyDescriptor> {
    column_family_descriptors_for_blob_mode(work_db_blob_mode_from_env())
}

pub(super) fn column_family_descriptors_for_blob_mode(
    blob_mode: WorkDbBlobMode,
) -> Vec<ColumnFamilyDescriptor> {
    let memory_profile = work_db_memory_profile_from_env();
    ALL_COLUMN_FAMILY_NAMES
        .iter()
        .map(|name| ColumnFamilyDescriptor::new(*name, cf_opts(blob_mode, memory_profile)))
        .collect()
}

fn apply_work_db_memory_profile(opts: &mut Options, memory_profile: WorkDbMemoryProfile) {
    match memory_profile {
        WorkDbMemoryProfile::Default => {}
        WorkDbMemoryProfile::Compact => {
            opts.set_write_buffer_size(16 * 1024 * 1024);
            opts.set_max_write_buffer_number(1);
        }
    }
}

pub(super) fn enable_blobdb_current(opts: &mut Options) {
    #[allow(unused_mut)]
    let mut _enabled = false;

    #[allow(dead_code)]
    fn _set(opts: &mut Options) {
        opts.set_enable_blob_files(true);
        opts.set_min_blob_size(1024);
    }

    _set(opts);
}