Fixed state pinning in state-db (#4557)

* Account for references when pinning

* Fixed pinned state issues

* Fixes

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

@@ -37,7 +37,9 @@ 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
+	//would be deleted but kept around because block is pinned, ref counted.
+	pinned: HashMap<BlockHash, u32>,
+	pinned_insertions: HashMap<BlockHash, Vec<Key>>,
 }
 
 #[derive(Encode, Decode)]
@@ -68,21 +70,14 @@ 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>,
-	mut into: Option<&mut HashMap<Key, DBValue>>,
-) {
+fn discard_values<Key: Hash>(values: &mut HashMap<Key, (u32, DBValue)>, inserted: Vec<Key>) {
 	for k in inserted {
 		match values.entry(k) {
 			Entry::Occupied(mut e) => {
 				let (ref mut counter, _) = e.get_mut();
 				*counter -= 1;
 				if *counter == 0 {
-					let (key, (_, value)) = e.remove_entry();
-					if let Some(ref mut into) = into {
-						into.insert(key, value);
-					}
+					e.remove_entry();
 				}
 			},
 			Entry::Vacant(_) => {
@@ -97,7 +92,8 @@ 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>>,
+	pinned: &HashMap<BlockHash, u32>,
+	pinned_insertions: &mut HashMap<BlockHash, Vec<Key>>,
 	hash: &BlockHash,
 ) {
 	let mut discarded = Vec::new();
@@ -105,9 +101,15 @@ fn discard_descendants<BlockHash: Hash, Key: Hash>(
 		*level = level.drain(..).filter_map(|overlay| {
 			let parent = parents.get(&overlay.hash).expect("there is a parent entry for each entry in levels; qed").clone();
 			if parent == *hash {
-				parents.remove(&overlay.hash);
-				discarded.push(overlay.hash);
-				discard_values(&mut values, overlay.inserted, pinned.get_mut(hash));
+				discarded.push(overlay.hash.clone());
+				if pinned.contains_key(&overlay.hash) {
+					// save to be discarded later.
+					pinned_insertions.insert(overlay.hash.clone(), overlay.inserted);
+				} else {
+					// discard immediatelly.
+					parents.remove(&overlay.hash);
+					discard_values(&mut values, overlay.inserted);
+				}
 				None
 			} else {
 				Some(overlay)
@@ -115,7 +117,7 @@ fn discard_descendants<BlockHash: Hash, Key: Hash>(
 		}).collect();
 	}
 	for hash in discarded {
-		discard_descendants(levels, values, index + 1, parents, pinned, &hash);
+		discard_descendants(levels, values, index + 1, parents, pinned, pinned_insertions, &hash);
 	}
 }
 
@@ -176,6 +178,7 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 			pending_canonicalizations: Default::default(),
 			pending_insertions: Default::default(),
 			pinned: Default::default(),
+			pinned_insertions: Default::default(),
 			values: values,
 		})
 	}
@@ -339,18 +342,23 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 
 			// discard unfinalized overlays and values
 			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,
-						&mut self.pinned,
+						&self.pinned,
+						&mut self.pinned_insertions,
 						&overlay.hash,
 					);
 				}
-				discard_values(&mut self.values, overlay.inserted, self.pinned.get_mut(&overlay.hash));
+				if self.pinned.contains_key(&overlay.hash) {
+					self.pinned_insertions.insert(overlay.hash.clone(), overlay.inserted);
+				} else {
+					self.parents.remove(&overlay.hash);
+					discard_values(&mut self.values, overlay.inserted);
+				}
 			}
 		}
 		if let Some(hash) = last {
@@ -364,11 +372,6 @@ 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
 	}
 
@@ -385,7 +388,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, None);
+				discard_values(&mut self.values, overlay.inserted);
 			}
 			commit
 		})
@@ -402,7 +405,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, None);
+			discard_values(&mut self.values, overlay.inserted);
 			if self.levels[level_index].is_empty() {
 				debug_assert_eq!(level_index, self.levels.len() - 1);
 				self.levels.pop_back();
@@ -424,12 +427,43 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 
 	/// Pin state values in memory
 	pub fn pin(&mut self, hash: &BlockHash) {
-		self.pinned.insert(hash.clone(), HashMap::default());
+		if self.pending_insertions.contains(hash) {
+			debug_assert!(false, "Trying to pin pending state");
+			return;
+		}
+		// Also pin all parents
+		let mut parent = Some(hash);
+		while let Some(hash) = parent {
+			let refs = self.pinned.entry(hash.clone()).or_default();
+			if *refs == 0 {
+				trace!(target: "state-db", "Pinned non-canon block: {:?}", hash);
+			}
+			*refs += 1;
+			parent = self.parents.get(hash);
+		}
 	}
 
 	/// Discard pinned state
 	pub fn unpin(&mut self, hash: &BlockHash) {
-		self.pinned.remove(hash);
+		// Also unpin all parents
+		let mut parent = Some(hash.clone());
+		while let Some(hash) = parent {
+			parent = self.parents.get(&hash).cloned();
+			match self.pinned.entry(hash.clone()) {
+				Entry::Occupied(mut entry) => {
+					*entry.get_mut() -= 1;
+					if *entry.get() == 0 {
+						entry.remove();
+						if let Some(inserted) = self.pinned_insertions.remove(&hash) {
+							trace!(target: "state-db", "Discarding unpinned non-canon block: {:?}", hash);
+							discard_values(&mut self.values, inserted);
+							self.parents.remove(&hash);
+						}
+					}
+				},
+				Entry::Vacant(_) => {},
+			}
+		}
 	}
 }
 
@@ -801,7 +835,7 @@ mod tests {
 	fn keeps_pinned() {
 		let mut db = make_db(&[]);
 
-		// - 1 - 1_1
+		// - 0 - 1_1
 		//     \ 1_2
 
 		let (h_1, c_1) = (H256::random(), make_changeset(&[1], &[]));
@@ -810,6 +844,7 @@ mod tests {
 		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.apply_pending();
 
 		overlay.pin(&h_1);
 
@@ -821,4 +856,65 @@ mod tests {
 		overlay.unpin(&h_1);
 		assert!(!contains(&overlay, 1));
 	}
+
+	#[test]
+	fn keeps_pinned_ref_count() {
+		let mut db = make_db(&[]);
+
+		// - 0 - 1_1
+		//     \ 1_2
+		//     \ 1_3
+
+		// 1_1 and 1_2 both make the same change
+		let (h_1, c_1) = (H256::random(), make_changeset(&[1], &[]));
+		let (h_2, c_2) = (H256::random(), make_changeset(&[1], &[]));
+		let (h_3, c_3) = (H256::random(), make_changeset(&[], &[]));
+
+		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());
+		db.commit(&overlay.insert::<io::Error>(&h_3, 1, &H256::default(), c_3).unwrap());
+		overlay.apply_pending();
+
+		overlay.pin(&h_1);
+
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_3, &mut commit).unwrap();
+		db.commit(&commit);
+		overlay.apply_pending(); // 1_2 should be discarded, 1_1 is pinned
+
+		assert!(contains(&overlay, 1));
+		overlay.unpin(&h_1);
+		assert!(!contains(&overlay, 1));
+	}
+
+	#[test]
+	fn pin_keeps_parent() {
+		let mut db = make_db(&[]);
+
+		// - 0 - 1_1 - 2_1
+		//     \ 1_2
+
+		let (h_11, c_11) = (H256::random(), make_changeset(&[1], &[]));
+		let (h_12, c_12) = (H256::random(), make_changeset(&[], &[]));
+		let (h_21, c_21) = (H256::random(), make_changeset(&[], &[]));
+
+		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
+		db.commit(&overlay.insert::<io::Error>(&h_11, 1, &H256::default(), c_11).unwrap());
+		db.commit(&overlay.insert::<io::Error>(&h_12, 1, &H256::default(), c_12).unwrap());
+		db.commit(&overlay.insert::<io::Error>(&h_21, 2, &h_11, c_21).unwrap());
+		overlay.apply_pending();
+
+		overlay.pin(&h_21);
+
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_12, &mut commit).unwrap();
+		db.commit(&commit);
+		overlay.apply_pending(); // 1_1 and 2_1 should be both pinned
+
+		assert!(contains(&overlay, 1));
+		overlay.unpin(&h_21);
+		assert!(!contains(&overlay, 1));
+		assert!(overlay.pinned.is_empty());
+	}
 }