add take_snapshot to snapshot service

ethcore/src/service.rs

@@ -46,6 +46,8 @@ pub enum ClientIoMessage {
 	FeedStateChunk(H256, Bytes),
 	/// Feed a block chunk to the snapshot service
 	FeedBlockChunk(H256, Bytes),
+	/// Take a snapshot for the block with given number.
+	TakeSnapshot(u64),
 }
 
 /// Client service setup. Creates and registers client and network services with the IO subsystem.
@@ -170,6 +172,11 @@ impl IoHandler<ClientIoMessage> for ClientIoHandler {
 			}
 			ClientIoMessage::FeedStateChunk(ref hash, ref chunk) => self.snapshot.feed_state_chunk(*hash, chunk),
 			ClientIoMessage::FeedBlockChunk(ref hash, ref chunk) => self.snapshot.feed_block_chunk(*hash, chunk),
+			ClientIoMessage::TakeSnapshot(num) => {
+				if let Err(e) = self.snapshot.take_snapshot(&*self.client, num) {
+					warn!("Failed to take snapshot at block #{}: {}", num, e);
+				}
+			}
 			_ => {} // ignore other messages
 		}
 	}

ethcore/src/snapshot/service.rs

@@ -27,8 +27,10 @@ use super::{ManifestData, StateRebuilder, BlockRebuilder};
 use super::io::{SnapshotReader, LooseReader, SnapshotWriter, LooseWriter};
 
 use blockchain::BlockChain;
+use client::Client;
 use engines::Engine;
 use error::Error;
+use ids::BlockID;
 use service::ClientIoMessage;
 use spec::Spec;
 
@@ -269,6 +271,13 @@ impl Service {
 		dir
 	}
 
+	// get the temporary snapshot dir.
+	fn temp_snapshot_dir(&self) -> PathBuf {
+		let mut dir = self.root_dir();
+		dir.push("in_progress");
+		dir
+	}
+
 	// get the restoration directory.
 	fn restoration_dir(&self) -> PathBuf {
 		let mut dir = self.root_dir();
@@ -328,6 +337,34 @@ impl Service {
 		}
 	}
 
+	/// Take a snapshot at the block with the given number.
+	/// calling this while a restoration is in progress or vice versa
+	/// will lead to a race condition where the first one to finish will
+	/// have their produced snapshot overwritten.
+	pub fn take_snapshot(&self, client: &Client, num: u64) -> Result<(), Error> {
+		info!("Taking snapshot at #{}", num);
+
+		let temp_dir = self.temp_snapshot_dir();
+		let snapshot_dir = self.snapshot_dir();
+
+		let _ = fs::remove_dir_all(&temp_dir);
+		let writer = try!(LooseWriter::new(temp_dir.clone()));
+		let progress = Default::default();
+
+		// Todo [rob] log progress.
+		try!(client.take_snapshot(writer, BlockID::Number(num), &progress));
+		let mut reader = self.reader.write();
+
+		// destroy the old snapshot reader.
+		*reader = None;
+
+		try!(fs::rename(temp_dir, &snapshot_dir));
+
+		*reader = Some(try!(LooseReader::new(snapshot_dir)));
+
+		Ok(())
+	}
+
 	/// Initialize the restoration synchronously.
 	pub fn init_restore(&self, manifest: ManifestData) -> Result<(), Error> {
 		let rest_dir = self.restoration_dir();
@@ -393,14 +430,7 @@ impl Service {
 		try!(fs::create_dir(&snapshot_dir));
 
 		trace!(target: "snapshot", "copying restored snapshot files over");
-		for maybe_file in try!(fs::read_dir(self.temp_recovery_dir())) {
-			let path = try!(maybe_file).path();
-			if let Some(name) = path.file_name().map(|x| x.to_owned()) {
-				let mut new_path = snapshot_dir.clone();
-				new_path.push(name);
-				try!(fs::rename(path, new_path));
-			}
-		}
+		try!(fs::rename(self.temp_recovery_dir(), &snapshot_dir));
 
 		let _ = fs::remove_dir_all(self.restoration_dir());