feat: Add lock_repo command

Author: aawsome

crates/core/src/backend.rs

@@ -6,6 +6,7 @@ pub(crate) mod dry_run;
 pub(crate) mod hotcold;
 pub(crate) mod ignore;
 pub(crate) mod local_destination;
+pub(crate) mod lock;
 pub(crate) mod node;
 pub(crate) mod stdin;
 pub(crate) mod warm_up;
@@ -14,8 +15,9 @@ use std::{io::Read, ops::Deref, path::PathBuf, sync::Arc};
 
 use anyhow::Result;
 use bytes::Bytes;
+use chrono::{DateTime, Local};
 use enum_map::Enum;
-use log::trace;
+use log::{debug, trace};
 
 #[cfg(test)]
 use mockall::mock;
@@ -337,6 +339,27 @@ pub trait WriteBackend: ReadBackend {
     ///
     /// The result of the removal.
     fn remove(&self, tpe: FileType, id: &Id, cacheable: bool) -> Result<()>;
+
+    /// Specify if the backend is able to lock files
+    fn can_lock(&self) -> bool {
+        false
+    }
+
+    /// Lock the given file.
+    ///
+    /// # Arguments
+    ///
+    /// * `tpe` - The type of the file.
+    /// * `id` - The id of the file.
+    /// * `until` - The date until when to lock. May be `None` which usually specifies a unlimited lock
+    ///
+    /// # Errors
+    ///
+    /// If the file could not be read.
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        debug!("no locking implemented. {tpe:?}, {id}, {until:?}");
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -374,6 +397,12 @@ impl WriteBackend for Arc<dyn WriteBackend> {
     fn remove(&self, tpe: FileType, id: &Id, cacheable: bool) -> Result<()> {
         self.deref().remove(tpe, id, cacheable)
     }
+    fn can_lock(&self) -> bool {
+        self.deref().can_lock()
+    }
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        self.deref().lock(tpe, id, until)
+    }
 }
 
 impl ReadBackend for Arc<dyn WriteBackend> {

crates/core/src/backend/cache.rs

@@ -8,6 +8,7 @@ use std::{
 
 use anyhow::Result;
 use bytes::Bytes;
+use chrono::{DateTime, Local};
 use dirs::cache_dir;
 use log::{trace, warn};
 use walkdir::WalkDir;
@@ -210,6 +211,14 @@ impl WriteBackend for CachedBackend {
         }
         self.be.remove(tpe, id, cacheable)
     }
+
+    fn can_lock(&self) -> bool {
+        self.be.can_lock()
+    }
+
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        self.be.lock(tpe, id, until)
+    }
 }
 
 /// Backend that caches data in a directory.

crates/core/src/backend/decrypt.rs

@@ -2,6 +2,7 @@ use std::{num::NonZeroU32, sync::Arc};
 
 use anyhow::Result;
 use bytes::Bytes;
+use chrono::{DateTime, Local};
 use crossbeam_channel::{unbounded, Receiver};
 use rayon::prelude::*;
 use zstd::stream::{copy_encode, decode_all, encode_all};
@@ -578,6 +579,14 @@ impl<C: CryptoKey> WriteBackend for DecryptBackend<C> {
     fn remove(&self, tpe: FileType, id: &Id, cacheable: bool) -> Result<()> {
         self.be.remove(tpe, id, cacheable)
     }
+
+    fn can_lock(&self) -> bool {
+        self.be.can_lock()
+    }
+
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        self.be.lock(tpe, id, until)
+    }
 }
 
 #[cfg(test)]

crates/core/src/backend/dry_run.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
 use bytes::Bytes;
+use chrono::{DateTime, Local};
 use zstd::decode_all;
 
 use crate::{
@@ -156,4 +157,12 @@ impl<BE: DecryptFullBackend> WriteBackend for DryRunBackend<BE> {
             self.be.remove(tpe, id, cacheable)
         }
     }
+
+    fn can_lock(&self) -> bool {
+        self.be.can_lock()
+    }
+
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        self.be.lock(tpe, id, until)
+    }
 }

crates/core/src/backend/hotcold.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use bytes::Bytes;
+use chrono::{DateTime, Local};
 
 use crate::{
     backend::{FileType, ReadBackend, WriteBackend},
@@ -98,4 +99,12 @@ impl WriteBackend for HotColdBackend {
         }
         Ok(())
     }
+
+    fn can_lock(&self) -> bool {
+        self.be.can_lock()
+    }
+
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        self.be.lock(tpe, id, until)
+    }
 }

crates/core/src/backend/lock.rs

@@ -0,0 +1,113 @@
+use std::{process::Command, sync::Arc};
+
+use anyhow::Result;
+use bytes::Bytes;
+use chrono::{DateTime, Local};
+use log::{debug, warn};
+
+use crate::{
+    backend::{FileType, ReadBackend, WriteBackend},
+    id::Id,
+    CommandInput,
+};
+
+/// A backend which warms up files by simply accessing them.
+#[derive(Clone, Debug)]
+pub struct LockBackend {
+    /// The backend to use.
+    be: Arc<dyn WriteBackend>,
+    /// The command to be called to lock files in the backend
+    command: CommandInput,
+}
+
+impl LockBackend {
+    /// Creates a new `WarmUpAccessBackend`.
+    ///
+    /// # Arguments
+    ///
+    /// * `be` - The backend to use.
+    pub fn new_lock(be: Arc<dyn WriteBackend>, command: CommandInput) -> Arc<dyn WriteBackend> {
+        Arc::new(Self { be, command })
+    }
+}
+
+impl ReadBackend for LockBackend {
+    fn location(&self) -> String {
+        self.be.location()
+    }
+
+    fn list_with_size(&self, tpe: FileType) -> Result<Vec<(Id, u32)>> {
+        self.be.list_with_size(tpe)
+    }
+
+    fn read_full(&self, tpe: FileType, id: &Id) -> Result<Bytes> {
+        self.be.read_full(tpe, id)
+    }
+
+    fn read_partial(
+        &self,
+        tpe: FileType,
+        id: &Id,
+        cacheable: bool,
+        offset: u32,
+        length: u32,
+    ) -> Result<Bytes> {
+        self.be.read_partial(tpe, id, cacheable, offset, length)
+    }
+
+    fn list(&self, tpe: FileType) -> Result<Vec<Id>> {
+        self.be.list(tpe)
+    }
+
+    fn needs_warm_up(&self) -> bool {
+        self.be.needs_warm_up()
+    }
+
+    fn warm_up(&self, tpe: FileType, id: &Id) -> Result<()> {
+        self.be.warm_up(tpe, id)
+    }
+}
+
+fn path(tpe: FileType, id: &Id) -> String {
+    let hex_id = id.to_hex();
+    match tpe {
+        FileType::Config => "config".into(),
+        FileType::Pack => format!("data/{}/{}", &hex_id[0..2], &*hex_id),
+        _ => format!("{}/{}", tpe.dirname(), &*hex_id),
+    }
+}
+
+impl WriteBackend for LockBackend {
+    fn create(&self) -> Result<()> {
+        self.be.create()
+    }
+
+    fn write_bytes(&self, tpe: FileType, id: &Id, cacheable: bool, buf: Bytes) -> Result<()> {
+        self.be.write_bytes(tpe, id, cacheable, buf)
+    }
+
+    fn remove(&self, tpe: FileType, id: &Id, cacheable: bool) -> Result<()> {
+        self.be.remove(tpe, id, cacheable)
+    }
+
+    fn can_lock(&self) -> bool {
+        true
+    }
+
+    fn lock(&self, tpe: FileType, id: &Id, until: Option<DateTime<Local>>) -> Result<()> {
+        let until = until.map_or_else(|| String::new(), |u| format!("{}", u.to_rfc3339()));
+        let path = path(tpe, id);
+        let args = self.command.args().iter().map(|c| {
+            c.replace("%id", &id.to_hex())
+                .replace("%type", tpe.dirname())
+                .replace("%path", &path)
+                .replace("%until", &until)
+        });
+        debug!("calling {:?}...", self.command);
+        let status = Command::new(self.command.command()).args(args).status()?;
+        if !status.success() {
+            warn!("lock command was not successful for {tpe:?}, id: {id}. {status}");
+        }
+        Ok(())
+    }
+}

crates/core/src/commands.rs

@@ -12,6 +12,7 @@ pub mod dump;
 pub mod forget;
 pub mod init;
 pub mod key;
+pub mod lock;
 pub mod merge;
 pub mod prune;
 /// The `repair` command.

crates/core/src/commands/lock.rs

@@ -0,0 +1,75 @@
+//! `lock` subcommand
+
+use chrono::{DateTime, Local};
+use log::error;
+use rayon::ThreadPoolBuilder;
+
+use crate::{
+    error::{CommandErrorKind, RepositoryErrorKind, RusticResult},
+    progress::{Progress, ProgressBars},
+    repofile::{configfile::ConfigId, IndexId, KeyId, PackId, RepoId, SnapshotId},
+    repository::Repository,
+    WriteBackend,
+};
+
+pub(super) mod constants {
+    /// The maximum number of reader threads to use for locking.
+    pub(super) const MAX_LOCKER_THREADS_NUM: usize = 20;
+}
+
+pub fn lock_repo<P: ProgressBars, S>(
+    repo: &Repository<P, S>,
+    until: Option<DateTime<Local>>,
+) -> RusticResult<()> {
+    lock_all_files::<P, S, ConfigId>(repo, until)?;
+    lock_all_files::<P, S, KeyId>(repo, until)?;
+    lock_all_files::<P, S, SnapshotId>(repo, until)?;
+    lock_all_files::<P, S, IndexId>(repo, until)?;
+    lock_all_files::<P, S, PackId>(repo, until)?;
+    Ok(())
+}
+
+pub fn lock_all_files<P: ProgressBars, S, ID: RepoId + std::fmt::Debug>(
+    repo: &Repository<P, S>,
+    until: Option<DateTime<Local>>,
+) -> RusticResult<()> {
+    if !repo.be.can_lock() {
+        return Err(CommandErrorKind::NoLockingConfigured.into());
+    }
+
+    let p = &repo
+        .pb
+        .progress_spinner(format!("listing {:?} files..", ID::TYPE));
+    let ids: Vec<ID> = repo.list()?.collect();
+    p.finish();
+    lock_files(repo, &ids, until)
+}
+
+fn lock_files<P: ProgressBars, S, ID: RepoId + std::fmt::Debug>(
+    repo: &Repository<P, S>,
+    ids: &[ID],
+    until: Option<DateTime<Local>>,
+) -> RusticResult<()> {
+    let pool = ThreadPoolBuilder::new()
+        .num_threads(constants::MAX_LOCKER_THREADS_NUM)
+        .build()
+        .map_err(RepositoryErrorKind::FromThreadPoolbilderError)?;
+    let p = &repo
+        .pb
+        .progress_counter(format!("locking {:?} files..", ID::TYPE));
+    p.set_length(ids.len().try_into().unwrap());
+    let backend = &repo.be;
+    pool.in_place_scope(|scope| {
+        for id in ids {
+            scope.spawn(move |_| {
+                if let Err(e) = backend.lock(ID::TYPE, id, until) {
+                    // FIXME: Use error handling
+                    error!("lock failed for {:?} {id:?}. {e}", ID::TYPE);
+                };
+                p.inc(1);
+            });
+        }
+    });
+    p.finish();
+    Ok(())
+}

crates/core/src/error.rs

@@ -229,6 +229,8 @@ pub enum CommandErrorKind {
     NoKeepOption,
     /// {0:?}
     FromParseError(#[from] shell_words::ParseError),
+    /// No locking capability configured for the backend
+    NoLockingConfigured,
 }
 
 /// [`CryptoErrorKind`] describes the errors that can happen while dealing with Cryptographic functions