[thin_check (rust)] change BTreeWalker to use a space map rather than seen bitset

src/pdata/btree.rs

@@ -1,11 +1,11 @@
 use anyhow::{anyhow, Result};
-use fixedbitset::FixedBitSet;
 use nom::{number::complete::*, IResult};
 use std::collections::BTreeMap;
 use std::sync::{Arc, Mutex};
 
 use crate::checksum;
 use crate::io_engine::*;
+use crate::pdata::space_map::*;
 use crate::pdata::unpack::*;
 
 // FIXME: check that keys are in ascending order between nodes.
@@ -153,7 +153,7 @@ pub trait NodeVisitor<V: Unpack> {
 #[derive(Clone)]
 pub struct BTreeWalker {
     pub engine: Arc<dyn IoEngine + Send + Sync>,
-    pub seen: Arc<Mutex<FixedBitSet>>,
+    pub sm: Arc<Mutex<dyn SpaceMap + Send + Sync>>,
     ignore_non_fatal: bool,
 }
 
@@ -162,37 +162,35 @@ impl BTreeWalker {
         let nr_blocks = engine.get_nr_blocks() as usize;
         let r: BTreeWalker = BTreeWalker {
             engine,
-            seen: Arc::new(Mutex::new(FixedBitSet::with_capacity(nr_blocks))),
+            sm: Arc::new(Mutex::new(RestrictedSpaceMap::new(nr_blocks as u64))),
             ignore_non_fatal,
         };
         r
     }
 
-    pub fn new_with_seen(
+    pub fn new_with_sm(
         engine: Arc<dyn IoEngine + Send + Sync>,
-        seen: Arc<Mutex<FixedBitSet>>,
+        sm: Arc<Mutex<dyn SpaceMap + Send + Sync>>,
         ignore_non_fatal: bool,
-    ) -> BTreeWalker {
+    ) -> Result<BTreeWalker> {
         {
-            let seen = seen.lock().unwrap();
+            let sm = sm.lock().unwrap();
-            assert_eq!(seen.len(), engine.get_nr_blocks() as usize);
+            assert_eq!(sm.get_nr_blocks()?, engine.get_nr_blocks());
         }
 
-        BTreeWalker {
+        Ok(BTreeWalker {
             engine,
-            seen,
+            sm,
             ignore_non_fatal,
-        }
+        })
     }
 
-    fn is_seen(&self, b: u64) -> bool {
+    // Atomically increments the ref count, and returns the _old_ count.
-        let mut seen = self.seen.lock().unwrap();
+    fn sm_inc(&self, b: u64) -> Result<u32> {
-        if !seen[b as usize] {
+        let mut sm = self.sm.lock().unwrap();
-            seen.insert(b as usize);
+        let count = sm.get(b)?;
-            return false;
+        sm.inc(b, 1)?;
-        }
+        Ok(count)
-
-        true
     }
 
     fn walk_nodes<NV, V>(&mut self, visitor: &mut NV, bs: &[u64]) -> Result<()>
@@ -201,14 +199,11 @@ impl BTreeWalker {
         V: Unpack,
     {
         let mut blocks = Vec::new();
-        let mut seen = self.seen.lock().unwrap();
         for b in bs {
-            if !seen[*b as usize] {
+            if self.sm_inc(*b)? == 0 {
                 blocks.push(Block::new(*b));
-                seen.insert(*b as usize);
             }
         }
-        drop(seen);
 
         self.engine.read_many(&mut blocks)?;
 
@@ -249,7 +244,7 @@ impl BTreeWalker {
         NV: NodeVisitor<V>,
         V: Unpack,
     {
-        if self.is_seen(root.loc) {
+        if self.sm_inc(root.loc)? > 0 {
             Ok(())
         } else {
             self.walk_node(visitor, &root, true)
@@ -261,7 +256,7 @@ impl BTreeWalker {
         NV: NodeVisitor<V>,
         V: Unpack,
     {
-        if self.is_seen(root) {
+        if self.sm_inc(root)? > 0 {
             Ok(())
         } else {
             let mut root = Block::new(root);
@@ -316,4 +311,17 @@ pub fn btree_to_map<V: Unpack + Clone>(
     Ok(visitor.values)
 }
 
+pub fn btree_to_map_with_sm<V: Unpack + Clone>(
+    engine: Arc<dyn IoEngine + Send + Sync>,
+    sm: Arc<Mutex<dyn SpaceMap + Send + Sync>>,
+    ignore_non_fatal: bool,
+    root: u64,
+) -> Result<BTreeMap<u64, V>> {
+    let mut walker = BTreeWalker::new_with_sm(engine, sm, ignore_non_fatal)?;
+    let mut visitor = ValueCollector::<V>::new();
+
+    walker.walk(&mut visitor, root)?;
+    Ok(visitor.values)
+}
+
 //------------------------------------------

src/pdata/space_map.rs

@@ -1,4 +1,5 @@
 use anyhow::{anyhow, Result};
+use fixedbitset::FixedBitSet;
 use nom::{number::complete::*, IResult};
 use std::sync::{Arc, Mutex};
 
@@ -154,10 +155,13 @@ impl Unpack for Bitmap {
 //------------------------------------------
 
 pub trait SpaceMap {
+    fn get_nr_blocks(&self) -> Result<u64>;
     fn get(&self, b: u64) -> Result<u32>;
     fn inc(&mut self, begin: u64, len: u64) -> Result<()>;
 }
 
+//------------------------------------------
+
 pub struct CoreSpaceMap<T> {
     counts: Vec<T>,
 }
@@ -177,6 +181,10 @@ impl<V> SpaceMap for CoreSpaceMap<V>
 where
     V: Copy + Default + std::ops::AddAssign + From<u8> + Into<u32>,
 {
+    fn get_nr_blocks(&self) -> Result<u64> {
+        Ok(self.counts.len() as u64)
+    }
+    
     fn get(&self, b: u64) -> Result<u32> {
         Ok(self.counts[b as usize].into())
     }
@@ -189,7 +197,7 @@ where
     }
 }
 
-pub fn core_sm(nr_entries: u64, max_count: u32) -> Arc<Mutex<dyn SpaceMap + Send>> {
+pub fn core_sm(nr_entries: u64, max_count: u32) -> Arc<Mutex<dyn SpaceMap + Send + Sync>> {
     if max_count <= u8::MAX as u32 {
         Arc::new(Mutex::new(CoreSpaceMap::<u8>::new(nr_entries)))
     } else if max_count <= u16::MAX as u32 {
@@ -200,3 +208,41 @@ pub fn core_sm(nr_entries: u64, max_count: u32) -> Arc<Mutex<dyn SpaceMap + Send
 }
 
 //------------------------------------------
+
+// This in core space map can only count to one, useful when walking
+// btrees when we want to avoid visiting a node more than once, but
+// aren't interested in counting how many times we've visited.
+pub struct RestrictedSpaceMap {
+    counts: FixedBitSet,
+}
+
+impl RestrictedSpaceMap {
+    pub fn new(nr_entries: u64) -> RestrictedSpaceMap {
+        RestrictedSpaceMap {
+            counts: FixedBitSet::with_capacity(nr_entries as usize),
+        }
+    }
+}
+
+impl SpaceMap for RestrictedSpaceMap {
+    fn get_nr_blocks(&self) -> Result<u64> {
+        Ok(self.counts.len() as u64)
+    }
+    
+    fn get(&self, b: u64) -> Result<u32> {
+        if self.counts.contains(b as usize) {
+            Ok(1)
+        } else {
+            Ok(0)
+        }
+    }
+
+    fn inc(&mut self, begin: u64, len: u64) -> Result<()> {
+        for b in begin..(begin + len) {
+            self.counts.insert(b as usize);
+        }
+        Ok(())
+    }
+}
+
+//------------------------------------------

src/thin/check.rs

@@ -1,5 +1,4 @@
 use anyhow::{anyhow, Result};
-use fixedbitset::FixedBitSet;
 use nom::{number::complete::*, IResult};
 use std::collections::BTreeMap;
 use std::path::Path;
@@ -8,9 +7,9 @@ use threadpool::ThreadPool;
 
 use crate::checksum;
 use crate::io_engine::{AsyncIoEngine, Block, IoEngine, SyncIoEngine};
-use crate::pdata::unpack::*;
+use crate::pdata::btree::{btree_to_map, btree_to_map_with_sm, BTreeWalker, Node, NodeVisitor};
-use crate::pdata::btree::{btree_to_map, BTreeWalker, Node, NodeVisitor};
 use crate::pdata::space_map::*;
+use crate::pdata::unpack::*;
 use crate::thin::superblock::*;
 
 //------------------------------------------
@@ -200,11 +199,26 @@ pub fn check(opts: &ThinCheckOptions) -> Result<()> {
     // superblock
     let sb = read_superblock(engine.as_ref(), SUPERBLOCK_LOCATION)?;
 
-    // device details
+    // Device details.   We read this once to get the number of thin devices, and hence the
+    // maximum metadata ref count.  Then create metadata space map, and reread to increment
+    // the ref counts for that metadata.
     let devs = btree_to_map::<DeviceDetail>(engine.clone(), false, sb.details_root)?;
     let nr_devs = devs.len();
+    let metadata_sm = core_sm(engine.get_nr_blocks(), nr_devs as u32);
+    let _devs = btree_to_map_with_sm::<DeviceDetail>(
+        engine.clone(),
+        metadata_sm.clone(),
+        false,
+        sb.details_root,
+    )?;
     println!("found {} devices", nr_devs);
 
+    // increment superblock
+    {
+        let mut sm = metadata_sm.lock().unwrap();
+        sm.inc(SUPERBLOCK_LOCATION, 1)?;
+    }
+    
     // mapping top level
     let roots = btree_to_map::<u64>(engine.clone(), false, sb.mapping_root)?;
 
@@ -214,15 +228,12 @@ pub fn check(opts: &ThinCheckOptions) -> Result<()> {
         // FIXME: with a thread pool we need to return errors another way.
         let nr_workers = nr_threads;
         let pool = ThreadPool::new(nr_workers);
-        let seen = Arc::new(Mutex::new(FixedBitSet::with_capacity(
-            engine.get_nr_blocks() as usize,
-        )));
 
         let root = unpack::<SMRoot>(&sb.data_sm_root[0..])?;
         data_sm = core_sm(root.nr_blocks, nr_devs as u32);
 
         for (thin_id, root) in roots {
-            let mut w = BTreeWalker::new_with_seen(engine.clone(), seen.clone(), false);
+            let mut w = BTreeWalker::new_with_sm(engine.clone(), metadata_sm.clone(), false)?;
             let data_sm = data_sm.clone();
             pool.execute(move || {
                 let mut v = BottomLevelVisitor { data_sm };
@@ -325,7 +336,6 @@ pub fn check(opts: &ThinCheckOptions) -> Result<()> {
     }
 
     // Check the metadata space map.
-    
 
     Ok(())
 }