Optimize finalization performance (paritytech#4922)

This PR largely fixes
paritytech#4903 by addressing it
from a few different directions.

The high-level observation is that complexity of finalization was
unfortunately roughly `O(n^3)`. Not only
`displaced_leaves_after_finalizing` was extremely inefficient on its
own, especially when large ranges of blocks were involved, it was called
once upfront and then on every single block that was finalized over and
over again.

The first commit refactores code adjacent to
`displaced_leaves_after_finalizing` to optimize memory allocations. For
example things like `BTreeMap<_, Vec<_>>` were very bad in terms of
number of allocations and after analyzing code paths was completely
unnecessary and replaced with `Vec<(_, _)>`. In other places allocations
of known size were not done upfront and some APIs required unnecessary
cloning of vectors.

I checked invariants and didn't find anything that was violated after
refactoring.

Second commit completely replaces `displaced_leaves_after_finalizing`
implementation with a much more efficient one. In my case with ~82k
blocks and ~13k leaves it takes ~5.4s to finish
`client.apply_finality()` now.

The idea is to avoid querying the same blocks over and over again as
well as introducing temporary local cache for blocks related to leaves
above block that is being finalized as well as local cache of the
finalized branch of the chain. I left some comments in the code and
wrote tests that I belive should check all code invariants for
correctness. `lowest_common_ancestor_multiblock` was removed as
unnecessary and not great in terms of performance API, domain-specific
code should be written instead like done in
`displaced_leaves_after_finalizing`.

After these changes I noticed finalization is still horribly slow,
turned out that even though `displaced_leaves_after_finalizing` was way
faster that before (probably order of magnitude), it was called for
every single of those 82k blocks 🤦

The quick hack I came up with in the third commit to handle this edge
case was to not call it when finalizing multiple blocks at once until
the very last moment. It works and allows to finish the whole
finalization in just 14 seconds (5.4+5.4 of which are two calls to
`displaced_leaves_after_finalizing`). I'm really not happy with the fact
that `displaced_leaves_after_finalizing` is called twice, but much
heavier refactoring would be necessary to get rid of second call.

---

Next steps:
* assuming the changes are acceptable I'll write prdoc
* paritytech#4920 or something
similar in spirit should be implemented to unleash efficient parallelsm
with rayon in `displaced_leaves_after_finalizing`, which will allow to
further (and significant!) scale its performance rather that being
CPU-bound on a single core, also reading database sequentially should
ideally be avoided
* someone should look into removal of the second
`displaced_leaves_after_finalizing` call
* further cleanups are possible if `undo_finalization` can be removed

---

Polkadot Address: 1vSxzbyz2cJREAuVWjhXUT1ds8vBzoxn2w4asNpusQKwjJd

---------

Co-authored-by: Sebastian Kunert <[email protected]>

prdoc/pr_4922.prdoc

@@ -0,0 +1,15 @@
+title: Optimize finalization performance
+
+doc:
+  - audience: Node Dev
+    description: |
+      Finalization algorithm was replaced with a more efficient version, data structures refactored to be faster and do
+      fewer memory allocations. As the result some APIs have changed in a minor, but incompatible way.
+
+crates:
+- name: sc-client-api
+  bump: major
+- name: sc-client-db
+  bump: major
+- name: sp-blockchain
+  bump: major

substrate/client/api/src/backend.rs

@@ -217,7 +217,8 @@ pub trait BlockImportOperation<Block: BlockT> {
 	where
 		I: IntoIterator<Item = (Vec<u8>, Option<Vec<u8>>)>;
 
-	/// Mark a block as finalized.
+	/// Mark a block as finalized, if multiple blocks are finalized in the same operation then they
+	/// must be marked in ascending order.
 	fn mark_finalized(
 		&mut self,
 		hash: Block::Hash,

substrate/client/api/src/leaves.rs

@@ -45,33 +45,20 @@ pub struct RemoveOutcome<H, N> {
 }
 
 /// Removed leaves after a finalization action.
-pub struct FinalizationOutcome<H, N> {
-	removed: BTreeMap<Reverse<N>, Vec<H>>,
+pub struct FinalizationOutcome<I, H, N>
+where
+	I: Iterator<Item = (N, H)>,
+{
+	removed: I,
 }
 
-impl<H: Copy, N: Ord> FinalizationOutcome<H, N> {
-	/// Merge with another. This should only be used for displaced items that
-	/// are produced within one transaction of each other.
-	pub fn merge(&mut self, mut other: Self) {
-		// this will ignore keys that are in duplicate, however
-		// if these are actually produced correctly via the leaf-set within
-		// one transaction, then there will be no overlap in the keys.
-		self.removed.append(&mut other.removed);
-	}
-
-	/// Iterate over all displaced leaves.
-	pub fn leaves(&self) -> impl Iterator<Item = &H> {
-		self.removed.values().flatten()
-	}
-
+impl<I, H: Ord, N: Ord> FinalizationOutcome<I, H, N>
+where
+	I: Iterator<Item = (N, H)>,
+{
 	/// Constructor
-	pub fn new(new_displaced: impl Iterator<Item = (H, N)>) -> Self {
-		let mut removed = BTreeMap::<Reverse<N>, Vec<H>>::new();
-		for (hash, number) in new_displaced {
-			removed.entry(Reverse(number)).or_default().push(hash);
-		}
-
-		FinalizationOutcome { removed }
+	pub fn new(new_displaced: I) -> Self {
+		FinalizationOutcome { removed: new_displaced }
 	}
 }
 
@@ -86,7 +73,7 @@ pub struct LeafSet<H, N> {
 impl<H, N> LeafSet<H, N>
 where
 	H: Clone + PartialEq + Decode + Encode,
-	N: std::fmt::Debug + Clone + AtLeast32Bit + Decode + Encode,
+	N: std::fmt::Debug + Copy + AtLeast32Bit + Decode + Encode,
 {
 	/// Construct a new, blank leaf set.
 	pub fn new() -> Self {
@@ -117,13 +104,13 @@ where
 		let number = Reverse(number);
 
 		let removed = if number.0 != N::zero() {
-			let parent_number = Reverse(number.0.clone() - N::one());
+			let parent_number = Reverse(number.0 - N::one());
 			self.remove_leaf(&parent_number, &parent_hash).then(|| parent_hash)
 		} else {
 			None
 		};
 
-		self.insert_leaf(number.clone(), hash.clone());
+		self.insert_leaf(number, hash.clone());
 
 		ImportOutcome { inserted: LeafSetItem { hash, number }, removed }
 	}
@@ -150,7 +137,7 @@ where
 
 		let inserted = parent_hash.and_then(|parent_hash| {
 			if number.0 != N::zero() {
-				let parent_number = Reverse(number.0.clone() - N::one());
+				let parent_number = Reverse(number.0 - N::one());
 				self.insert_leaf(parent_number, parent_hash.clone());
 				Some(parent_hash)
 			} else {
@@ -162,11 +149,12 @@ where
 	}
 
 	/// Remove all leaves displaced by the last block finalization.
-	pub fn remove_displaced_leaves(&mut self, displaced_leaves: &FinalizationOutcome<H, N>) {
-		for (number, hashes) in &displaced_leaves.removed {
-			for hash in hashes.iter() {
-				self.remove_leaf(number, hash);
-			}
+	pub fn remove_displaced_leaves<I>(&mut self, displaced_leaves: FinalizationOutcome<I, H, N>)
+	where
+		I: Iterator<Item = (N, H)>,
+	{
+		for (number, hash) in displaced_leaves.removed {
+			self.remove_leaf(&Reverse(number), &hash);
 		}
 	}
 
@@ -186,13 +174,13 @@ where
 		let items = self
 			.storage
 			.iter()
-			.flat_map(|(number, hashes)| hashes.iter().map(move |h| (h.clone(), number.clone())))
+			.flat_map(|(number, hashes)| hashes.iter().map(move |h| (h.clone(), *number)))
 			.collect::<Vec<_>>();
 
-		for (hash, number) in &items {
+		for (hash, number) in items {
 			if number.0 > best_number {
 				assert!(
-					self.remove_leaf(number, hash),
+					self.remove_leaf(&number, &hash),
 					"item comes from an iterator over storage; qed",
 				);
 			}
@@ -207,7 +195,7 @@ where
 		// we need to make sure that the best block exists in the leaf set as
 		// this is an invariant of regular block import.
 		if !leaves_contains_best {
-			self.insert_leaf(best_number.clone(), best_hash.clone());
+			self.insert_leaf(best_number, best_hash.clone());
 		}
 	}
 
@@ -229,7 +217,7 @@ where
 		column: u32,
 		prefix: &[u8],
 	) {
-		let leaves: Vec<_> = self.storage.iter().map(|(n, h)| (n.0.clone(), h.clone())).collect();
+		let leaves: Vec<_> = self.storage.iter().map(|(n, h)| (n.0, h.clone())).collect();
 		tx.set_from_vec(column, prefix, leaves.encode());
 	}
 
@@ -274,7 +262,7 @@ where
 
 	/// Returns the highest leaf and all hashes associated to it.
 	pub fn highest_leaf(&self) -> Option<(N, &[H])> {
-		self.storage.iter().next().map(|(k, v)| (k.0.clone(), &v[..]))
+		self.storage.iter().next().map(|(k, v)| (k.0, &v[..]))
 	}
 }
 
@@ -286,13 +274,13 @@ pub struct Undo<'a, H: 'a, N: 'a> {
 impl<'a, H: 'a, N: 'a> Undo<'a, H, N>
 where
 	H: Clone + PartialEq + Decode + Encode,
-	N: std::fmt::Debug + Clone + AtLeast32Bit + Decode + Encode,
+	N: std::fmt::Debug + Copy + AtLeast32Bit + Decode + Encode,
 {
 	/// Undo an imported block by providing the import operation outcome.
 	/// No additional operations should be performed between import and undo.
 	pub fn undo_import(&mut self, outcome: ImportOutcome<H, N>) {
 		if let Some(removed_hash) = outcome.removed {
-			let removed_number = Reverse(outcome.inserted.number.0.clone() - N::one());
+			let removed_number = Reverse(outcome.inserted.number.0 - N::one());
 			self.inner.insert_leaf(removed_number, removed_hash);
 		}
 		self.inner.remove_leaf(&outcome.inserted.number, &outcome.inserted.hash);
@@ -302,16 +290,21 @@ where
 	/// No additional operations should be performed between remove and undo.
 	pub fn undo_remove(&mut self, outcome: RemoveOutcome<H, N>) {
 		if let Some(inserted_hash) = outcome.inserted {
-			let inserted_number = Reverse(outcome.removed.number.0.clone() - N::one());
+			let inserted_number = Reverse(outcome.removed.number.0 - N::one());
 			self.inner.remove_leaf(&inserted_number, &inserted_hash);
 		}
 		self.inner.insert_leaf(outcome.removed.number, outcome.removed.hash);
 	}
 
 	/// Undo a finalization operation by providing the displaced leaves.
 	/// No additional operations should be performed between finalization and undo.
-	pub fn undo_finalization(&mut self, mut outcome: FinalizationOutcome<H, N>) {
-		self.inner.storage.append(&mut outcome.removed);
+	pub fn undo_finalization<I>(&mut self, outcome: FinalizationOutcome<I, H, N>)
+	where
+		I: Iterator<Item = (N, H)>,
+	{
+		for (number, hash) in outcome.removed {
+			self.inner.storage.entry(Reverse(number)).or_default().push(hash);
+		}
 	}
 }