Skip to content

Commit

Permalink
WIP: feat: Add methods for repairing hot/cold repositories
Browse files Browse the repository at this point in the history
  • Loading branch information
aawsome committed Nov 18, 2024
1 parent 749879f commit 1b5fb66
Show file tree
Hide file tree
Showing 5 changed files with 335 additions and 20 deletions.
25 changes: 25 additions & 0 deletions crates/core/src/commands/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
15 changes: 8 additions & 7 deletions crates/core/src/commands/init.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
///
Expand All @@ -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,
Expand Down
1 change: 1 addition & 0 deletions crates/core/src/commands/repair.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
pub mod hotcold;
pub mod index;
pub mod snapshots;
180 changes: 180 additions & 0 deletions crates/core/src/commands/repair/hotcold.rs
Original file line number Diff line number Diff line change
@@ -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))
}
Loading

0 comments on commit 1b5fb66

Please sign in to comment.