[thin/cache_restore (rust)] Build the metadata space map in-place

That avoids cloning the source space map

src/cache/restore.rs

@@ -3,7 +3,6 @@ use anyhow::{anyhow, Result};
 use std::convert::TryInto;
 use std::fs::OpenOptions;
 use std::io::Cursor;
-use std::ops::Deref;
 use std::path::Path;
 use std::sync::Arc;
 
@@ -230,8 +229,7 @@ impl<'a> MetadataVisitor for Restorer<'a> {
 fn build_metadata_sm(w: &mut WriteBatcher) -> Result<Vec<u8>> {
     let mut sm_root = vec![0u8; SPACE_MAP_ROOT_SIZE];
     let mut cur = Cursor::new(&mut sm_root);
-    let sm_without_meta = clone_space_map(w.sm.lock().unwrap().deref())?;
-    let r = write_metadata_sm(w, sm_without_meta.deref())?;
+    let r = write_metadata_sm(w)?;
     r.pack(&mut cur)?;
 
     Ok(sm_root)

src/pdata/space_map.rs

@@ -152,16 +152,6 @@ pub fn core_sm_without_mutex(nr_entries: u64, max_count: u32) -> Box<dyn SpaceMa
     }
 }
 
-// FIXME: replace it by using the Clone trait
-pub fn clone_space_map(src: &dyn SpaceMap) -> Result<Box<dyn SpaceMap>> {
-    let nr_blocks = src.get_nr_blocks()?;
-    let mut dest = Box::new(CoreSpaceMap::<u32>::new(nr_blocks));
-    for i in 0..nr_blocks {
-        dest.set(i, src.get(i)?)?;
-    }
-    Ok(dest)
-}
-
 //------------------------------------------
 
 // This in core space map can only count to one, useful when walking

src/pdata/space_map_common.rs

@@ -250,6 +250,90 @@ pub fn write_common(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<(Vec<Inde
     Ok((index_entries, ref_count_root))
 }
 
+pub fn write_metadata_common(w: &mut WriteBatcher) -> Result<(Vec<IndexEntry>, u64)> {
+    use BitmapEntry::*;
+
+    let mut index_entries = Vec::new();
+    let mut overflow_builder: BTreeBuilder<u32> = BTreeBuilder::new(Box::new(NoopRC {}));
+
+    // how many bitmaps do we need?
+    let nr_blocks = w.sm.lock().unwrap().get_nr_blocks()?;
+    let nr_bitmaps = div_up(nr_blocks, ENTRIES_PER_BITMAP as u64) as usize;
+
+    // how many blocks are allocated or reserved so far?
+    let reserved = w.get_reserved_range();
+    if reserved.end < reserved.start {
+        return Err(anyhow!("unsupported allocation pattern"));
+    }
+    let nr_used_bitmaps = div_up(reserved.end, ENTRIES_PER_BITMAP as u64) as usize;
+
+    for bm in 0..nr_used_bitmaps {
+        let begin = bm as u64 * ENTRIES_PER_BITMAP as u64;
+        let len = std::cmp::min(nr_blocks - begin, ENTRIES_PER_BITMAP as u64);
+        let mut entries = Vec::with_capacity(ENTRIES_PER_BITMAP);
+        let mut first_free: Option<u32> = None;
+
+        // blocks beyond the limit won't be checked right now, thus are marked as freed
+        let limit = std::cmp::min(reserved.end - begin, ENTRIES_PER_BITMAP as u64);
+        let mut nr_free: u32 = (len - limit) as u32;
+
+        for i in 0..limit {
+            let b = begin + i;
+            let rc = w.sm.lock().unwrap().get(b)?;
+            let e = match rc {
+                0 => {
+                    nr_free += 1;
+                    if first_free.is_none() {
+                        first_free = Some(i as u32);
+                    }
+                    Small(0)
+                }
+                1 => Small(1),
+                2 => Small(2),
+                _ => {
+                    overflow_builder.push_value(w, b as u64, rc)?;
+                    Overflow
+                }
+            };
+            entries.push(e);
+        }
+
+        // Fill unused entries with zeros
+        if limit < len {
+            entries.resize_with(len as usize, || BitmapEntry::Small(0));
+        }
+
+        let blocknr = write_bitmap(w, entries)?;
+
+        // Insert into the index list
+        let ie = IndexEntry {
+            blocknr,
+            nr_free,
+            none_free_before: first_free.unwrap_or(limit as u32),
+        };
+        index_entries.push(ie);
+    }
+
+    // Fill the rest of the bitmaps with zeros
+    for bm in nr_used_bitmaps..nr_bitmaps {
+        let begin = bm as u64 * ENTRIES_PER_BITMAP as u64;
+        let len = std::cmp::min(nr_blocks - begin, ENTRIES_PER_BITMAP as u64);
+        let entries = vec![BitmapEntry::Small(0); ENTRIES_PER_BITMAP];
+        let blocknr = write_bitmap(w, entries)?;
+
+        // Insert into the index list
+        let ie = IndexEntry {
+            blocknr,
+            nr_free: len as u32,
+            none_free_before: 0,
+        };
+        index_entries.push(ie);
+    }
+
+    let ref_count_root = overflow_builder.complete(w)?;
+    Ok((index_entries, ref_count_root))
+}
+
 fn write_bitmap(w: &mut WriteBatcher, entries: Vec<BitmapEntry>) -> Result<u64> {
     // allocate a new block
     let b = w.alloc_zeroed()?;

src/pdata/space_map_metadata.rs

@@ -6,7 +6,6 @@ use std::io::Cursor;
 
 use crate::checksum;
 use crate::io_engine::*;
-use crate::pdata::space_map::*;
 use crate::pdata::space_map_common::*;
 use crate::pdata::unpack::*;
 use crate::write_batcher::*;
@@ -103,9 +102,10 @@ fn adjust_counts(w: &mut WriteBatcher, ie: &IndexEntry, allocs: &[u64]) -> Resul
     })
 }
 
-pub fn write_metadata_sm(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<SMRoot> {
+pub fn write_metadata_sm(w: &mut WriteBatcher) -> Result<SMRoot> {
     w.clear_allocations();
-    let (mut indexes, ref_count_root) = write_common(w, sm)?;
+
+    let (mut indexes, ref_count_root) = write_metadata_common(w)?;
 
     let bitmap_root = w.alloc_zeroed()?;
 
@@ -135,6 +135,7 @@ pub fn write_metadata_sm(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<SMRo
     w.write(bitmap_root, checksum::BT::INDEX)?;
     w.flush()?;
 
+    let sm = w.sm.lock().unwrap();
     Ok(SMRoot {
         nr_blocks: sm.get_nr_blocks()?,
         nr_allocated: sm.get_nr_allocated()?,

src/thin/restore.rs

@@ -251,8 +251,7 @@ fn build_data_sm(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<Vec<u8>> {
 fn build_metadata_sm(w: &mut WriteBatcher) -> Result<Vec<u8>> {
     let mut sm_root = vec![0u8; SPACE_MAP_ROOT_SIZE];
     let mut cur = Cursor::new(&mut sm_root);
-    let sm_without_meta = clone_space_map(w.sm.lock().unwrap().deref())?;
-    let r = write_metadata_sm(w, sm_without_meta.deref())?;
+    let r = write_metadata_sm(w)?;
     r.pack(&mut cur)?;
 
     Ok(sm_root)

src/write_batcher.rs

@@ -114,6 +114,13 @@ impl WriteBatcher {
         tmp
     }
 
+    pub fn get_reserved_range(&self) -> std::ops::Range<u64> {
+        std::ops::Range {
+            start: self.reserved.start,
+            end: self.reserved.end,
+        }
+    }
+
     pub fn write(&mut self, b: Block, kind: checksum::BT) -> Result<()> {
         checksum::write_checksum(&mut b.get_data(), kind)?;