From c142cd0d485e2bea2a45205ce1eb03865634d6c6 Mon Sep 17 00:00:00 2001
From: Ming-Hung Tsai <mtsai@redhat.com>
Date: Wed, 26 May 2021 01:41:37 +0800
Subject: [PATCH] [space_map (rust)] Fix space map building

- Fix out-of-bounds index
- Automatically flush queued writes before function return
- Track allocated blocks in write_batcher (might be space consuming)
---
 src/pdata/space_map_common.rs   | 2 +-
 src/pdata/space_map_disk.rs     | 1 +
 src/pdata/space_map_metadata.rs | 1 +
 src/write_batcher.rs            | 1 +
 4 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/pdata/space_map_common.rs b/src/pdata/space_map_common.rs
index 33e3a51..e2f8918 100644
--- a/src/pdata/space_map_common.rs
+++ b/src/pdata/space_map_common.rs
@@ -209,7 +209,7 @@ pub fn write_common(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<(Vec<Inde
         let mut nr_free: u32 = 0;
         for i in 0..ENTRIES_PER_BITMAP {
             let b: u64 = ((bm * ENTRIES_PER_BITMAP) as u64) + i as u64;
-            if b > sm.get_nr_blocks()? {
+            if b >= sm.get_nr_blocks()? {
                 break;
             }
             let rc = sm.get(b)?;
diff --git a/src/pdata/space_map_disk.rs b/src/pdata/space_map_disk.rs
index 1389d78..88e1749 100644
--- a/src/pdata/space_map_disk.rs
+++ b/src/pdata/space_map_disk.rs
@@ -16,6 +16,7 @@ pub fn write_disk_sm(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<SMRoot>
     }
 
     let bitmap_root = index_builder.complete(w)?;
+    w.flush()?;
 
     Ok(SMRoot {
         nr_blocks: sm.get_nr_blocks()?,
diff --git a/src/pdata/space_map_metadata.rs b/src/pdata/space_map_metadata.rs
index 85f6828..4567b05 100644
--- a/src/pdata/space_map_metadata.rs
+++ b/src/pdata/space_map_metadata.rs
@@ -126,6 +126,7 @@ pub fn write_metadata_sm(w: &mut WriteBatcher, sm: &dyn SpaceMap) -> Result<SMRo
     metadata_index.pack(&mut cur)?;
     let loc = bitmap_root.loc;
     w.write(bitmap_root, checksum::BT::INDEX)?;
+    w.flush()?;
 
     Ok(SMRoot {
         nr_blocks: sm.get_nr_blocks()?,
diff --git a/src/write_batcher.rs b/src/write_batcher.rs
index b911be5..74a04ec 100644
--- a/src/write_batcher.rs
+++ b/src/write_batcher.rs
@@ -43,6 +43,7 @@ impl WriteBatcher {
             return Err(anyhow!("out of metadata space"));
         }
 
+        self.allocations.insert(b.unwrap());
         Ok(Block::new(b.unwrap()))
     }