Snapshot fixes and optimizations (#2863)

ethcore/src/snapshot/mod.rs

@@ -38,6 +38,7 @@ use util::kvdb::Database;
 use util::trie::{TrieDB, TrieDBMut, Trie, TrieMut};
 use util::sha3::SHA3_NULL_RLP;
 use rlp::{RlpStream, Stream, UntrustedRlp, View};
+use bloom_journal::Bloom;
 
 use self::account::Account;
 use self::block::AbridgedBlock;
@@ -390,6 +391,7 @@ pub struct StateRebuilder {
 	state_root: H256,
 	code_map: HashMap<H256, Bytes>, // maps code hashes to code itself.
 	missing_code: HashMap<H256, Vec<H256>>, // maps code hashes to lists of accounts missing that code.
+	bloom: Bloom,
 }
 
 impl StateRebuilder {
@@ -400,6 +402,7 @@ impl StateRebuilder {
 			state_root: SHA3_NULL_RLP,
 			code_map: HashMap::new(),
 			missing_code: HashMap::new(),
+			bloom: StateDB::load_bloom(&*db),
 		}
 	}
 
@@ -462,9 +465,6 @@ impl StateRebuilder {
 
 		let backing = self.db.backing().clone();
 
-		// bloom has to be updated
-		let mut bloom = StateDB::load_bloom(&backing);
-
 		// batch trie writes
 		{
 			let mut account_trie = if self.state_root != SHA3_NULL_RLP {
@@ -475,17 +475,17 @@ impl StateRebuilder {
 
 			for (hash, thin_rlp) in pairs {
 				if &thin_rlp[..] != &empty_rlp[..] {
-					bloom.set(&*hash);
+					self.bloom.set(&*hash);
 				}
 				try!(account_trie.insert(&hash, &thin_rlp));
 			}
 		}
 
-		let bloom_journal = bloom.drain_journal();
+		let bloom_journal = self.bloom.drain_journal();
 		let mut batch = backing.transaction();
 		try!(StateDB::commit_bloom(&mut batch, bloom_journal));
 		try!(self.db.inject(&mut batch));
-		try!(backing.write(batch).map_err(::util::UtilError::SimpleString));
+		backing.write_buffered(batch);
 		trace!(target: "snapshot", "current state root: {:?}", self.state_root);
 		Ok(())
 	}
@@ -628,7 +628,7 @@ impl BlockRebuilder {
 			} else {
 				self.chain.insert_unordered_block(&mut batch, &block_bytes, receipts, None, is_best, false);
 			}
-			self.db.write(batch).expect("Error writing to the DB");
+			self.db.write_buffered(batch);
 			self.chain.commit();
 
 			parent_hash = BlockView::new(&block_bytes).hash();

ethcore/src/snapshot/service.rs

@@ -74,6 +74,7 @@ struct Restoration {
 	snappy_buffer: Bytes,
 	final_state_root: H256,
 	guard: Guard,
+	db: Arc<Database>,
 }
 
 struct RestorationParams<'a> {
@@ -105,12 +106,13 @@ impl Restoration {
 			manifest: manifest,
 			state_chunks_left: state_chunks,
 			block_chunks_left: block_chunks,
-			state: StateRebuilder::new(raw_db, params.pruning),
+			state: StateRebuilder::new(raw_db.clone(), params.pruning),
 			blocks: blocks,
 			writer: params.writer,
 			snappy_buffer: Vec::new(),
 			final_state_root: root,
 			guard: params.guard,
+			db: raw_db,
 		})
 	}
 
@@ -467,39 +469,46 @@ impl Service {
 	/// Feed a chunk of either kind. no-op if no restoration or status is wrong.
 	fn feed_chunk(&self, hash: H256, chunk: &[u8], is_state: bool) -> Result<(), Error> {
 		// TODO: be able to process block chunks and state chunks at same time?
-		let mut restoration = self.restoration.lock();
+		let (result, db) = {
+			let mut restoration = self.restoration.lock();
 
-		match self.status() {
-			RestorationStatus::Inactive | RestorationStatus::Failed => Ok(()),
-			RestorationStatus::Ongoing { .. } => {
-				let res = {
-					let rest = match *restoration {
-						Some(ref mut r) => r,
-						None => return Ok(()),
-					};
-
-					match is_state {
-						true => rest.feed_state(hash, chunk),
-						false => rest.feed_blocks(hash, chunk, &*self.engine),
-					}.map(|_| rest.is_done())
-				};
-
-				match res {
-					Ok(is_done) => {
-						match is_state {
-							true => self.state_chunks.fetch_add(1, Ordering::SeqCst),
-							false => self.block_chunks.fetch_add(1, Ordering::SeqCst),
+			match self.status() {
+				RestorationStatus::Inactive | RestorationStatus::Failed => return Ok(()),
+				RestorationStatus::Ongoing { .. } => {
+					let (res, db) = {
+						let rest = match *restoration {
+							Some(ref mut r) => r,
+							None => return Ok(()),
 						};
 
-						match is_done {
-							true => self.finalize_restoration(&mut *restoration),
-							false => Ok(())
+						(match is_state {
+							true => rest.feed_state(hash, chunk),
+							false => rest.feed_blocks(hash, chunk, &*self.engine),
+						}.map(|_| rest.is_done()), rest.db.clone())
+					};
+
+					let res = match res {
+						Ok(is_done) => {
+							match is_state {
+								true => self.state_chunks.fetch_add(1, Ordering::SeqCst),
+								false => self.block_chunks.fetch_add(1, Ordering::SeqCst),
+							};
+
+							match is_done {
+								true => {
+									try!(db.flush().map_err(::util::UtilError::SimpleString));
+									self.finalize_restoration(&mut *restoration)
+								},
+								false => Ok(())
+							}
 						}
-					}
-					other => other.map(drop),
+						other => other.map(drop),
+					};
+					(res, db)
 				}
 			}
-		}
+		};
+		result.and_then(|_| db.flush().map_err(|e| ::util::UtilError::SimpleString(e).into()))
 	}
 
 	/// Feed a state chunk to be processed synchronously.
@@ -549,8 +558,9 @@ impl SnapshotService for Service {
 	}
 
 	fn begin_restore(&self, manifest: ManifestData) {
-		self.io_channel.send(ClientIoMessage::BeginRestoration(manifest))
-			.expect("snapshot service and io service are kept alive by client service; qed");
+		if let Err(e) = self.io_channel.send(ClientIoMessage::BeginRestoration(manifest)) {
+			trace!("Error sending snapshot service message: {:?}", e);
+		}
 	}
 
 	fn abort_restore(&self) {
@@ -559,13 +569,15 @@ impl SnapshotService for Service {
 	}
 
 	fn restore_state_chunk(&self, hash: H256, chunk: Bytes) {
-		self.io_channel.send(ClientIoMessage::FeedStateChunk(hash, chunk))
-			.expect("snapshot service and io service are kept alive by client service; qed");
+		if let Err(e) = self.io_channel.send(ClientIoMessage::FeedStateChunk(hash, chunk)) {
+			trace!("Error sending snapshot service message: {:?}", e);
+		}
 	}
 
 	fn restore_block_chunk(&self, hash: H256, chunk: Bytes) {
-		self.io_channel.send(ClientIoMessage::FeedBlockChunk(hash, chunk))
-			.expect("snapshot service and io service are kept alive by client service; qed");
+		if let Err(e) = self.io_channel.send(ClientIoMessage::FeedBlockChunk(hash, chunk)) {
+			trace!("Error sending snapshot service message: {:?}", e);
+		}
 	}
 }