Resumable warp-sync / Seed downloaded snapshots

ethcore/src/snapshot/service.rs

@@ -17,8 +17,8 @@
 //! Snapshot network service implementation.
 
 use std::collections::HashSet;
-use std::io::ErrorKind;
-use std::fs;
+use std::io::{self, Read, ErrorKind};
+use std::fs::{self, File};
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
@@ -30,6 +30,7 @@ use blockchain::BlockChain;
 use client::{Client, ChainInfo, ClientIoMessage};
 use engines::EthEngine;
 use error::Error;
+use hash::keccak;
 use ids::BlockId;
 
 use io::IoChannel;
@@ -240,6 +241,7 @@ pub struct Service {
 	progress: super::Progress,
 	taking_snapshot: AtomicBool,
 	restoring_snapshot: AtomicBool,
+	restoration_ready: AtomicBool,
 }
 
 impl Service {
@@ -261,6 +263,7 @@ impl Service {
 			progress: Default::default(),
 			taking_snapshot: AtomicBool::new(false),
 			restoring_snapshot: AtomicBool::new(false),
+			restoration_ready: AtomicBool::new(false),
 		};
 
 		// create the root snapshot dir if it doesn't exist.
@@ -270,8 +273,8 @@ impl Service {
 			}
 		}
 
-		// delete the temporary restoration dir if it does exist.
-		if let Err(e) = fs::remove_dir_all(service.restoration_dir()) {
+		// delete the temporary restoration DB dir if it does exist.
+		if let Err(e) = fs::remove_dir_all(service.restoration_db()) {
 			if e.kind() != ErrorKind::NotFound {
 				return Err(e.into())
 			}
@@ -325,6 +328,13 @@ impl Service {
 		dir
 	}
 
+	// previous snapshot chunks path.
+	fn prev_chunks_dir(&self) -> PathBuf {
+		let mut dir = self.snapshot_root.clone();
+		dir.push("prev_chunks");
+		dir
+	}
+
 	// replace one the client's database with our own.
 	fn replace_client_db(&self) -> Result<(), Error> {
 		let our_db = self.restoration_db();
@@ -407,54 +417,129 @@ impl Service {
 	/// The recover flag indicates whether to recover the restored snapshot.
 	pub fn init_restore(&self, manifest: ManifestData, recover: bool) -> Result<(), Error> {
 		let rest_dir = self.restoration_dir();
-
-		let mut res = self.restoration.lock();
-
-		self.state_chunks.store(0, Ordering::SeqCst);
-		self.block_chunks.store(0, Ordering::SeqCst);
-
-		// tear down existing restoration.
-		*res = None;
+		let rest_db = self.restoration_db();
+		let recovery_temp = self.temp_recovery_dir();
+		let prev_chunks = self.prev_chunks_dir();
 
 		// delete and restore the restoration dir.
-		if let Err(e) = fs::remove_dir_all(&rest_dir) {
+		if let Err(e) = fs::remove_dir_all(&prev_chunks) {
 			match e.kind() {
 				ErrorKind::NotFound => {},
 				_ => return Err(e.into()),
 			}
 		}
 
-		fs::create_dir_all(&rest_dir)?;
+		{
+			self.restoration_ready.store(false, Ordering::SeqCst);
+			let mut res = self.restoration.lock();
 
-		// make new restoration.
-		let writer = match recover {
-			true => Some(LooseWriter::new(self.temp_recovery_dir())?),
-			false => None
-		};
+			// Move the previous recovery temp directory
+			// to `prev_chunks` to be able to restart restoring
+			// with previously downloaded blocks
+			// This step is optional, so don't fail on error
+			fs::rename(&recovery_temp, &prev_chunks).ok();
 
-		let params = RestorationParams {
-			manifest: manifest,
-			pruning: self.pruning,
-			db: self.restoration_db_handler.open(&self.restoration_db())?,
-			writer: writer,
-			genesis: &self.genesis_block,
-			guard: Guard::new(rest_dir),
-			engine: &*self.engine,
-		};
+			self.state_chunks.store(0, Ordering::SeqCst);
+			self.block_chunks.store(0, Ordering::SeqCst);
 
-		let state_chunks = params.manifest.state_hashes.len();
-		let block_chunks = params.manifest.block_hashes.len();
+			// tear down existing restoration.
+			*res = None;
 
-		*res = Some(Restoration::new(params)?);
+			// delete and restore the restoration dir.
+			if let Err(e) = fs::remove_dir_all(&rest_dir) {
+				match e.kind() {
+					ErrorKind::NotFound => {},
+					_ => return Err(e.into()),
+				}
+			}
 
-		*self.status.lock() = RestorationStatus::Ongoing {
-			state_chunks: state_chunks as u32,
-			block_chunks: block_chunks as u32,
-			state_chunks_done: self.state_chunks.load(Ordering::SeqCst) as u32,
-			block_chunks_done: self.block_chunks.load(Ordering::SeqCst) as u32,
-		};
+			fs::create_dir_all(&rest_dir)?;
+
+			// make new restoration.
+			let writer = match recover {
+				true => Some(LooseWriter::new(self.temp_recovery_dir())?),
+				false => None
+			};
+
+			let params = RestorationParams {
+				manifest: manifest.clone(),
+				pruning: self.pruning,
+				db: self.restoration_db_handler.open(&self.restoration_db())?,
+				writer: writer,
+				genesis: &self.genesis_block,
+				guard: Guard::new(rest_db),
+				engine: &*self.engine,
+			};
+
+			let state_chunks = manifest.state_hashes.len();
+			let block_chunks = manifest.block_hashes.len();
+
+			*res = Some(Restoration::new(params)?);
+
+			*self.status.lock() = RestorationStatus::Ongoing {
+				state_chunks: state_chunks as u32,
+				block_chunks: block_chunks as u32,
+				state_chunks_done: self.state_chunks.load(Ordering::SeqCst) as u32,
+				block_chunks_done: self.block_chunks.load(Ordering::SeqCst) as u32,
+			};
+		}
 
 		self.restoring_snapshot.store(true, Ordering::SeqCst);
+
+		// Import previous chunks, continue if it fails
+		self.import_prev_chunks(manifest).ok();
+		self.restoration_ready.store(true, Ordering::SeqCst);
+
+		Ok(())
+	}
+
+	/// Import the previous chunks into the current restoration
+	fn import_prev_chunks(&self, manifest: ManifestData) -> Result<(), Error> {
+		let prev_chunks = self.prev_chunks_dir();
+
+		// Restore previous snapshot chunks
+		let files = fs::read_dir(prev_chunks.as_path())?;
+		let mut num_temp_chunks = 0;
+
+		for prev_chunk_file in files {
+			// Import the chunk, don't fail and continue if one fails
+			match self.import_prev_chunk(&manifest, prev_chunk_file) {
+				Ok(_) => num_temp_chunks += 1,
+				Err(e) => trace!(target: "snapshot", "Error importing chunk: {:?}", e),
+			}
+		}
+
+		trace!(target:"snapshot", "Imported {} previous chunks", num_temp_chunks);
+
+		// Remove the prev temp directory
+		fs::remove_dir_all(&prev_chunks)?;
+
+		Ok(())
+	}
+
+	/// Import a previous chunk at the given path
+	fn import_prev_chunk(&self, manifest: &ManifestData, file: io::Result<fs::DirEntry>) -> Result<(), Error> {
+		let file = file?;
+		let path = file.path();
+
+		let mut file = File::open(path.clone())?;
+		let mut buffer = Vec::new();
+		file.read_to_end(&mut buffer)?;
+
+		let hash = keccak(&buffer);
+
+		let is_state = if manifest.block_hashes.contains(&hash) {
+			false
+		} else if manifest.state_hashes.contains(&hash) {
+			true
+		} else {
+			return Ok(());
+		};
+
+		self.feed_chunk(hash, &buffer, is_state)?;
+
+		trace!(target: "snapshot", "Fed chunk {:?}", hash);
+
 		Ok(())
 	}
 
@@ -492,6 +577,7 @@ impl Service {
 
 		let _ = fs::remove_dir_all(self.restoration_dir());
 		*self.status.lock() = RestorationStatus::Inactive;
+		self.restoration_ready.store(false, Ordering::SeqCst);
 
 		Ok(())
 	}
@@ -503,7 +589,10 @@ impl Service {
 			let mut restoration = self.restoration.lock();
 
 			match self.status() {
-				RestorationStatus::Inactive | RestorationStatus::Failed => return Ok(()),
+				RestorationStatus::Inactive | RestorationStatus::Failed => {
+					trace!(target: "snapshot", "Tried to restore chunk {:x} while inactive or failed", hash);
+					return Ok(());
+				},
 				RestorationStatus::Ongoing { .. } => {
 					let (res, db) = {
 						let rest = match *restoration {
@@ -583,6 +672,32 @@ impl SnapshotService for Service {
 		self.reader.read().as_ref().and_then(|r| r.chunk(hash).ok())
 	}
 
+	fn restoration_ready(&self) -> bool {
+		self.restoration_ready.load(Ordering::SeqCst)
+	}
+
+	fn completed_chunks(&self) -> Option<Vec<H256>> {
+		let restoration = self.restoration.lock();
+
+		match *restoration {
+			Some(ref restoration) => {
+				let completed_chunks = restoration.manifest.block_hashes
+					.iter()
+					.filter(|h| !restoration.block_chunks_left.contains(h))
+					.chain(
+						restoration.manifest.state_hashes
+							.iter()
+							.filter(|h| !restoration.state_chunks_left.contains(h))
+					)
+					.map(|h| *h)
+					.collect();
+
+				Some(completed_chunks)
+			},
+			None => None,
+		}
+	}
+
 	fn status(&self) -> RestorationStatus {
 		let mut cur_status = self.status.lock();
 		if let RestorationStatus::Ongoing { ref mut state_chunks_done, ref mut block_chunks_done, .. } = *cur_status {
@@ -600,7 +715,9 @@ impl SnapshotService for Service {
 	}
 
 	fn abort_restore(&self) {
+		trace!(target: "snapshot", "Aborting restore");
 		self.restoring_snapshot.store(false, Ordering::SeqCst);
+		self.restoration_ready.store(false, Ordering::SeqCst);
 		*self.restoration.lock() = None;
 		*self.status.lock() = RestorationStatus::Inactive;
 	}

ethcore/src/snapshot/tests/service.rs

@@ -130,12 +130,16 @@ fn guards_delete_folders() {
 	service.init_restore(manifest.clone(), true).unwrap();
 	assert!(path.exists());
 
+	// The `db` folder should have been deleted,
+	// while the `temp` one kept
 	service.abort_restore();
-	assert!(!path.exists());
+	assert!(!path.join("db").exists());
+	assert!(path.join("temp").exists());
 
 	service.init_restore(manifest.clone(), true).unwrap();
 	assert!(path.exists());
 
 	drop(service);
-	assert!(!path.exists());
+	assert!(!path.join("db").exists());
+	assert!(path.join("temp").exists());
 }

ethcore/src/snapshot/traits.rs

@@ -30,6 +30,12 @@ pub trait SnapshotService : Sync + Send {
 	/// `None` indicates warp sync isn't supported by the consensus engine.
 	fn supported_versions(&self) -> Option<(u64, u64)>;
 
+	/// Returns whether the Snapshot Service restoration is ready
+	fn restoration_ready(&self) -> bool;
+
+	/// Returns a list of the completed chunks
+	fn completed_chunks(&self) -> Option<Vec<H256>>;
+
 	/// Get raw chunk for a given hash.
 	fn chunk(&self, hash: H256) -> Option<Bytes>;