agave-snapshots 4.2.0

use {
    crate::{
        ArchiveFormat, Result,
        error::SnapshotError,
        snapshot_archive_info::{
            FullSnapshotArchiveInfo, IncrementalSnapshotArchiveInfo, SnapshotArchiveInfoGetter as _,
        },
        snapshot_hash::SnapshotHash,
    },
    log::*,
    regex::Regex,
    solana_clock::Slot,
    solana_hash::Hash,
    std::{
        fs,
        path::{Path, PathBuf},
        sync::LazyLock,
    },
};

pub const SNAPSHOT_STATUS_CACHE_FILENAME: &str = "status_cache";
pub const SNAPSHOT_VERSION_FILENAME: &str = "version";
pub const SNAPSHOT_FASTBOOT_VERSION_FILENAME: &str = "fastboot_version";
pub const SNAPSHOT_ACCOUNTS_HARDLINKS: &str = "accounts_hardlinks";
pub const SNAPSHOT_STORAGES_LIST_FILENAME: &str = "storages_list";
pub const SNAPSHOT_ARCHIVE_DOWNLOAD_DIR: &str = "remote";
pub const SNAPSHOT_OBSOLETE_ACCOUNTS_FILENAME: &str = "obsolete_accounts";
/// When a snapshot is taken of a bank, the state is serialized under this directory.
/// Specifically in `BANK_SNAPSHOTS_DIR/SLOT/`.
/// This is also where the bank state is located in the snapshot archive.
pub const BANK_SNAPSHOTS_DIR: &str = "snapshots";
pub const TMP_SNAPSHOT_ARCHIVE_PREFIX: &str = "tmp-snapshot-archive-";
pub const FULL_SNAPSHOT_ARCHIVE_FILENAME_REGEX: &str =
    r"^snapshot-(?P<slot>[[:digit:]]+)-(?P<hash>[[:alnum:]]+)\.(?P<ext>tar\.zst|tar\.lz4)$";
pub const INCREMENTAL_SNAPSHOT_ARCHIVE_FILENAME_REGEX: &str = r"^incremental-snapshot-(?P<base>[[:digit:]]+)-(?P<slot>[[:digit:]]+)-(?P<hash>[[:alnum:]]+)\.(?P<ext>tar\.zst|tar\.lz4)$";

/// Get the `&str` from a `&Path`
pub fn path_to_file_name_str(path: &Path) -> Result<&str> {
    path.file_name()
        .ok_or_else(|| SnapshotError::PathToFileNameError(path.to_path_buf()))?
        .to_str()
        .ok_or_else(|| SnapshotError::FileNameToStrError(path.to_path_buf()))
}

pub fn build_snapshot_archives_remote_dir(snapshot_archives_dir: impl AsRef<Path>) -> PathBuf {
    snapshot_archives_dir
        .as_ref()
        .join(SNAPSHOT_ARCHIVE_DOWNLOAD_DIR)
}

/// Build the full snapshot archive path from its components: the snapshot archives directory, the
/// snapshot slot, the accounts hash, and the archive format.
pub fn build_full_snapshot_archive_path(
    full_snapshot_archives_dir: impl AsRef<Path>,
    slot: Slot,
    hash: &SnapshotHash,
    archive_format: ArchiveFormat,
) -> PathBuf {
    full_snapshot_archives_dir.as_ref().join(format!(
        "snapshot-{}-{}.{}",
        slot,
        hash.0,
        archive_format.extension(),
    ))
}

/// Build the incremental snapshot archive path from its components: the snapshot archives
/// directory, the snapshot base slot, the snapshot slot, the accounts hash, and the archive
/// format.
pub fn build_incremental_snapshot_archive_path(
    incremental_snapshot_archives_dir: impl AsRef<Path>,
    base_slot: Slot,
    slot: Slot,
    hash: &SnapshotHash,
    archive_format: ArchiveFormat,
) -> PathBuf {
    incremental_snapshot_archives_dir.as_ref().join(format!(
        "incremental-snapshot-{}-{}-{}.{}",
        base_slot,
        slot,
        hash.0,
        archive_format.extension(),
    ))
}

/// Parse a full snapshot archive filename into its Slot, Hash, and Archive Format
pub fn parse_full_snapshot_archive_filename(
    archive_filename: &str,
) -> Result<(Slot, SnapshotHash, ArchiveFormat)> {
    static RE: LazyLock<Regex> =
        LazyLock::new(|| Regex::new(FULL_SNAPSHOT_ARCHIVE_FILENAME_REGEX).unwrap());

    let do_parse = || {
        RE.captures(archive_filename).and_then(|captures| {
            let slot = captures
                .name("slot")
                .map(|x| x.as_str().parse::<Slot>())?
                .ok()?;
            let hash = captures
                .name("hash")
                .map(|x| x.as_str().parse::<Hash>())?
                .ok()?;
            let archive_format = captures
                .name("ext")
                .map(|x| x.as_str().parse::<ArchiveFormat>())?
                .ok()?;

            Some((slot, SnapshotHash(hash), archive_format))
        })
    };

    do_parse().ok_or_else(|| {
        SnapshotError::ParseSnapshotArchiveFileNameError(archive_filename.to_string())
    })
}

/// Parse an incremental snapshot archive filename into its base Slot, actual Slot, Hash, and Archive Format
pub fn parse_incremental_snapshot_archive_filename(
    archive_filename: &str,
) -> Result<(Slot, Slot, SnapshotHash, ArchiveFormat)> {
    static RE: LazyLock<Regex> =
        LazyLock::new(|| Regex::new(INCREMENTAL_SNAPSHOT_ARCHIVE_FILENAME_REGEX).unwrap());

    let do_parse = || {
        RE.captures(archive_filename).and_then(|captures| {
            let base_slot = captures
                .name("base")
                .map(|x| x.as_str().parse::<Slot>())?
                .ok()?;
            let slot = captures
                .name("slot")
                .map(|x| x.as_str().parse::<Slot>())?
                .ok()?;
            let hash = captures
                .name("hash")
                .map(|x| x.as_str().parse::<Hash>())?
                .ok()?;
            let archive_format = captures
                .name("ext")
                .map(|x| x.as_str().parse::<ArchiveFormat>())?
                .ok()?;

            Some((base_slot, slot, SnapshotHash(hash), archive_format))
        })
    };

    do_parse().ok_or_else(|| {
        SnapshotError::ParseSnapshotArchiveFileNameError(archive_filename.to_string())
    })
}

/// Returns the file name of the bank snapshot for `slot`
pub fn get_snapshot_file_name(slot: Slot) -> String {
    slot.to_string()
}

/// Constructs the path to the bank snapshot directory for `slot` within `bank_snapshots_dir`
pub fn get_bank_snapshot_dir(bank_snapshots_dir: impl AsRef<Path>, slot: Slot) -> PathBuf {
    bank_snapshots_dir
        .as_ref()
        .join(get_snapshot_file_name(slot))
}

/// Walk down the snapshot archive to yield snapshot archive file info
fn snapshot_archives_iter<T>(
    snapshot_archives_dir: &Path,
    cb: fn(PathBuf) -> Result<T>,
) -> impl Iterator<Item = T> + use<T> {
    let remote_dir = build_snapshot_archives_remote_dir(snapshot_archives_dir);

    [
        Some(snapshot_archives_dir.to_path_buf()),
        remote_dir.exists().then_some(remote_dir),
    ]
    .into_iter()
    .flatten()
    .flat_map(move |dir| {
        fs::read_dir(&dir)
            .map_err(|err| {
                info!(
                    "Unable to read snapshot archives directory '{}': {err}",
                    dir.display(),
                );
                err
            })
            .ok()
            .into_iter()
            .flat_map(move |entries| {
                entries.filter_map(move |entry| entry.ok().and_then(|entry| cb(entry.path()).ok()))
            })
    })
}

/// Get an iterator of the full snapshot archives from a directory
pub fn full_snapshot_archives_iter(
    full_snapshot_archives_dir: &Path,
) -> impl Iterator<Item = FullSnapshotArchiveInfo> + use<> {
    snapshot_archives_iter(
        full_snapshot_archives_dir,
        FullSnapshotArchiveInfo::new_from_path,
    )
}

/// Get the incremental snapshot archives from a directory
pub fn incremental_snapshot_archives_iter(
    incremental_snapshot_archives_dir: &Path,
) -> impl Iterator<Item = IncrementalSnapshotArchiveInfo> + use<> {
    snapshot_archives_iter(
        incremental_snapshot_archives_dir,
        IncrementalSnapshotArchiveInfo::new_from_path,
    )
}

/// Get the highest slot of the full snapshot archives in a directory
pub fn get_highest_full_snapshot_archive_slot(
    full_snapshot_archives_dir: impl AsRef<Path>,
) -> Option<Slot> {
    get_highest_full_snapshot_archive_info(full_snapshot_archives_dir)
        .map(|full_snapshot_archive_info| full_snapshot_archive_info.slot())
}

/// Get the highest slot of the incremental snapshot archives in a directory, for a given full
/// snapshot slot
pub fn get_highest_incremental_snapshot_archive_slot(
    incremental_snapshot_archives_dir: impl AsRef<Path>,
    full_snapshot_slot: Slot,
) -> Option<Slot> {
    get_highest_incremental_snapshot_archive_info(
        incremental_snapshot_archives_dir,
        full_snapshot_slot,
    )
    .map(|incremental_snapshot_archive_info| incremental_snapshot_archive_info.slot())
}

/// Get the path (and metadata) for the full snapshot archive with the highest slot in a directory
pub fn get_highest_full_snapshot_archive_info(
    full_snapshot_archives_dir: impl AsRef<Path>,
) -> Option<FullSnapshotArchiveInfo> {
    full_snapshot_archives_iter(full_snapshot_archives_dir.as_ref()).max()
}

/// Get the path for the incremental snapshot archive with the highest slot, for a given full
/// snapshot slot, in a directory
pub fn get_highest_incremental_snapshot_archive_info(
    incremental_snapshot_archives_dir: impl AsRef<Path>,
    full_snapshot_slot: Slot,
) -> Option<IncrementalSnapshotArchiveInfo> {
    // Since we want to filter down to only the incremental snapshot archives that have the same
    // full snapshot slot as the value passed in, perform the filtering before sorting to avoid
    // doing unnecessary work.
    let mut incremental_snapshot_archives =
        incremental_snapshot_archives_iter(incremental_snapshot_archives_dir.as_ref())
            .filter(|incremental_snapshot_archive_info| {
                incremental_snapshot_archive_info.base_slot() == full_snapshot_slot
            })
            .collect::<Vec<_>>();
    incremental_snapshot_archives.sort_unstable();
    incremental_snapshot_archives.into_iter().next_back()
}

#[cfg(test)]
mod tests {
    use {super::*, crate::ZstdConfig};

    #[test]
    fn test_parse_full_snapshot_archive_filename() {
        assert_eq!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-43-{}.tar.zst",
                Hash::default()
            ))
            .unwrap(),
            (
                43,
                SnapshotHash(Hash::default()),
                ArchiveFormat::TarZstd {
                    config: ZstdConfig::default(),
                }
            )
        );
        assert_eq!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-45-{}.tar.lz4",
                Hash::default()
            ))
            .unwrap(),
            (45, SnapshotHash(Hash::default()), ArchiveFormat::TarLz4)
        );

        assert!(parse_full_snapshot_archive_filename("invalid").is_err());
        assert!(
            parse_full_snapshot_archive_filename("snapshot-bad!slot-bad!hash.bad!ext").is_err()
        );

        assert!(
            parse_full_snapshot_archive_filename("snapshot-12345678-bad!hash.bad!ext").is_err()
        );
        assert!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-12345678-{}.bad!ext",
                Hash::new_unique()
            ))
            .is_err()
        );
        assert!(
            parse_full_snapshot_archive_filename("snapshot-12345678-bad!hash.tar.zst").is_err()
        );

        assert!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-bad!slot-{}.bad!ext",
                Hash::new_unique()
            ))
            .is_err()
        );
        assert!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-12345678-{}.bad!ext",
                Hash::new_unique()
            ))
            .is_err()
        );
        assert!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-bad!slot-{}.tar.zst",
                Hash::new_unique()
            ))
            .is_err()
        );

        assert!(
            parse_full_snapshot_archive_filename("snapshot-bad!slot-bad!hash.tar.zst").is_err()
        );
        assert!(
            parse_full_snapshot_archive_filename("snapshot-12345678-bad!hash.tar.zst").is_err()
        );
        assert!(
            parse_full_snapshot_archive_filename(&format!(
                "snapshot-bad!slot-{}.tar.zst",
                Hash::new_unique()
            ))
            .is_err()
        );
    }

    #[test]
    fn test_parse_incremental_snapshot_archive_filename() {
        assert_eq!(
            parse_incremental_snapshot_archive_filename(&format!(
                "incremental-snapshot-43-234-{}.tar.zst",
                Hash::default()
            ))
            .unwrap(),
            (
                43,
                234,
                SnapshotHash(Hash::default()),
                ArchiveFormat::TarZstd {
                    config: ZstdConfig::default(),
                }
            )
        );
        assert_eq!(
            parse_incremental_snapshot_archive_filename(&format!(
                "incremental-snapshot-45-456-{}.tar.lz4",
                Hash::default()
            ))
            .unwrap(),
            (
                45,
                456,
                SnapshotHash(Hash::default()),
                ArchiveFormat::TarLz4
            )
        );

        assert!(parse_incremental_snapshot_archive_filename("invalid").is_err());
        assert!(
            parse_incremental_snapshot_archive_filename(&format!(
                "snapshot-42-{}.tar.zst",
                Hash::new_unique()
            ))
            .is_err()
        );
        assert!(
            parse_incremental_snapshot_archive_filename(
                "incremental-snapshot-bad!slot-bad!slot-bad!hash.bad!ext"
            )
            .is_err()
        );

        assert!(
            parse_incremental_snapshot_archive_filename(&format!(
                "incremental-snapshot-bad!slot-56785678-{}.tar.zst",
                Hash::new_unique()
            ))
            .is_err()
        );

        assert!(
            parse_incremental_snapshot_archive_filename(&format!(
                "incremental-snapshot-12345678-bad!slot-{}.tar.zst",
                Hash::new_unique()
            ))
            .is_err()
        );

        assert!(
            parse_incremental_snapshot_archive_filename(
                "incremental-snapshot-12341234-56785678-bad!HASH.tar.zst"
            )
            .is_err()
        );

        assert!(
            parse_incremental_snapshot_archive_filename(&format!(
                "incremental-snapshot-12341234-56785678-{}.bad!ext",
                Hash::new_unique()
            ))
            .is_err()
        );
    }
}