abort snapshot restoration faster (#3356)

* abort snapshot restoration faster

* flag-checking tests

ethcore/src/snapshot/tests/blocks.rs

@@ -17,10 +17,11 @@
 //! Block chunker and rebuilder tests.
 
 use devtools::RandomTempPath;
+use error::Error;
 
 use blockchain::generator::{ChainGenerator, ChainIterator, BlockFinalizer};
 use blockchain::BlockChain;
-use snapshot::{chunk_blocks, BlockRebuilder, Progress};
+use snapshot::{chunk_blocks, BlockRebuilder, Error as SnapshotError, Progress};
 use snapshot::io::{PackedReader, PackedWriter, SnapshotReader, SnapshotWriter};
 
 use util::{Mutex, snappy};
@@ -28,6 +29,7 @@ use util::kvdb::{Database, DatabaseConfig};
 
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
 
 fn chunk_and_restore(amount: u64) {
 	let mut canon_chain = ChainGenerator::default();
@@ -75,10 +77,11 @@ fn chunk_and_restore(amount: u64) {
 	let mut rebuilder = BlockRebuilder::new(new_chain, new_db.clone(), &manifest).unwrap();
 	let reader = PackedReader::new(&snapshot_path).unwrap().unwrap();
 	let engine = ::engines::NullEngine::new(Default::default(), Default::default());
+	let flag = AtomicBool::new(true);
 	for chunk_hash in &reader.manifest().block_hashes {
 		let compressed = reader.chunk(*chunk_hash).unwrap();
 		let chunk = snappy::decompress(&compressed).unwrap();
-		rebuilder.feed(&chunk, &engine).unwrap();
+		rebuilder.feed(&chunk, &engine, &flag).unwrap();
 	}
 
 	rebuilder.finalize(HashMap::new()).unwrap();
@@ -93,3 +96,46 @@ fn chunk_and_restore_500() { chunk_and_restore(500) }
 
 #[test]
 fn chunk_and_restore_40k() { chunk_and_restore(40000) }
+
+#[test]
+fn checks_flag() {
+	use ::rlp::{RlpStream, Stream};
+	use util::H256;
+
+	let mut stream = RlpStream::new_list(5);
+
+	stream.append(&100u64)
+		.append(&H256::default())
+		.append(&(!0u64));
+
+	stream.append_empty_data().append_empty_data();
+
+	let genesis = {
+		let mut canon_chain = ChainGenerator::default();
+		let mut finalizer = BlockFinalizer::default();
+		canon_chain.generate(&mut finalizer).unwrap()
+	};
+
+	let chunk = stream.out();
+	let path = RandomTempPath::create_dir();
+
+	let db_cfg = DatabaseConfig::with_columns(::db::NUM_COLUMNS);
+	let db = Arc::new(Database::open(&db_cfg, path.as_str()).unwrap());
+	let chain = BlockChain::new(Default::default(), &genesis, db.clone());
+	let engine = ::engines::NullEngine::new(Default::default(), Default::default());
+
+	let manifest = ::snapshot::ManifestData {
+		state_hashes: Vec::new(),
+		block_hashes: Vec::new(),
+		state_root: ::util::sha3::SHA3_NULL_RLP,
+		block_number: 102,
+		block_hash: H256::default(),
+	};
+
+	let mut rebuilder = BlockRebuilder::new(chain, db.clone(), &manifest).unwrap();
+
+	match rebuilder.feed(&chunk, &engine, &AtomicBool::new(false)) {
+		Err(Error::Snapshot(SnapshotError::RestorationAborted)) => {}
+		_ => panic!("Wrong result on abort flag set")
+	}
+}

ethcore/src/snapshot/tests/state.rs

@@ -16,10 +16,12 @@
 
 //! State snapshotting tests.
 
-use snapshot::{chunk_state, Progress, StateRebuilder};
+use snapshot::{chunk_state, Error as SnapshotError, Progress, StateRebuilder};
 use snapshot::io::{PackedReader, PackedWriter, SnapshotReader, SnapshotWriter};
 use super::helpers::{compare_dbs, StateProducer};
 
+use error::Error;
+
 use rand::{XorShiftRng, SeedableRng};
 use util::hash::H256;
 use util::journaldb::{self, Algorithm};
@@ -29,6 +31,7 @@ use util::Mutex;
 use devtools::RandomTempPath;
 
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
 
 #[test]
 fn snap_and_restore() {
@@ -65,11 +68,13 @@ fn snap_and_restore() {
 		let mut rebuilder = StateRebuilder::new(new_db.clone(), Algorithm::Archive);
 		let reader = PackedReader::new(&snap_file).unwrap().unwrap();
 
+		let flag = AtomicBool::new(true);
+
 		for chunk_hash in &reader.manifest().state_hashes {
 			let raw = reader.chunk(*chunk_hash).unwrap();
 			let chunk = ::util::snappy::decompress(&raw).unwrap();
 
-			rebuilder.feed(&chunk).unwrap();
+			rebuilder.feed(&chunk, &flag).unwrap();
 		}
 
 		assert_eq!(rebuilder.state_root(), state_root);
@@ -82,3 +87,52 @@ fn snap_and_restore() {
 
 	compare_dbs(&old_db, new_db.as_hashdb());
 }
+
+#[test]
+fn checks_flag() {
+	let mut producer = StateProducer::new();
+	let mut rng = XorShiftRng::from_seed([5, 6, 7, 8]);
+	let mut old_db = MemoryDB::new();
+	let db_cfg = DatabaseConfig::with_columns(::db::NUM_COLUMNS);
+
+	for _ in 0..10 {
+		producer.tick(&mut rng, &mut old_db);
+	}
+
+	let snap_dir = RandomTempPath::create_dir();
+	let mut snap_file = snap_dir.as_path().to_owned();
+	snap_file.push("SNAP");
+
+	let state_root = producer.state_root();
+	let writer = Mutex::new(PackedWriter::new(&snap_file).unwrap());
+
+	let state_hashes = chunk_state(&old_db, &state_root, &writer, &Progress::default()).unwrap();
+
+	writer.into_inner().finish(::snapshot::ManifestData {
+		state_hashes: state_hashes,
+		block_hashes: Vec::new(),
+		state_root: state_root,
+		block_number: 0,
+		block_hash: H256::default(),
+	}).unwrap();
+
+	let mut db_path = snap_dir.as_path().to_owned();
+	db_path.push("db");
+	{
+		let new_db = Arc::new(Database::open(&db_cfg, &db_path.to_string_lossy()).unwrap());
+		let mut rebuilder = StateRebuilder::new(new_db.clone(), Algorithm::Archive);
+		let reader = PackedReader::new(&snap_file).unwrap().unwrap();
+
+		let flag = AtomicBool::new(false);
+
+		for chunk_hash in &reader.manifest().state_hashes {
+			let raw = reader.chunk(*chunk_hash).unwrap();
+			let chunk = ::util::snappy::decompress(&raw).unwrap();
+
+			match rebuilder.feed(&chunk, &flag) {
+				Err(Error::Snapshot(SnapshotError::RestorationAborted)) => {},
+				_ => panic!("unexpected result when feeding with flag off"),
+			}
+		}
+	}
+}