Fix a problem with restore when used with --pick.

Our simple offset system was not correct for all cases, revert to the slower but simpler solution... Also add more tests.
andrewchambers · Aug 2, 2021 · 53519c0 · 53519c0
1 parent 41e07af
commit 53519c0
Show file tree

Hide file tree

Showing 6 changed files with 154 additions and 41 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "bupstash"
-version = "0.10.1"
+version = "0.10.2"
 authors = ["Andrew Chambers <ac@acha.ninja>"]
 edition = "2018"
 license = "MIT"

diff --git a/cli-tests/cli-tests.bats b/cli-tests/cli-tests.bats
@@ -774,13 +774,130 @@ _concurrent_modify_worker () {
   test 0 = "$(bupstash diff --relaxed $SCRATCH/d/b :: $SCRATCH/restore | expr $(wc -l))"
 }
 
-@test "pick fuzz torture" {
+@test "restore torture" {
+
+  if test -z "$BUPSTASH_TORTURE_DIR"
+  then
+    skip "Set BUPSTASH_TORTURE_DIR to run this test."
+  fi
+
+  restore_dir="$SCRATCH/restore_dir"
+  copy_dir="$SCRATCH/copy_dir"
+  # Put twice so we test caching code paths.
+  id1=$(bupstash put :: "$BUPSTASH_TORTURE_DIR")
+  id2=$(bupstash put :: "$BUPSTASH_TORTURE_DIR")
+
+  for id in $(echo $id1 $id2)
+  do
+    for d in $(cd "$BUPSTASH_TORTURE_DIR" ; find . -type d | sed 's,^\./,,g')
+    do
+      rm -rf "$copy_dir" "$restore_dir"
+      mkdir "$copy_dir" "$restore_dir"
+
+      tar -C "$BUPSTASH_TORTURE_DIR/$d" -cf - . | tar -C "$copy_dir" -xf -
+      bupstash restore --into "$restore_dir" --pick "$d" "id=$id"
+
+      diff -u \
+        <(cd "$restore_dir" ; find . | sort) \
+        <(cd "$copy_dir" ; find . | sort)
+
+      for f in $(cd "$copy_dir" ; find . -type f | cut -c 3-)
+      do
+        cmp "$copy_dir/$f" "$restore_dir/$f"
+      done
+    done
+  done
+}
+
+@test "restore fuzz torture" {
+
+  rand_dir="$SCRATCH/random_dir"
+  restore_dir="$SCRATCH/restore_dir"
+  copy_dir="$SCRATCH/copy_dir"
+
+  for i in `seq 50`
+  do
+    rm -rf "$rand_dir"
+
+    "$BATS_TEST_DIRNAME/mk-random-dir.py" "$rand_dir"
+
+    # Put twice so we test caching code paths.
+    id1=$(bupstash put :: "$rand_dir")
+    id2=$(bupstash put :: "$rand_dir")
+
+    for id in $(echo $id1 $id2)
+    do
+      for d in $(cd "$rand_dir" ; find . -type d | sed 's,^\./,,g')
+      do
+        rm -rf "$copy_dir" "$restore_dir"
+        mkdir "$copy_dir" "$restore_dir"
+
+        tar -C "$rand_dir/$d" -cf - . | tar -C "$copy_dir" -xf -
+        bupstash restore --into "$restore_dir" --pick "$d" "id=$id"
+
+        diff -u \
+          <(cd "$restore_dir" ; find . | sort) \
+          <(cd "$copy_dir" ; find . | sort)
+
+        for f in $(cd "$copy_dir" ; find . -type f | cut -c 3-)
+        do
+          cmp "$copy_dir/$f" "$restore_dir/$f"
+        done
+      done
+      bupstash rm id=$id
+    done
+
+    bupstash gc
+  done
+}
+
+@test "get pick torture" {
+
+  if test -z "$BUPSTASH_TORTURE_DIR"
+  then
+    skip "Set BUPSTASH_TORTURE_DIR to run this test."
+  fi
+
+  restore_dir="$SCRATCH/restore_dir"
+  copy_dir="$SCRATCH/copy_dir"
+  # Put twice so we test caching code paths.
+  id1=$(bupstash put :: "$BUPSTASH_TORTURE_DIR")
+  id2=$(bupstash put :: "$BUPSTASH_TORTURE_DIR")
+
+  for id in $(echo $id1 $id2)
+  do
+    for f in $(cd "$BUPSTASH_TORTURE_DIR" ; find . -type f | cut -c 3-)
+    do
+      cmp <(bupstash get --pick "$f"  "id=$id") "$BUPSTASH_TORTURE_DIR/$f"
+    done
+
+    for d in $(cd "$BUPSTASH_TORTURE_DIR" ; find . -type d | sed 's,^\./,,g')
+    do
+      rm -rf "$copy_dir" "$restore_dir"
+      mkdir "$copy_dir" "$restore_dir"
+
+      tar -C "$BUPSTASH_TORTURE_DIR/$d" -cf - . | tar -C "$copy_dir" -xf -
+      bupstash get --pick "$d" "id=$id" | tar -C "$restore_dir" -xf -
+
+      diff -u \
+        <(cd "$restore_dir" ; find . | sort) \
+        <(cd "$copy_dir" ; find . | sort)
+
+      for f in $(cd "$copy_dir" ; find . -type f | cut -c 3-)
+      do
+        cmp "$copy_dir/$f" "$restore_dir/$f"
+      done
+    done
+  done
+}
+
+@test "get pick fuzz torture" {
 
   rand_dir="$SCRATCH/random_dir"
   restore_dir="$SCRATCH/restore_dir"
   copy_dir="$SCRATCH/copy_dir"
 
-  for i in `seq 100`
+  for i in `seq 50`
   do
     rm -rf "$rand_dir"
 

diff --git a/src/client.rs b/src/client.rs
@@ -1377,11 +1377,24 @@ pub fn restore_to_local_dir(
     progress: &indicatif::ProgressBar,
     ctx: RestoreContext,
     content_index: index::CompressedIndex,
-    data_map: Option<index::DataMap>,
+    pick: Option<String>,
     serve_out: &mut dyn std::io::Read,
     serve_in: &mut dyn std::io::Write,
     to_dir: &Path,
 ) -> Result<(), anyhow::Error> {
+    let sub_index = match pick {
+        Some(ref pick) => {
+            progress.set_message("building pick index...");
+            Some(index::pick_dir_without_data(pick, &content_index)?)
+        }
+        None => None,
+    };
+
+    let index_to_diff = match sub_index {
+        Some(ref sub_index) => sub_index,
+        None => &content_index,
+    };
+
     // Initially reset the permissions and groups on everything
     // so we don't need to worry about read only files or other access.
     progress.set_message("preparing directory...");
@@ -1442,7 +1455,7 @@ pub fn restore_to_local_dir(
         ciw.finish()
     };
 
-    progress.set_message("computing content diff...");
+    progress.set_message("calculating content diff...");
     let mut to_remove = Vec::with_capacity(512);
     let mut new_dirs = Vec::with_capacity(512);
     let mut create_path_set = HashSet::with_capacity(512);
@@ -1451,9 +1464,15 @@ pub fn restore_to_local_dir(
     let mut downloads = Vec::with_capacity(512);
 
     {
+        let download_path_prefix = match pick {
+            Some(ref pick) if pick == "." => "".to_string(),
+            Some(ref pick) => format!("{}/", pick),
+            None => "".to_string(),
+        };
+
         index::diff(
             &to_dir_index,
-            &content_index,
+            index_to_diff,
             !(index::INDEX_COMPARE_MASK_TYPE
                 | index::INDEX_COMPARE_MASK_LINK_TARGET
                 | index::INDEX_COMPARE_MASK_DEVNOS
@@ -1468,7 +1487,9 @@ pub fn restore_to_local_dir(
                         if e.is_dir() {
                             new_dirs.push(PathBuf::from(&e.path));
                         } else if e.is_file() {
-                            download_index_path_set.insert(e.path.clone());
+                            // Mapping back to the path in the unpicked index.
+                            download_index_path_set
+                                .insert(format!("{}{}", download_path_prefix, e.path));
                             downloads.push((e.path.clone(), e.size.0, e.data_hash));
                         } else {
                             create_path_set.insert(e.path.clone());
@@ -1512,17 +1533,15 @@ pub fn restore_to_local_dir(
     std::mem::drop(new_dirs);
 
     if !download_index_path_set.is_empty() {
-        progress.set_message("fetching files...");
+        progress.set_message("calculating fetch set...");
 
-        let mut fetch_data_map = index::data_map_for_predicate(&content_index, &|e| {
+        let fetch_data_map = index::data_map_for_predicate(&content_index, &|e| {
             download_index_path_set.contains(&e.path)
         })?;
 
-        if let Some(data_map) = data_map {
-            if let Some(start_range) = data_map.data_chunk_ranges.first() {
-                fetch_data_map.add_offset(start_range.start_idx.0);
-            }
-        }
+        std::mem::drop(download_index_path_set);
+
+        progress.set_message("fetching files...");
 
         let (mut r, mut w) = ioutil::buffered_pipe(3 * 1024 * 1024); // Sized to cover most packets in one allocation.
 
@@ -1584,7 +1603,6 @@ pub fn restore_to_local_dir(
             Err(err) => return Err(err.into()),
         };
     }
-    std::mem::drop(download_index_path_set);
 
     if !create_path_set.is_empty() {
         progress.set_message("creating special files and devices...");
@@ -1722,7 +1740,7 @@ pub fn restore_to_local_dir(
     {
         index::diff(
             &to_dir_index,
-            &content_index,
+            index_to_diff,
             !(index::INDEX_COMPARE_MASK_PERMS | index::INDEX_COMPARE_MASK_XATTRS),
             &mut |ds: index::DiffStat, ent: &index::IndexEntry| -> Result<(), anyhow::Error> {
                 if matches!(ds, index::DiffStat::Removed) {

diff --git a/src/index.rs b/src/index.rs
@@ -307,25 +307,12 @@ pub struct HTreeDataRange {
     pub end_idx: serde_bare::Uint,
 }
 
+#[derive(Debug)]
 pub struct DataMap {
     pub data_chunk_ranges: Vec<HTreeDataRange>,
     pub incomplete_data_chunks: HashMap<u64, rangemap::RangeSet<usize>>,
 }
 
-impl DataMap {
-    pub fn add_offset(&mut self, offset: u64) {
-        self.incomplete_data_chunks = self
-            .incomplete_data_chunks
-            .drain()
-            .map(|(k, v)| (k + offset, v))
-            .collect();
-        for range in self.data_chunk_ranges.iter_mut() {
-            range.start_idx.0 += offset;
-            range.end_idx.0 += offset;
-        }
-    }
-}
-
 fn add_ent_to_data_map(
     ent: &IndexEntry,
     cur_chunk_idx: u64,

diff --git a/src/main.rs b/src/main.rs
@@ -2306,15 +2306,6 @@ fn restore_main(args: Vec<String>) -> Result<(), anyhow::Error> {
         anyhow::bail!("restore is only supported for directory snapshots created by bupstash");
     };
 
-    let (content_index, data_map) = if let Some(ref pick_path) = matches.opt_str("pick") {
-        match index::pick(pick_path, &content_index)? {
-            (Some(content_index), data_map) => (content_index, Some(data_map)),
-            _ => anyhow::bail!("the given pick was not a directory"),
-        }
-    } else {
-        (content_index, None)
-    };
-
     client::restore_to_local_dir(
         &progress,
         client::RestoreContext {
@@ -2330,7 +2321,7 @@ fn restore_main(args: Vec<String>) -> Result<(), anyhow::Error> {
             restore_xattrs: matches.opt_present("xattrs"),
         },
         content_index,
-        data_map,
+        matches.opt_str("pick"),
         &mut serve_out,
         &mut serve_in,
         &into_dir,