Lot's of work on thin_restore

src/write_batcher.rs

@@ -1,4 +1,5 @@
 use anyhow::{anyhow, Result};
+use std::collections::BTreeSet;
 use std::sync::{Arc, Mutex};
 
 use crate::checksum;
@@ -10,10 +11,13 @@ use crate::pdata::space_map::*;
 #[derive(Clone)]
 pub struct WriteBatcher {
     pub engine: Arc<dyn IoEngine + Send + Sync>,
+
+    // FIXME: this doesn't need to be in a mutex
     pub sm: Arc<Mutex<dyn SpaceMap>>,
 
     batch_size: usize,
     queue: Vec<Block>,
+    allocations: BTreeSet<u64>,
 }
 
 impl WriteBatcher {
@@ -27,10 +31,11 @@ impl WriteBatcher {
             sm,
             batch_size,
             queue: Vec::with_capacity(batch_size),
+            allocations: BTreeSet::new(),
         }
     }
 
-    pub fn alloc(&mut self) -> Result<u64> {
+    pub fn alloc(&mut self) -> Result<Block> {
         let mut sm = self.sm.lock().unwrap();
         let b = sm.alloc()?;
 
@@ -38,23 +43,37 @@ impl WriteBatcher {
             return Err(anyhow!("out of metadata space"));
         }
 
-        Ok(b.unwrap())
+        Ok(Block::new(b.unwrap()))
+    }
+
+    pub fn clear_allocations(&mut self) -> BTreeSet<u64> {
+        let mut tmp = BTreeSet::new();
+        std::mem::swap(&mut tmp, &mut self.allocations);
+        tmp
     }
 
     pub fn write(&mut self, b: Block, kind: checksum::BT) -> Result<()> {
         checksum::write_checksum(&mut b.get_data(), kind)?;
 
         if self.queue.len() == self.batch_size {
-            self.flush()?;
+            let mut tmp = Vec::new();
+            std::mem::swap(&mut tmp, &mut self.queue);
+            self.flush_(tmp)?;
         }
 
         self.queue.push(b);
         Ok(())
     }
 
+    pub fn flush_(&mut self, queue: Vec<Block>) -> Result<()> {
+        self.engine.write_many(&queue)?;
+        Ok(())
+    }
+
     pub fn flush(&mut self) -> Result<()> {
-        self.engine.write_many(&self.queue)?;
-        self.queue.clear();
+        let mut tmp = Vec::new();
+        std::mem::swap(&mut tmp, &mut self.queue);
+        self.flush_(tmp)?;
         Ok(())
     }
 }