More robust state pinning (#3355)

* Better state pinning

* Fixed pinning race

* Update core/state-db/src/noncanonical.rs

Co-Authored-By: Robert Habermeier <rphmeier@gmail.com>

substrate/core/state-db/src/noncanonical.rs

@@ -37,6 +37,7 @@ pub struct NonCanonicalOverlay<BlockHash: Hash, Key: Hash> {
 	pending_canonicalizations: Vec<BlockHash>,
 	pending_insertions: Vec<BlockHash>,
 	values: HashMap<Key, (u32, DBValue)>, //ref counted
+	pinned: HashMap<BlockHash, HashMap<Key, DBValue>>, //would be deleted but kept around because block is pinned
 }
 
 #[derive(Encode, Decode)]
@@ -67,14 +68,21 @@ fn insert_values<Key: Hash>(values: &mut HashMap<Key, (u32, DBValue)>, inserted:
 	}
 }
 
-fn discard_values<Key: Hash>(values: &mut HashMap<Key, (u32, DBValue)>, inserted: Vec<Key>) {
+fn discard_values<Key: Hash>(
+	values: &mut HashMap<Key, (u32, DBValue)>,
+	inserted: Vec<Key>,
+	mut into: Option<&mut HashMap<Key, DBValue>>,
+) {
 	for k in inserted {
 		match values.entry(k) {
 			Entry::Occupied(mut e) => {
 				let (ref mut counter, _) = e.get_mut();
 				*counter -= 1;
 				if *counter == 0 {
-					e.remove();
+					let (key, (_, value)) = e.remove_entry();
+					if let Some(ref mut into) = into {
+						into.insert(key, value);
+					}
 				}
 			},
 			Entry::Vacant(_) => {
@@ -89,8 +97,9 @@ fn discard_descendants<BlockHash: Hash, Key: Hash>(
 	mut values: &mut HashMap<Key, (u32, DBValue)>,
 	index: usize,
 	parents: &mut HashMap<BlockHash, BlockHash>,
+	pinned: &mut HashMap<BlockHash, HashMap<Key, DBValue>>,
 	hash: &BlockHash,
-	) {
+) {
 	let mut discarded = Vec::new();
 	if let Some(level) = levels.get_mut(index) {
 		*level = level.drain(..).filter_map(|overlay| {
@@ -98,7 +107,7 @@ fn discard_descendants<BlockHash: Hash, Key: Hash>(
 			if parent == *hash {
 				parents.remove(&overlay.hash);
 				discarded.push(overlay.hash);
-				discard_values(&mut values, overlay.inserted);
+				discard_values(&mut values, overlay.inserted, pinned.get_mut(hash));
 				None
 			} else {
 				Some(overlay)
@@ -106,7 +115,7 @@ fn discard_descendants<BlockHash: Hash, Key: Hash>(
 		}).collect();
 	}
 	for hash in discarded {
-		discard_descendants(levels, values, index + 1, parents, &hash);
+		discard_descendants(levels, values, index + 1, parents, pinned, &hash);
 	}
 }
 
@@ -166,6 +175,7 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 			parents,
 			pending_canonicalizations: Default::default(),
 			pending_insertions: Default::default(),
+			pinned: Default::default(),
 			values: values,
 		})
 	}
@@ -278,7 +288,6 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 	pub fn canonicalize<E: fmt::Debug>(
 		&mut self,
 		hash: &BlockHash,
-		pinned: &HashMap<BlockHash, u32>,
 		commit: &mut CommitSet<Key>,
 	) -> Result<(), Error<E>> {
 		trace!(target: "state-db", "Canonicalizing {:?}", hash);
@@ -303,13 +312,6 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 			discarded_blocks.push(overlay.hash.clone());
 		}
 
-		for hash in discarded_blocks.into_iter() {
-			if pinned.contains_key(&hash) {
-				trace!(target: "state-db", "Refusing to discard pinned state {:?}", hash);
-				return Err(Error::DiscardingPinned)
-			}
-		}
-
 		// get the one we need to canonicalize
 		let overlay = &level[index];
 		commit.data.inserted.extend(overlay.inserted.iter()
@@ -339,9 +341,16 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 			for (i, overlay) in level.into_iter().enumerate() {
 				self.parents.remove(&overlay.hash);
 				if i != index {
-					discard_descendants(&mut self.levels, &mut self.values, 0, &mut self.parents, &overlay.hash);
+					discard_descendants(
+						&mut self.levels,
+						&mut self.values,
+						0,
+						&mut self.parents,
+						&mut self.pinned,
+						&overlay.hash,
+					);
 				}
-				discard_values(&mut self.values, overlay.inserted);
+				discard_values(&mut self.values, overlay.inserted, self.pinned.get_mut(&overlay.hash));
 			}
 		}
 		if let Some(hash) = last {
@@ -355,6 +364,11 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 		if let Some((_, value)) = self.values.get(&key) {
 			return Some(value.clone());
 		}
+		for pinned in self.pinned.values() {
+			if let Some(value) = pinned.get(&key) {
+				return Some(value.clone());
+			}
+		}
 		None
 	}
 
@@ -371,7 +385,7 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 			for overlay in level.into_iter() {
 				commit.meta.deleted.push(overlay.journal_key);
 				self.parents.remove(&overlay.hash);
-				discard_values(&mut self.values, overlay.inserted);
+				discard_values(&mut self.values, overlay.inserted, None);
 			}
 			commit
 		})
@@ -388,7 +402,7 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 				.expect("Hash is added in insert");
 
 			let	overlay = self.levels[level_index].pop().expect("Empty levels are not allowed in self.levels");
-			discard_values(&mut self.values, overlay.inserted);
+			discard_values(&mut self.values, overlay.inserted, None);
 			if self.levels[level_index].is_empty() {
 				debug_assert_eq!(level_index, self.levels.len() - 1);
 				self.levels.pop_back();
@@ -407,11 +421,21 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 		self.pending_canonicalizations.clear();
 		self.revert_insertions();
 	}
+
+	/// Pin state values in memory
+	pub fn pin(&mut self, hash: &BlockHash) {
+		self.pinned.insert(hash.clone(), HashMap::default());
+	}
+
+	/// Discard pinned state
+	pub fn unpin(&mut self, hash: &BlockHash) {
+		self.pinned.remove(hash);
+	}
 }
 
 #[cfg(test)]
 mod tests {
-	use std::{collections::HashMap, io};
+	use std::io;
 	use primitives::H256;
 	use super::{NonCanonicalOverlay, to_journal_key};
 	use crate::{ChangeSet, CommitSet};
@@ -436,7 +460,7 @@ mod tests {
 		let db = make_db(&[]);
 		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&H256::default(), &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&H256::default(), &mut commit).unwrap();
 	}
 
 	#[test]
@@ -481,7 +505,7 @@ mod tests {
 		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
 		overlay.insert::<io::Error>(&h1, 1, &H256::default(), ChangeSet::default()).unwrap();
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h2, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h2, &mut commit).unwrap();
 	}
 
 	#[test]
@@ -497,7 +521,7 @@ mod tests {
 		assert_eq!(insertion.meta.deleted.len(), 0);
 		db.commit(&insertion);
 		let mut finalization = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut finalization).unwrap();
+		overlay.canonicalize::<io::Error>(&h1, &mut finalization).unwrap();
 		assert_eq!(finalization.data.inserted.len(), changeset.inserted.len());
 		assert_eq!(finalization.data.deleted.len(), changeset.deleted.len());
 		assert_eq!(finalization.meta.inserted.len(), 1);
@@ -531,7 +555,7 @@ mod tests {
 		db.commit(&overlay.insert::<io::Error>(&h1, 10, &H256::default(), make_changeset(&[3, 4], &[2])).unwrap());
 		db.commit(&overlay.insert::<io::Error>(&h2, 11, &h1, make_changeset(&[5], &[3])).unwrap());
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h1, &mut commit).unwrap();
 		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 1);
@@ -558,7 +582,7 @@ mod tests {
 		assert_eq!(overlay.levels.len(), 2);
 		assert_eq!(overlay.parents.len(), 2);
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h1, &mut commit).unwrap();
 		db.commit(&commit);
 		assert!(contains(&overlay, 5));
 		assert_eq!(overlay.levels.len(), 2);
@@ -569,7 +593,7 @@ mod tests {
 		assert!(!contains(&overlay, 5));
 		assert!(contains(&overlay, 7));
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h2, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h2, &mut commit).unwrap();
 		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 0);
@@ -588,7 +612,7 @@ mod tests {
 		db.commit(&overlay.insert::<io::Error>(&h_2, 1, &H256::default(), c_2).unwrap());
 		assert!(contains(&overlay, 1));
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h_1, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h_1, &mut commit).unwrap();
 		db.commit(&commit);
 		assert!(contains(&overlay, 1));
 		overlay.apply_pending();
@@ -607,8 +631,8 @@ mod tests {
 		db.commit(&overlay.insert::<io::Error>(&h2, 2, &h1, changeset.clone()).unwrap());
 		overlay.apply_pending();
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut commit).unwrap();
-		overlay.canonicalize::<io::Error>(&h2, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h1, &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h2, &mut commit).unwrap();
 		db.commit(&commit);
 		db.commit(&overlay.insert::<io::Error>(&h3, 3, &h2, changeset.clone()).unwrap());
 		overlay.apply_pending();
@@ -679,7 +703,7 @@ mod tests {
 
 		// canonicalize 1. 2 and all its children should be discarded
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h_1, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h_1, &mut commit).unwrap();
 		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 2);
@@ -698,14 +722,9 @@ mod tests {
 		assert!(db.get_meta(&to_journal_key(2, 2)).unwrap().is_none());
 		assert!(db.get_meta(&to_journal_key(2, 3)).unwrap().is_none());
 
-		// check that discarding pinned state produces an error.
-		let mut commit = CommitSet::default();
-		let pinned = vec![(h_1_1_1, 1)].into_iter().collect();
-		assert!(overlay.canonicalize::<io::Error>(&h_1_2, &pinned, &mut commit).is_err());
-
 		// canonicalize 1_2. 1_1 and all its children should be discarded
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h_1_2, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h_1_2, &mut commit).unwrap();
 		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 1);
@@ -722,7 +741,7 @@ mod tests {
 
 		// canonicalize 1_2_2
 		let mut commit = CommitSet::default();
-		overlay.canonicalize::<io::Error>(&h_1_2_2, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h_1_2_2, &mut commit).unwrap();
 		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 0);
@@ -777,5 +796,29 @@ mod tests {
 		assert_eq!(overlay.levels.len(), 0);
 		assert_eq!(overlay.parents.len(), 0);
 	}
-}
 
+	#[test]
+	fn keeps_pinned() {
+		let mut db = make_db(&[]);
+
+		// - 1 - 1_1
+		//     \ 1_2
+
+		let (h_1, c_1) = (H256::random(), make_changeset(&[1], &[]));
+		let (h_2, c_2) = (H256::random(), make_changeset(&[2], &[]));
+
+		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
+		db.commit(&overlay.insert::<io::Error>(&h_1, 1, &H256::default(), c_1).unwrap());
+		db.commit(&overlay.insert::<io::Error>(&h_2, 1, &H256::default(), c_2).unwrap());
+
+		overlay.pin(&h_1);
+
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_2, &mut commit).unwrap();
+		db.commit(&commit);
+		overlay.apply_pending();
+		assert!(contains(&overlay, 1));
+		overlay.unpin(&h_1);
+		assert!(!contains(&overlay, 1));
+	}
+}