Keep last finalized overlay in memory until DB commit (#409)

substrate/substrate/state-db/src/lib.rs

@@ -194,6 +194,8 @@ impl<BlockHash: Hash, Key: Hash> StateDbSync<BlockHash, Key> {
 	}
 
 	pub fn finalize_block(&mut self, hash: &BlockHash) -> CommitSet<Key> {
+		// clear the temporary overlay from the previous finalization.
+		self.unfinalized.clear_overlay();
 		let mut commit = match self.mode {
 			PruningMode::ArchiveAll => {
 				CommitSet::default()

substrate/substrate/state-db/src/unfinalized.rs

@@ -16,6 +16,9 @@
 
 //! Finalization window.
 //! Maintains trees of block overlays and allows discarding trees/roots
+//! The overlays are added in `insert` and removed in `finalize`.
+//! Last finalized overlay is kept in memory until next call to `finalize` or
+//! `clear_overlay`
 
 use std::collections::{HashMap, VecDeque};
 use super::{Error, DBValue, ChangeSet, CommitSet, MetaDb, Hash, to_meta_key};
@@ -29,6 +32,7 @@ pub struct UnfinalizedOverlay<BlockHash: Hash, Key: Hash> {
 	last_finalized: Option<(BlockHash, u64)>,
 	levels: VecDeque<Vec<BlockOverlay<BlockHash, Key>>>,
 	parents: HashMap<BlockHash, BlockHash>,
+	last_finalized_overlay: HashMap<Key, DBValue>,
 }
 
 struct JournalRecord<BlockHash: Hash, Key: Hash> {
@@ -121,6 +125,7 @@ impl<BlockHash: Hash, Key: Hash> UnfinalizedOverlay<BlockHash, Key> {
 			last_finalized: last_finalized,
 			levels,
 			parents,
+			last_finalized_overlay: Default::default(),
 		})
 	}
 
@@ -199,6 +204,11 @@ impl<BlockHash: Hash, Key: Hash> UnfinalizedOverlay<BlockHash, Key> {
 		self.last_finalized.as_ref().map(|&(_, n)| n + 1).unwrap_or(0)
 	}
 
+	/// This may be called when the last finalization commit was applied to the database.
+	pub fn clear_overlay(&mut self) {
+		self.last_finalized_overlay.clear();
+	}
+
 	/// Select a top-level root and finalized it. Discards all sibling subtrees and the root.
 	/// Returns a set of changes that need to be added to the DB.
 	pub fn finalize(&mut self, hash: &BlockHash) -> CommitSet<Key> {
@@ -212,8 +222,9 @@ impl<BlockHash: Hash, Key: Hash> UnfinalizedOverlay<BlockHash, Key> {
 		for (i, overlay) in level.into_iter().enumerate() {
 			self.parents.remove(&overlay.hash);
 			if i == index {
+				self.last_finalized_overlay = overlay.values;
 				// that's the one we need to finalize
-				commit.data.inserted = overlay.values.into_iter().collect();
+				commit.data.inserted = self.last_finalized_overlay.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
 				commit.data.deleted = overlay.deleted;
 			} else {
 				// TODO: borrow checker won't allow us to split out mutable refernces
@@ -236,6 +247,9 @@ impl<BlockHash: Hash, Key: Hash> UnfinalizedOverlay<BlockHash, Key> {
 
 	/// Get a value from the node overlay. This searches in every existing changeset.
 	pub fn get(&self, key: &Key) -> Option<DBValue> {
+		if let Some(value) = self.last_finalized_overlay.get(&key) {
+			return Some(value.clone());
+		}
 		for level in self.levels.iter() {
 			for overlay in level.iter() {
 				if let Some(value) = overlay.values.get(&key) {
@@ -374,9 +388,12 @@ mod tests {
 		db.commit(&overlay.finalize(&h1));
 		assert_eq!(overlay.levels.len(), 1);
 		assert_eq!(overlay.parents.len(), 1);
+		assert!(contains(&overlay, 5));
+		overlay.clear_overlay();
 		assert!(!contains(&overlay, 5));
 		assert!(contains(&overlay, 7));
 		db.commit(&overlay.finalize(&h2));
+		overlay.clear_overlay();
 		assert_eq!(overlay.levels.len(), 0);
 		assert_eq!(overlay.parents.len(), 0);
 		assert!(db.data_eq(&make_db(&[1, 4, 6, 7, 8])));
@@ -446,6 +463,7 @@ mod tests {
 
 		// finalize 1. 2 and all its children should be discarded
 		db.commit(&overlay.finalize(&h_1));
+		overlay.clear_overlay();
 		assert_eq!(overlay.levels.len(), 2);
 		assert_eq!(overlay.parents.len(), 6);
 		assert!(!contains(&overlay, 1));
@@ -457,6 +475,7 @@ mod tests {
 
 		// finalize 1_2. 1_1 and all its children should be discarded
 		db.commit(&overlay.finalize(&h_1_2));
+		overlay.clear_overlay();
 		assert_eq!(overlay.levels.len(), 1);
 		assert_eq!(overlay.parents.len(), 3);
 		assert!(!contains(&overlay, 11));
@@ -467,6 +486,7 @@ mod tests {
 
 		// finalize 1_2_2
 		db.commit(&overlay.finalize(&h_1_2_2));
+		overlay.clear_overlay();
 		assert_eq!(overlay.levels.len(), 0);
 		assert_eq!(overlay.parents.len(), 0);
 		assert!(db.data_eq(&make_db(&[1, 12, 122])));