prune: Allow to read (nearly) contiguous blobs at once during rewrite

Author: aawsome

crates/core/src/blob/packer.rs

@@ -846,6 +846,25 @@ impl Actor {
     }
 }
 
+#[derive(Debug)]
+pub struct RepackPackBlobs {
+    pub pack_id: PackId,
+    pub offset: u32,
+    pub length: u32,
+    pub blobs: Vec<IndexBlob>,
+}
+
+impl RepackPackBlobs {
+    pub fn from_index_blob(pack_id: PackId, blob: IndexBlob) -> Self {
+        Self {
+            pack_id,
+            offset: blob.offset,
+            length: blob.length,
+            blobs: vec![blob],
+        }
+    }
+}
+
 /// The `Repacker` is responsible for repacking blobs into pack files.
 ///
 /// # Type Parameters
@@ -862,6 +881,8 @@ where
     packer: Packer<BE>,
     /// The size limit of the pack file.
     size_limit: u32,
+    /// the blob type
+    blob_type: BlobType,
 }
 
 impl<BE: DecryptFullBackend> Repacker<BE> {
@@ -895,6 +916,7 @@ impl<BE: DecryptFullBackend> Repacker<BE> {
             be,
             packer,
             size_limit,
+            blob_type,
         })
     }
 
@@ -909,30 +931,38 @@ impl<BE: DecryptFullBackend> Repacker<BE> {
     ///
     /// * If the blob could not be added
     /// * If reading the blob from the backend fails
-    pub fn add_fast(&self, pack_id: &PackId, blob: &IndexBlob) -> RusticResult<()> {
+    pub fn add_fast(&self, pack_blobs: RepackPackBlobs) -> RusticResult<()> {
+        let offset = pack_blobs.offset;
         let data = self.be.read_partial(
             FileType::Pack,
-            pack_id,
-            blob.tpe.is_cacheable(),
-            blob.offset,
-            blob.length,
+            &pack_blobs.pack_id,
+            self.blob_type.is_cacheable(),
+            offset,
+            pack_blobs.length,
         )?;
 
-        self.packer
-            .add_raw(
-                &data,
-                &blob.id,
-                0,
-                blob.uncompressed_length,
-                Some(self.size_limit),
-            )
-            .map_err(|err| {
-                err.overwrite_kind(ErrorKind::Internal)
-                    .prepend_guidance_line(
-                        "Failed to fast-add (unchecked) blob `{blob_id}` to packfile.",
-                    )
-                    .attach_context("blob_id", blob.id.to_string())
-            })?;
+        // TODO: write in parallel
+        for blob in pack_blobs.blobs {
+            let start = usize::try_from(blob.offset - offset)
+                .expect("convert from u32 to usize should not fail!");
+            let end = usize::try_from(blob.offset + blob.length - offset)
+                .expect("convert from u32 to usize should not fail!");
+            self.packer
+                .add_raw(
+                    &data[start..end],
+                    &blob.id,
+                    0,
+                    blob.uncompressed_length,
+                    Some(self.size_limit),
+                )
+                .map_err(|err| {
+                    err.overwrite_kind(ErrorKind::Internal)
+                        .prepend_guidance_line(
+                            "Failed to fast-add (unchecked) blob `{blob_id}` to packfile.",
+                        )
+                        .attach_context("blob_id", blob.id.to_string())
+                })?;
+        }
 
         Ok(())
     }
@@ -948,25 +978,36 @@ impl<BE: DecryptFullBackend> Repacker<BE> {
     ///
     /// * If the blob could not be added
     /// * If reading the blob from the backend fails
-    pub fn add(&self, pack_id: &PackId, blob: &IndexBlob) -> RusticResult<()> {
-        let data = self.be.read_encrypted_partial(
+    pub fn add(&self, pack_blobs: RepackPackBlobs) -> RusticResult<()> {
+        let offset = pack_blobs.offset;
+        let read_data = self.be.read_partial(
             FileType::Pack,
-            pack_id,
-            blob.tpe.is_cacheable(),
-            blob.offset,
-            blob.length,
-            blob.uncompressed_length,
+            &pack_blobs.pack_id,
+            self.blob_type.is_cacheable(),
+            offset,
+            pack_blobs.length,
         )?;
 
-        self.packer
-            .add_with_sizelimit(data, blob.id, Some(self.size_limit))
-            .map_err(|err| {
-                RusticError::with_source(
-                    ErrorKind::Internal,
-                    "Failed to add blob to packfile.",
-                    err,
-                )
-            })?;
+        // TODO: write in parallel
+        for blob in pack_blobs.blobs {
+            let start = usize::try_from(blob.offset - offset)
+                .expect("convert from u32 to usize should not fail!");
+            let end = usize::try_from(blob.offset + blob.length - offset)
+                .expect("convert from u32 to usize should not fail!");
+            let data = self
+                .be
+                .read_encrypted_from_partial(&read_data[start..end], blob.uncompressed_length)?;
+
+            self.packer
+                .add_with_sizelimit(data, blob.id, Some(self.size_limit))
+                .map_err(|err| {
+                    RusticError::with_source(
+                        ErrorKind::Internal,
+                        "Failed to add blob to packfile.",
+                        err,
+                    )
+                })?;
+        }
 
         Ok(())
     }

crates/core/src/commands/prune.rs

@@ -26,7 +26,7 @@ use crate::{
     },
     blob::{
         BlobId, BlobType, BlobTypeMap, Initialize,
-        packer::{PackSizer, Repacker},
+        packer::{PackSizer, RepackPackBlobs, Repacker},
         tree::TreeStreamerOnce,
     },
     error::{ErrorKind, RusticError, RusticResult},
@@ -46,6 +46,11 @@ use crate::{
 pub(super) mod constants {
     /// Minimum size of an index file to be considered for pruning
     pub(super) const MIN_INDEX_LEN: usize = 10_000;
+    /// The maximum size of pack-part which is read at once from the backend.
+    /// (needed to limit the memory size used for large backends)
+    pub(crate) const LIMIT_PACK_READ: u32 = 40 * 1024 * 1024; // 40 MiB
+    /// The maximum size of holes which are still read when repacking
+    pub(crate) const MAX_HOLESIZE: u32 = 256 * 1024; // 256 kiB
 }
 
 #[allow(clippy::struct_excessive_bools)]
@@ -1266,6 +1271,11 @@ pub(crate) fn prune_repository<P: ProgressBars, S: Open>(
         }
     }
 
+    if prune_plan.index_files.is_empty() {
+        info!("nothing to do!");
+        return Ok(());
+    }
+
     let mut indexes_remove = Vec::new();
     let mut tree_packs_remove = Vec::new();
     let mut data_packs_remove = Vec::new();
@@ -1282,12 +1292,8 @@ pub(crate) fn prune_repository<P: ProgressBars, S: Open>(
     let mut repack_packs = Vec::new();
 
     // process packs by index_file
-    let p = if prune_plan.index_files.is_empty() {
-        info!("nothing to do!");
-        pb.progress_hidden()
-    } else {
-        pb.progress_spinner("rebuilding index...")
-    };
+    let p = pb.progress_counter("rebuilding index...");
+    p.set_length(u64::try_from(prune_plan.index_files.len()).unwrap_or_default());
     for index in prune_plan.index_files {
         indexes_remove.push(index.id);
         for mut pack in index.packs {
@@ -1397,18 +1403,39 @@ pub(crate) fn prune_repository<P: ProgressBars, S: Open>(
         repack_packs
             .into_par_iter()
             .try_for_each(|pack| -> RusticResult<_> {
+                let repacker = match pack.blob_type {
+                    BlobType::Data => &data_repacker,
+                    BlobType::Tree => &tree_repacker,
+                };
+                let blobs: Vec<_> = pack
+                    .blobs
+                    .into_iter()
+                    .map(|blob| RepackPackBlobs::from_index_blob(pack.id, blob))
+                    .coalesce(|mut x, mut y| {
+                        // if the blobs are (almost) contiguous and we don't trespass the limit, read blobs one partial read
+                        if y.offset <= x.offset + x.length + constants::MAX_HOLESIZE
+                            && y.offset > x.offset
+                            && y.offset + y.length - x.offset <= constants::LIMIT_PACK_READ
+                        {
+                            x.length = y.offset + y.length - x.offset; // read till the end of y
+                            x.blobs.append(&mut y.blobs);
+                            Ok(x)
+                        } else {
+                            Err((x, y))
+                        }
+                    })
+                    .collect();
+
                 // TODO: repack in parallel
-                for blob in &pack.blobs {
-                    let repacker = match blob.tpe {
-                        BlobType::Data => &data_repacker,
-                        BlobType::Tree => &tree_repacker,
-                    };
+                for blobs in blobs {
+                    dbg!(&blobs);
+                    let length = u64::from(blobs.length);
                     if opts.fast_repack {
-                        repacker.add_fast(&pack.id, blob)?;
+                        repacker.add_fast(blobs)?;
                     } else {
-                        repacker.add(&pack.id, blob)?;
+                        repacker.add(blobs)?;
                     }
-                    p.inc(u64::from(blob.length));
+                    p.inc(length);
                 }
                 Ok(())
             })?;