[thin_restore] first pass at btree_builder.

No tests yet

src/bin/thin_explore.rs

@@ -301,7 +301,7 @@ impl Adjacent for BlockTime {
 }
 
 impl Adjacent for DeviceDetail {
-    fn adjacent(&self, rhs: &Self) -> bool {
+    fn adjacent(&self, _rhs: &Self) -> bool {
         false
     }
 }
@@ -561,7 +561,7 @@ impl Panel for DeviceDetailPanel {
                 btree::Node::Internal { values, .. } => {
                     Some(PushDeviceDetail(values[self.state.selected().unwrap()]))
                 }
-                btree::Node::Leaf { values, keys, .. } => None,
+                btree::Node::Leaf { .. } => None,
             },
             Key::Char('h') | Key::Left => Some(PopPanel),
             _ => None,

src/bin/thin_restore.rs

@@ -0,0 +1,75 @@
+extern crate clap;
+extern crate thinp;
+
+use atty::Stream;
+use clap::{App, Arg};
+use std::path::Path;
+use std::process;
+use std::process::exit;
+use std::sync::Arc;
+use thinp::file_utils;
+use thinp::report::*;
+use thinp::thin::restore::{restore, ThinRestoreOptions};
+
+fn main() {
+    let parser = App::new("thin_restore")
+        .version(thinp::version::TOOLS_VERSION)
+        .about("Convert XML format metadata to binary.")
+        .arg(
+            Arg::with_name("OVERRIDE_MAPPING_ROOT")
+                .help("Specify a mapping root to use")
+                .long("override-mapping-root")
+                .value_name("OVERRIDE_MAPPING_ROOT")
+                .takes_value(true),
+        )
+        .arg(
+            Arg::with_name("INPUT")
+                .help("Specify the input xml")
+                .short("i")
+                .long("input")
+                .required(true),
+        )
+        .arg(
+            Arg::with_name("OUTPUT")
+                .help("Specify the output device to check")
+                .short("o")
+                .long("output")
+                .required(true),
+        )
+        .arg(
+            Arg::with_name("SYNC_IO")
+                .help("Force use of synchronous io")
+                .long("sync-io"),
+        );
+
+    let matches = parser.get_matches();
+    let input_file = Path::new(matches.value_of("INPUT").unwrap());
+    let output_file = Path::new(matches.value_of("OUTPUT").unwrap());
+
+    if !file_utils::file_exists(input_file) {
+        eprintln!("Couldn't find input file '{:?}'.", &input_file);
+        exit(1);
+    }
+
+    let report;
+
+    if matches.is_present("QUIET") {
+        report = std::sync::Arc::new(mk_quiet_report());
+    } else if atty::is(Stream::Stdout) {
+        report = std::sync::Arc::new(mk_progress_bar_report());
+    } else {
+        report = Arc::new(mk_simple_report());
+    }
+
+    let opts = ThinRestoreOptions {
+        input: &input_file,
+        output: &output_file,
+        async_io: !matches.is_present("SYNC_IO"),
+        report,
+    };
+
+    if let Err(reason) = restore(opts) {
+        println!("{}", reason);
+        process::exit(1);
+    }
+}

src/pdata/btree.rs

@@ -461,6 +461,26 @@ impl<V: Unpack> Node<V> {
             Leaf { header, .. } => header,
         }
     }
+
+    fn get_mut_header(&mut self) -> &mut NodeHeader {
+        use Node::*;
+        match self {
+            Internal { header, .. } => header,
+            Leaf { header, .. } => header,
+        }
+    }
+    
+    pub fn get_keys(&self) -> &[u64] {
+        use Node::*;
+        match self {
+            Internal { keys, .. } => &keys[0..],
+            Leaf { keys, .. } => &keys[0..],
+        }
+    }
+
+    pub fn set_block(&mut self, b: u64) {
+        self.get_mut_header().block = b;
+    }
 }
 
 pub fn convert_result<'a, V>(path: &Vec<u64>, r: IResult<&'a [u8], V>) -> Result<(&'a [u8], V)> {

src/pdata/btree_builder.rs

@@ -0,0 +1,364 @@
+use anyhow::{anyhow, Result};
+use byteorder::{LittleEndian, WriteBytesExt};
+use std::collections::VecDeque;
+use std::io::Cursor;
+use std::sync::{Arc, Mutex};
+
+use crate::checksum;
+use crate::io_engine::*;
+use crate::pdata::btree::*;
+use crate::pdata::space_map::*;
+use crate::pdata::unpack::*;
+
+//------------------------------------------
+
+fn pack_node<W: WriteBytesExt, V: Pack + Unpack>(node: &Node<V>, w: &mut W) -> Result<()> {
+    match node {
+        Node::Internal {
+            header,
+            keys,
+            values,
+        } => {
+            header.pack(w)?;
+            for k in keys {
+                w.write_u64::<LittleEndian>(*k)?;
+            }
+
+            // pad with zeroes
+            for _i in keys.len()..header.max_entries as usize {
+                w.write_u64::<LittleEndian>(0)?;
+            }
+
+            for v in values {
+                v.pack(w)?;
+            }
+        }
+        Node::Leaf {
+            header,
+            keys,
+            values,
+        } => {
+            header.pack(w)?;
+            for k in keys {
+                w.write_u64::<LittleEndian>(*k)?;
+            }
+
+            // pad with zeroes
+            for _i in keys.len()..header.max_entries as usize {
+                w.write_u64::<LittleEndian>(0)?;
+            }
+
+            for v in values {
+                v.pack(w)?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+//------------------------------------------
+
+fn calc_max_entries<V: Unpack>() -> usize {
+    let elt_size = 8 + V::disk_size() as usize;
+    ((BLOCK_SIZE - NodeHeader::disk_size() as usize) / elt_size) as usize
+}
+
+//------------------------------------------
+
+struct Entries<V> {
+    max_entries: usize,
+    entries: VecDeque<(u64, V)>,
+}
+
+enum Action<V> {
+    Noop,
+    WriteSingle {
+        keys: Vec<u64>,
+        values: Vec<V>,
+    },
+    WritePair {
+        keys1: Vec<u64>,
+        values1: Vec<V>,
+        keys2: Vec<u64>,
+        values2: Vec<V>,
+    },
+}
+
+impl<V> Entries<V> {
+    pub fn new(max_entries: usize) -> Entries<V> {
+        Entries {
+            max_entries,
+            entries: VecDeque::new(),
+        }
+    }
+
+    pub fn add_entry(&mut self, k: u64, v: V) -> Action<V> {
+        let result = if self.full() {
+            let (keys, values) = self.pop(self.max_entries);
+            Action::WriteSingle { keys, values }
+        } else {
+            Action::Noop
+        };
+
+        self.entries.push_back((k, v));
+
+        result
+    }
+
+    pub fn complete(&mut self) -> Action<V> {
+        let n = self.entries.len();
+
+        if n >= self.max_entries {
+            let n1 = n / 2;
+            let n2 = n - n1;
+            let (keys1, values1) = self.pop(n1);
+            let (keys2, values2) = self.pop(n2);
+
+            Action::WritePair {
+                keys1,
+                values1,
+                keys2,
+                values2,
+            }
+        } else if n > 0 {
+            let (keys, values) = self.pop(n);
+            Action::WriteSingle { keys, values }
+        } else {
+            Action::Noop
+        }
+    }
+
+    fn full(&self) -> bool {
+        self.entries.len() >= 2 * self.max_entries
+    }
+
+    fn pop(&mut self, count: usize) -> (Vec<u64>, Vec<V>) {
+        let mut keys = Vec::new();
+        let mut values = Vec::new();
+
+        for _i in 0..count {
+            let (k, v) = self.entries.pop_front().unwrap();
+            keys.push(k);
+            values.push(v);
+        }
+
+        (keys, values)
+    }
+}
+
+//------------------------------------------
+
+struct WriteBatcher {
+    engine: Arc<Box<dyn IoEngine>>,
+    sm: Arc<Mutex<dyn SpaceMap>>,
+
+    batch_size: usize,
+    queue: Vec<Block>,
+}
+
+impl WriteBatcher {
+    fn new(
+        engine: Arc<Box<dyn IoEngine>>,
+        sm: Arc<Mutex<dyn SpaceMap>>,
+        batch_size: usize,
+    ) -> WriteBatcher {
+        WriteBatcher {
+            engine,
+            sm,
+            batch_size,
+            queue: Vec::with_capacity(batch_size),
+        }
+    }
+
+    fn alloc(&mut self) -> Result<u64> {
+        let mut sm = self.sm.lock().unwrap();
+        let b = sm.alloc()?;
+
+        if b.is_none() {
+            return Err(anyhow!("out of metadata space"));
+        }
+
+        Ok(b.unwrap())
+    }
+
+    fn write(&mut self, b: Block) -> Result<()> {
+        checksum::write_checksum(&mut b.get_data(), checksum::BT::NODE)?;
+        
+        if self.queue.len() == self.batch_size {
+            self.flush()?;
+        }
+
+        self.queue.push(b);
+        Ok(())
+    }
+
+    fn flush(&mut self) -> Result<()> {
+        self.engine.write_many(&self.queue)?;
+        self.queue.clear();
+        Ok(())
+    }
+}
+
+//------------------------------------------
+
+fn write_node_<V: Unpack + Pack>(w: &mut WriteBatcher, mut node: Node<V>) -> Result<(u64, u64)> {
+    let keys = node.get_keys();
+    let first_key = keys.first().unwrap_or(&0u64).clone();
+
+    let loc = w.alloc()?;
+    node.set_block(loc);
+    
+    let b = Block::new(loc);
+    let mut cursor = Cursor::new(b.get_data());
+    pack_node(&node, &mut cursor)?;
+    w.write(b)?;
+
+    Ok((first_key, loc))
+}
+
+fn write_leaf<V: Unpack + Pack>(
+    w: &mut WriteBatcher,
+    keys: Vec<u64>,
+    values: Vec<V>,
+) -> Result<(u64, u64)> {
+    let header = NodeHeader {
+        block: 0,
+        is_leaf: true,
+        nr_entries: keys.len() as u32,
+        max_entries: calc_max_entries::<V>() as u32,
+        value_size: V::disk_size(),
+    };
+
+    let node = Node::Leaf {
+        header,
+        keys,
+        values,
+    };
+
+    write_node_(w, node)
+}
+
+fn write_internal(w: &mut WriteBatcher, keys: Vec<u64>, values: Vec<u64>) -> Result<(u64, u64)> {
+    let header = NodeHeader {
+        block: 0,
+        is_leaf: false,
+        nr_entries: keys.len() as u32,
+        max_entries: calc_max_entries::<u64>() as u32,
+        value_size: u64::disk_size(),
+    };
+
+    let node: Node<u64> = Node::Internal {
+        header,
+        keys,
+        values,
+    };
+
+    write_node_(w, node)
+}
+
+pub struct Builder<V: Unpack + Pack> {
+    w: WriteBatcher,
+    entries: Entries<V>,
+
+    max_internal_entries: usize,
+    internal_entries: Vec<Entries<u64>>,
+
+    root: u64,
+}
+
+impl<V: Unpack + Pack> Builder<V> {
+    pub fn new(engine: Arc<Box<dyn IoEngine>>, sm: Arc<Mutex<dyn SpaceMap>>) -> Builder<V> {
+        let max_entries = calc_max_entries::<V>();
+        let max_internal_entries = calc_max_entries::<u64>();
+
+        Builder {
+            w: WriteBatcher::new(engine, sm, 256),
+            entries: Entries::new(max_entries),
+            max_internal_entries,
+            internal_entries: Vec::new(),
+            root: 0,
+        }
+    }
+
+    pub fn add_entry(&mut self, k: u64, v: V) -> Result<()> {
+        let action = self.entries.add_entry(k, v);
+        self.perform_action(action)
+    }
+
+    pub fn complete(mut self) -> Result<u64> {
+        let action = self.entries.complete();
+        self.perform_action(action)?;
+        self.w.flush()?;
+        Ok(self.root)
+    }
+
+    //--------------------
+
+    fn add_internal_entry(&mut self, level: usize, k: u64, v: u64) -> Result<()> {
+        if self.internal_entries.len() == level {
+            self.internal_entries.push(Entries::new(self.max_internal_entries));
+        }
+
+        let action = self.internal_entries[level].add_entry(k, v);
+        self.perform_internal_action(level, action)
+    }
+
+    fn perform_internal_action(&mut self, level: usize, action: Action<u64>) -> Result<()> {
+        match action {
+            Action::Noop => {}
+            Action::WriteSingle { keys, values } => {
+                let (k, loc) = write_internal(&mut self.w, keys, values)?;
+                self.add_internal_entry(level + 1, k, loc)?;
+                self.root = loc;
+            }
+            Action::WritePair {
+                keys1,
+                values1,
+                keys2,
+                values2,
+            } => {
+                let (k, loc) = write_leaf(&mut self.w, keys1, values1)?;
+                self.add_internal_entry(level + 1, k, loc)?;
+
+                let (k, loc) = write_leaf(&mut self.w, keys2, values2)?;
+                self.add_internal_entry(level + 1, k, loc)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn perform_action<V2: Unpack + Pack>(&mut self, action: Action<V2>) -> Result<()> {
+        match action {
+            Action::Noop => {}
+            Action::WriteSingle { keys, values } => {
+                let (k, loc) = write_leaf(&mut self.w, keys, values)?;
+                self.add_internal_entry(0, k, loc)?;
+            }
+            Action::WritePair {
+                keys1,
+                values1,
+                keys2,
+                values2,
+            } => {
+                let (k, loc) = write_leaf(&mut self.w, keys1, values1)?;
+                self.add_internal_entry(0, k, loc)?;
+
+                let (k, loc) = write_leaf(&mut self.w, keys2, values2)?;
+                self.add_internal_entry(0, k, loc)?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+//------------------------------------------
+
+#[test]
+fn fail() {
+    assert!(false);
+}
+
+//------------------------------------------

src/pdata/mod.rs

@@ -1,4 +1,5 @@
 pub mod btree;
+pub mod btree_builder;
 pub mod space_map;
 pub mod unpack;
 

src/pdata/space_map.rs

@@ -1,8 +1,8 @@
 use anyhow::{anyhow, Result};
+use byteorder::{LittleEndian, WriteBytesExt};
 use fixedbitset::FixedBitSet;
 use nom::{multi::count, number::complete::*, IResult};
 use std::sync::{Arc, Mutex};
-use byteorder::{LittleEndian, WriteBytesExt};
 
 use crate::io_engine::*;
 use crate::pdata::unpack::{Pack, Unpack};
@@ -96,7 +96,7 @@ impl Unpack for MetadataIndex {
         let (i, _blocknr) = le_u64(i)?;
         let (i, indexes) = count(IndexEntry::unpack, MAX_METADATA_BITMAPS)(i)?;
 
-        Ok((i, MetadataIndex {indexes}))
+        Ok((i, MetadataIndex { indexes }))
     }
 }
 
@@ -195,17 +195,16 @@ impl Pack for Bitmap {
             for e in chunk {
                 w >>= 2;
                 match e {
-                    Small(0) => {
-                    },
+                    Small(0) => {}
                     Small(1) => {
                         w |= 0x2 << 62;
-                    },
+                    }
                     Small(2) => {
                         w |= 0x1 << 62;
-                    },
+                    }
                     Small(_) => {
                         return Err(anyhow!("Bad small value in bitmap entry"));
-                    },
+                    }
                     Overflow => {
                         w |= 0x3 << 62;
                     }
@@ -228,8 +227,12 @@ pub trait SpaceMap {
 
     // Returns the old ref count
     fn set(&mut self, b: u64, v: u32) -> Result<u32>;
-    
+
     fn inc(&mut self, begin: u64, len: u64) -> Result<()>;
+
+    // Finds a block with a zero reference count.  Increments the
+    // count.
+    fn alloc(&mut self) -> Result<Option<u64>>;
 }
 
 pub type ASpaceMap = Arc<Mutex<dyn SpaceMap + Sync + Send>>;
@@ -238,6 +241,7 @@ pub type ASpaceMap = Arc<Mutex<dyn SpaceMap + Sync + Send>>;
 
 pub struct CoreSpaceMap<T> {
     nr_allocated: u64,
+    first_free: u64,
     counts: Vec<T>,
 }
 
@@ -248,6 +252,7 @@ where
     pub fn new(nr_entries: u64) -> CoreSpaceMap<V> {
         CoreSpaceMap {
             nr_allocated: 0,
+            first_free: 0,
             counts: vec![V::default(); nr_entries as usize],
         }
     }
@@ -271,13 +276,16 @@ where
 
     fn set(&mut self, b: u64, v: u32) -> Result<u32> {
         let old = self.counts[b as usize];
-        assert!(v < 0xff);   // FIXME: we can't assume this
+        assert!(v < 0xff); // FIXME: we can't assume this
         self.counts[b as usize] = V::from(v as u8);
 
         if old == V::from(0u8) && v != 0 {
             self.nr_allocated += 1;
         } else if old != V::from(0u8) && v == 0 {
             self.nr_allocated -= 1;
+            if b < self.first_free {
+                self.first_free = b;
+            }
         }
 
         Ok(old.into())
@@ -295,6 +303,19 @@ where
         }
         Ok(())
     }
+
+    fn alloc(&mut self) -> Result<Option<u64>> {
+        for b in self.first_free..(self.counts.len() as u64) {
+            if self.counts[b as usize] == V::from(0u8) {
+                self.counts[b as usize] = V::from(1u8);
+                self.first_free = b + 1;
+                return Ok(Some(b));
+            }
+        }
+
+        self.first_free = self.counts.len() as u64;
+        Ok(None)
+    }
 }
 
 pub fn core_sm(nr_entries: u64, max_count: u32) -> Arc<Mutex<dyn SpaceMap + Send + Sync>> {
@@ -314,6 +335,7 @@ pub fn core_sm(nr_entries: u64, max_count: u32) -> Arc<Mutex<dyn SpaceMap + Send
 // aren't interested in counting how many times we've visited.
 pub struct RestrictedSpaceMap {
     nr_allocated: u64,
+    first_free: usize,
     counts: FixedBitSet,
 }
 
@@ -322,6 +344,7 @@ impl RestrictedSpaceMap {
         RestrictedSpaceMap {
             nr_allocated: 0,
             counts: FixedBitSet::with_capacity(nr_entries as usize),
+            first_free: 0,
         }
     }
 }
@@ -354,11 +377,14 @@ impl SpaceMap for RestrictedSpaceMap {
         } else {
             if old {
                 self.nr_allocated -= 1;
+                if b < self.first_free as u64 {
+                    self.first_free = b as usize;
+                }
             }
             self.counts.set(b as usize, false);
         }
 
-        Ok(if old {1} else {0})
+        Ok(if old { 1 } else { 0 })
     }
 
     fn inc(&mut self, begin: u64, len: u64) -> Result<()> {
@@ -370,6 +396,19 @@ impl SpaceMap for RestrictedSpaceMap {
         }
         Ok(())
     }
+
+    fn alloc(&mut self) -> Result<Option<u64>> {
+        for b in self.first_free..self.counts.len() {
+            if !self.counts.contains(b) {
+                self.counts.insert(b);
+                self.first_free = b + 1;
+                return Ok(Some(b as u64));
+            }
+        }
+
+        self.first_free = self.counts.len();
+        Ok(None)
+    }
 }
 
 //------------------------------------------

src/thin/device_detail.rs

@@ -19,7 +19,7 @@ impl fmt::Display for DeviceDetail {
               self.mapped_blocks,
               self.transaction_id,
               self.creation_time,
-              self.snapshotted_time);
+              self.snapshotted_time)?;
         Ok(())
     }
 }

src/thin/mod.rs

@@ -3,4 +3,5 @@ pub mod device_detail;
 pub mod superblock;
 pub mod check;
 pub mod dump;
+pub mod restore;
 pub mod xml;

src/thin/restore.rs

@@ -0,0 +1,31 @@
+use anyhow::Result;
+use std::collections::{BTreeMap, BTreeSet};
+use std::path::Path;
+use std::sync::{Arc, Mutex};
+
+use crate::io_engine::{AsyncIoEngine, IoEngine, SyncIoEngine};
+use crate::pdata::btree::{self, *};
+use crate::pdata::space_map::*;
+use crate::pdata::unpack::*;
+use crate::report::*;
+use crate::thin::block_time::*;
+use crate::thin::device_detail::*;
+use crate::thin::superblock::*;
+use crate::thin::xml::{self, MetadataVisitor};
+
+//------------------------------------------
+
+pub struct ThinRestoreOptions<'a> {
+    pub input: &'a Path,
+    pub output: &'a Path,
+    pub async_io: bool,
+    pub report: Arc<Report>,
+}
+
+//------------------------------------------
+
+pub fn restore(opts: ThinRestoreOptions) -> Result<()> {
+    todo!();
+}
+
+//------------------------------------------