add guard for temporary directories

This commit is contained in:
Robert Habermeier 2016-09-02 19:00:20 +02:00
parent 1c450f616d
commit d9eb87cae7

View File

@ -41,6 +41,23 @@ use util::journaldb::Algorithm;
use util::kvdb::{Database, DatabaseConfig}; use util::kvdb::{Database, DatabaseConfig};
use util::snappy; use util::snappy;
/// Helper for removing directories in case of error.
struct Guard(bool, PathBuf);
impl Guard {
fn new(path: PathBuf) -> Self { Guard(true, path) }
fn disarm(mut self) { self.0 = false }
}
impl Drop for Guard {
fn drop(&mut self) {
if self.0 {
let _ = fs::remove_dir_all(&self.1);
}
}
}
/// Statuses for restorations. /// Statuses for restorations.
#[derive(PartialEq, Clone, Copy, Debug)] #[derive(PartialEq, Clone, Copy, Debug)]
pub enum RestorationStatus { pub enum RestorationStatus {
@ -98,6 +115,7 @@ struct Restoration {
writer: LooseWriter, writer: LooseWriter,
snappy_buffer: Bytes, snappy_buffer: Bytes,
final_state_root: H256, final_state_root: H256,
guard: Guard,
} }
struct RestorationParams<'a> { struct RestorationParams<'a> {
@ -106,6 +124,7 @@ struct RestorationParams<'a> {
db_path: PathBuf, // database path db_path: PathBuf, // database path
writer: LooseWriter, // writer for recovered snapshot. writer: LooseWriter, // writer for recovered snapshot.
genesis: &'a [u8], // genesis block of the chain. genesis: &'a [u8], // genesis block of the chain.
guard: Guard, // guard for the restoration directory.
} }
impl Restoration { impl Restoration {
@ -133,6 +152,7 @@ impl Restoration {
writer: params.writer, writer: params.writer,
snappy_buffer: Vec::new(), snappy_buffer: Vec::new(),
final_state_root: root, final_state_root: root,
guard: params.guard,
}) })
} }
@ -181,6 +201,7 @@ impl Restoration {
try!(self.writer.finish(self.manifest)); try!(self.writer.finish(self.manifest));
self.guard.disarm();
Ok(()) Ok(())
} }
@ -348,10 +369,12 @@ impl Service {
let snapshot_dir = self.snapshot_dir(); let snapshot_dir = self.snapshot_dir();
let _ = fs::remove_dir_all(&temp_dir); let _ = fs::remove_dir_all(&temp_dir);
let writer = try!(LooseWriter::new(temp_dir.clone())); let writer = try!(LooseWriter::new(temp_dir.clone()));
let progress = Default::default(); let progress = Default::default();
// Todo [rob] log progress. // Todo [rob] log progress.
let guard = Guard::new(temp_dir.clone());
try!(client.take_snapshot(writer, BlockID::Number(num), &progress)); try!(client.take_snapshot(writer, BlockID::Number(num), &progress));
let mut reader = self.reader.write(); let mut reader = self.reader.write();
@ -362,6 +385,7 @@ impl Service {
*reader = Some(try!(LooseReader::new(snapshot_dir))); *reader = Some(try!(LooseReader::new(snapshot_dir)));
guard.disarm();
Ok(()) Ok(())
} }
@ -393,6 +417,7 @@ impl Service {
db_path: self.restoration_db(), db_path: self.restoration_db(),
writer: writer, writer: writer,
genesis: &self.genesis_block, genesis: &self.genesis_block,
guard: Guard::new(rest_dir),
}; };
*res = Some(try!(Restoration::new(params))); *res = Some(try!(Restoration::new(params)));