feat: Add methods for repairing hot/cold repositories

crates/core/src/backend.rs

@@ -390,6 +390,12 @@ impl ReadBackend for Arc<dyn WriteBackend> {
     fn warmup_path(&self, tpe: FileType, id: &Id) -> String {
         self.deref().warmup_path(tpe, id)
     }
+    fn needs_warm_up(&self) -> bool {
+        self.deref().needs_warm_up()
+    }
+    fn warm_up(&self, tpe: FileType, id: &Id) -> RusticResult<()> {
+        self.deref().warm_up(tpe, id)
+    }
 }
 
 impl std::fmt::Debug for dyn WriteBackend {

crates/core/src/commands/config.rs

@@ -15,7 +15,6 @@ use crate::{
 ///
 /// # Type Parameters
 ///
-/// * `P` - The progress bar type.
 /// * `S` - The state the repository is in.
 ///
 /// # Arguments
@@ -63,7 +62,6 @@ pub(crate) fn apply_config<S: Open>(
 ///
 /// # Type Parameters
 ///
-/// * `P` - The progress bar type.
 /// * `S` - The state the repository is in.
 ///
 /// # Arguments
@@ -84,7 +82,30 @@ pub(crate) fn save_config<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
+///
+/// * `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
+///
+/// * If the file could not be serialized to json.
+/// * If the file could not be saved to the backend.
+pub(crate) fn save_config_hot<S>(
+    repo: &Repository<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

@@ -20,7 +20,6 @@ use crate::{
 ///
 /// # Type Parameters
 ///
-/// * `P` - The progress bar type.
 /// * `S` - The state the repository is in.
 ///
 /// # Arguments
@@ -60,17 +59,16 @@ pub(crate) fn init<S>(
     Ok((key, key_id, config))
 }
 
-/// Initialize a new repository with a given config.
+/// Initializes a new repository using the given [`ConfigFile`]
 ///
 /// # Type Parameters
 ///
-/// * `P` - The progress bar type.
 /// * `S` - The state the repository is in.
 ///
 /// # Arguments
 ///
 /// * `repo` - The repository to initialize.
-/// * `pass` - The password to encrypt the key with.
+/// * `credentials` - The credentials to use.
 /// * `key_opts` - The options to create the key with.
 /// * `config` - The config to use.
 ///

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,181 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use log::{debug, info, warn};
+use rayon::prelude::{IntoParallelIterator, ParallelIterator};
+
+use crate::{
+    ALL_FILE_TYPES, ErrorKind, FileType, Id, Progress, ReadBackend, Repository, RusticError,
+    RusticResult, WriteBackend,
+    backend::decrypt::DecryptReadBackend,
+    repofile::{BlobType, IndexFile, PackId},
+    repository::{Open, warm_up::warm_up_wait},
+};
+
+/// Repairs a hot/cold repository by copying missing files (except pack files) over from one to the other part.
+pub(crate) fn repair_hotcold<S>(repo: &Repository<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(())
+}
+
+/// Repairs a hot/cold repository by copying missing tree packs over from one to the other part.
+pub(crate) fn repair_hotcold_packs<S: Open>(
+    repo: &Repository<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,
+    )
+}
+
+/// Copy relevant+missing files in a hot/cold repository from one to the other part.
+pub(crate) fn correct_missing_files<S>(
+    repo: &Repository<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.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)?;
+            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 {
+            warm_up_wait(repo, file_type, missing_hot.iter().copied())?;
+            // copy missing files from cold to hot repo
+            let p = repo.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)?;
+            p.finish();
+        }
+    }
+
+    Ok(())
+}
+
+/// copy a list of files from one repository part to the other
+fn copy(
+    files: Vec<Id>,
+    file_type: FileType,
+    from: &impl ReadBackend,
+    to: &impl WriteBackend,
+    p: &Progress,
+) -> RusticResult<()> {
+    files.into_par_iter().try_for_each(|id| {
+        let file = from.read_full(file_type, &id)?;
+        let length = u64::try_from(file.len()).expect("file len should fit into u64");
+        to.write_bytes(file_type, &id, false, file)?;
+        p.inc(length);
+        Ok(())
+    })
+}
+
+/// Get all tree packs from within the repository
+pub(crate) fn get_tree_packs<S: Open>(repo: &Repository<S>) -> RusticResult<BTreeSet<PackId>> {
+    let p = repo.progress_counter("reading index...");
+    let mut tree_packs = BTreeSet::new();
+    for index in repo.dbe().stream_all::<IndexFile>(&p)? {
+        let index = index?.1;
+        for (pack, _) in index.all_packs() {
+            let blob_type = pack.blob_type();
+            if blob_type == BlobType::Tree {
+                _ = tree_packs.insert(pack.id);
+            }
+        }
+    }
+    Ok(tree_packs)
+}
+
+/// Find missing files in the hot or cold part of a repository
+///
+/// # Returns
+///
+/// A tuble containing missing ids in the hot part; the corresponding total size; missing ids in the cold part; the corresponding total size.
+pub(crate) fn get_missing_files<S>(
+    repo: &Repository<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.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.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))
+}