WIP: feat: Add methods for repairing hot/cold repositories

Author: aawsome

crates/core/src/commands/config.rs

@@ -82,7 +82,32 @@ pub(crate) fn save_config<P, S>(
     let dbe = DecryptBackend::new(repo.be.clone(), key);
     // for hot/cold backend, this only saves the config to the cold repo.
     _ = dbe.save_file_uncompressed(&new_config)?;
+    save_config_hot(repo, new_config, key)
+}
 
+/// Save a [`ConfigFile`] only to the hot part of a repository
+///
+/// # Type Parameters
+///
+/// * `P` - The progress bar type.
+/// * `S` - The state the repository is in.
+///
+/// # Arguments
+///
+/// * `repo` - The repository to save the config to
+/// * `new_config` - The config to save
+/// * `key` - The key to encrypt the config with
+///
+/// # Errors
+///
+/// * [`CryptBackendErrorKind::SerializingToJsonByteVectorFailed`] - If the file could not be serialized to json.
+///
+/// [`CryptBackendErrorKind::SerializingToJsonByteVectorFailed`]: crate::error::CryptBackendErrorKind::SerializingToJsonByteVectorFailed
+pub(crate) fn save_config_hot<P, S>(
+    repo: &Repository<P, S>,
+    mut new_config: ConfigFile,
+    key: impl CryptoKey,
+) -> RusticResult<()> {
     if let Some(hot_be) = repo.be_hot.clone() {
         // save config to hot repo
         let dbe = DecryptBackend::new(hot_be.clone(), key);

crates/core/src/commands/init.rs

@@ -60,7 +60,7 @@ pub(crate) fn init<P, S>(
     Ok((key, config))
 }
 
-/// Initialize a new repository with a given config.
+/// Save a [`ConfigFile`] only to the hot part of a repository
 ///
 /// # Type Parameters
 ///
@@ -69,14 +69,15 @@ pub(crate) fn init<P, S>(
 ///
 /// # Arguments
 ///
-/// * `repo` - The repository to initialize.
-/// * `pass` - The password to encrypt the key with.
-/// * `key_opts` - The options to create the key with.
-/// * `config` - The config to use.
+/// * `repo` - The repository to save the config to
+/// * `new_config` - The config to save
+/// * `key` - The key to encrypt the config with
 ///
-/// # Returns
+/// # Errors
+///
+/// * [`CryptBackendErrorKind::SerializingToJsonByteVectorFailed`] - If the file could not be serialized to json.
 ///
-/// The key used to encrypt the config.
+/// [`CryptBackendErrorKind::SerializingToJsonByteVectorFailed`]: crate::error::CryptBackendErrorKind::SerializingToJsonByteVectorFailed
 pub(crate) fn init_with_config<P, S>(
     repo: &Repository<P, S>,
     pass: &str,

crates/core/src/commands/repair.rs

@@ -1,2 +1,3 @@
+pub mod hotcold;
 pub mod index;
 pub mod snapshots;

crates/core/src/commands/repair/hotcold.rs

@@ -0,0 +1,180 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use log::{debug, info, warn};
+
+use crate::{
+    backend::decrypt::DecryptReadBackend,
+    repofile::{BlobType, IndexFile, PackId},
+    repository::Open,
+    ErrorKind, FileType, Id, Progress, ProgressBars, ReadBackend, Repository, RusticError,
+    RusticResult, WriteBackend, ALL_FILE_TYPES,
+};
+
+pub(crate) fn repair_hotcold<P: ProgressBars, S>(
+    repo: &Repository<P, S>,
+    dry_run: bool,
+) -> RusticResult<()> {
+    for file_type in ALL_FILE_TYPES {
+        if file_type != FileType::Pack {
+            correct_missing_files(repo, file_type, |_| true, dry_run)?;
+        }
+    }
+    Ok(())
+}
+
+pub(crate) fn repair_hotcold_packs<P: ProgressBars, S: Open>(
+    repo: &Repository<P, S>,
+    dry_run: bool,
+) -> RusticResult<()> {
+    let tree_packs = get_tree_packs(repo)?;
+    correct_missing_files(
+        repo,
+        FileType::Pack,
+        |id| tree_packs.contains(&PackId::from(*id)),
+        dry_run,
+    )
+}
+
+pub(crate) fn correct_missing_files<P: ProgressBars, S>(
+    repo: &Repository<P, S>,
+    file_type: FileType,
+    is_relevant: impl Fn(&Id) -> bool,
+    dry_run: bool,
+) -> RusticResult<()> {
+    let Some(repo_hot) = &repo.be_hot else {
+        return Err(RusticError::new(
+            ErrorKind::Repository,
+            "Repository is no hot/cold repository.",
+        ));
+    };
+
+    let (missing_hot, missing_hot_size, missing_cold, missing_cold_size) =
+        get_missing_files(repo, file_type, is_relevant)?;
+
+    // copy missing files from hot to cold repo
+    if !missing_cold.is_empty() {
+        if dry_run {
+            info!(
+                "would have copied {} hot {file_type:?} files to cold",
+                missing_cold.len()
+            );
+            debug!("files: {missing_cold:?}");
+        } else {
+            let p = repo
+                .pb
+                .progress_bytes(format!("copying missing cold {file_type:?} files..."));
+            p.set_length(missing_cold_size);
+            copy(missing_cold, file_type, repo_hot, &repo.be_cold)?;
+            p.finish();
+        }
+    }
+
+    if !missing_hot.is_empty() {
+        if dry_run {
+            info!(
+                "would have copied {} cold {file_type:?} files to hot",
+                missing_hot.len()
+            );
+            debug!("files: {missing_hot:?}");
+        } else {
+            // TODO: warm-up
+            // copy missing files from cold to hot repo
+            let p = repo
+                .pb
+                .progress_bytes(format!("copying missing hot {file_type:?} files..."));
+            p.set_length(missing_hot_size);
+            copy(missing_hot, file_type, &repo.be_cold, repo_hot)?;
+            p.finish();
+        }
+    }
+
+    Ok(())
+}
+
+fn copy(
+    files: Vec<Id>,
+    file_type: FileType,
+    from: &impl ReadBackend,
+    to: &impl WriteBackend,
+) -> RusticResult<()> {
+    for id in files {
+        let file = from.read_full(file_type, &id)?;
+        to.write_bytes(file_type, &id, false, file)?;
+    }
+    Ok(())
+}
+
+pub(crate) fn get_tree_packs<P: ProgressBars, S: Open>(
+    repo: &Repository<P, S>,
+) -> RusticResult<BTreeSet<PackId>> {
+    let p = repo.pb.progress_counter("reading index...");
+    let mut tree_packs = BTreeSet::new();
+    for index in repo.dbe().stream_all::<IndexFile>(&p)? {
+        let index = index?.1;
+        for (p, _) in index.all_packs() {
+            let blob_type = p.blob_type();
+            if blob_type == BlobType::Tree {
+                _ = tree_packs.insert(p.id);
+            }
+        }
+    }
+    Ok(tree_packs)
+}
+
+pub(crate) fn get_missing_files<P: ProgressBars, S>(
+    repo: &Repository<P, S>,
+    file_type: FileType,
+    is_relevant: impl Fn(&Id) -> bool,
+) -> RusticResult<(Vec<Id>, u64, Vec<Id>, u64)> {
+    let Some(repo_hot) = &repo.be_hot else {
+        return Err(RusticError::new(
+            ErrorKind::Repository,
+            "Repository is no hot/cold repository.",
+        ));
+    };
+
+    let p = repo
+        .pb
+        .progress_spinner(format!("listing hot {file_type:?} files..."));
+    let hot_files: BTreeMap<_, _> = repo_hot.list_with_size(file_type)?.into_iter().collect();
+    p.finish();
+
+    let p = repo
+        .pb
+        .progress_spinner(format!("listing cold {file_type:?} files..."));
+    let cold_files: BTreeMap<_, _> = repo
+        .be_cold
+        .list_with_size(file_type)?
+        .into_iter()
+        .collect();
+    p.finish();
+
+    let common: BTreeSet<_> = hot_files
+            .iter()
+            .filter_map(|(id, size_hot)| match cold_files.get(id) {
+                Some(size_cold) if size_cold == size_hot => Some(*id),
+                Some(size_cold) => {
+                     warn!("sizes mismatch: type {file_type:?}, id: {id}, size hot: {size_hot}, size cold: {size_cold}. Ignoring...");
+                    None
+                }
+                None => None,
+            })
+            .collect();
+
+    let retain = |files: BTreeMap<_, _>| {
+        let mut retain_size: u64 = 0;
+        let only: Vec<_> = files
+            .into_iter()
+            .filter(|(id, _)| !common.contains(id) && is_relevant(id))
+            .map(|(id, size)| {
+                retain_size += u64::from(size);
+                id
+            })
+            .collect();
+        (only, retain_size)
+    };
+
+    let (cold_only, cold_only_size) = retain(cold_files);
+    let (hot_only, hot_only_size) = retain(hot_files);
+    Ok((cold_only, cold_only_size, hot_only, hot_only_size))
+}