From b181951bcc15db9246a80c599dd27a5c6ae01f8a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Silva?= <andre.beat@gmail.com>
Date: Sat, 5 Oct 2019 09:44:07 +0100
Subject: [PATCH] babe: prune the epoch tree on startup (#3768)

---
 substrate/core/consensus/babe/src/lib.rs | 69 +++++++++++++++++-------
 1 file changed, 49 insertions(+), 20 deletions(-)

diff --git a/substrate/core/consensus/babe/src/lib.rs b/substrate/core/consensus/babe/src/lib.rs
index 000c14269f..4f69c176ce 100644
--- a/substrate/core/consensus/babe/src/lib.rs
+++ b/substrate/core/consensus/babe/src/lib.rs
@@ -112,7 +112,7 @@ mod tests;
 pub use babe_primitives::{
 	AuthorityId, AuthorityPair, AuthoritySignature, Epoch, NextEpochDescriptor,
 };
-pub use epoch_changes::{EpochChanges, SharedEpochChanges};
+pub use epoch_changes::{EpochChanges, EpochChangesFor, SharedEpochChanges};
 
 macro_rules! babe_err {
 	($($i: expr),+) => {
@@ -889,24 +889,10 @@ impl<B, E, Block, I, RA, PRA> BlockImport<Block> for BabeBlockImport<B, E, Block
 			// used by pruning may not know about the block that is being
 			// imported.
 			let prune_and_import = || {
-				let finalized_slot = {
-					let finalized_header = self.client.header(&BlockId::Hash(info.finalized_hash))
-						.map_err(|e| ConsensusError::ClientImport(format!("{:?}", e)))?
-						.expect("best finalized hash was given by client; \
-								finalized headers must exist in db; qed");
-
-					find_pre_digest::<Block::Header>(&finalized_header)
-						.expect("finalized header must be valid; \
-								valid blocks have a pre-digest; qed")
-						.slot_number()
-				};
-
-				epoch_changes.prune_finalized(
-					descendent_query(&*self.client),
-					&info.finalized_hash,
-					info.finalized_number,
-					finalized_slot,
-				).map_err(|e| ConsensusError::ClientImport(format!("{:?}", e)))?;
+				prune_finalized(
+					&self.client,
+					&mut epoch_changes,
+				)?;
 
 				epoch_changes.import(
 					descendent_query(&*self.client),
@@ -989,6 +975,40 @@ impl<B, E, Block, I, RA, PRA> BlockImport<Block> for BabeBlockImport<B, E, Block
 	}
 }
 
+/// Gets the best finalized block and its slot, and prunes the given epoch tree.
+fn prune_finalized<B, E, Block, RA>(
+	client: &Client<B, E, Block, RA>,
+	epoch_changes: &mut EpochChangesFor<Block>,
+) -> Result<(), ConsensusError> where
+	Block: BlockT<Hash=H256>,
+	E: CallExecutor<Block, Blake2Hasher> + Send + Sync,
+	B: Backend<Block, Blake2Hasher>,
+	RA: Send + Sync,
+{
+	let info = client.info().chain;
+
+	let finalized_slot = {
+		let finalized_header = client.header(&BlockId::Hash(info.finalized_hash))
+			.map_err(|e| ConsensusError::ClientImport(format!("{:?}", e)))?
+			.expect("best finalized hash was given by client; \
+				 finalized headers must exist in db; qed");
+
+		find_pre_digest::<Block::Header>(&finalized_header)
+			.expect("finalized header must be valid; \
+					 valid blocks have a pre-digest; qed")
+			.slot_number()
+	};
+
+	epoch_changes.prune_finalized(
+		descendent_query(&*client),
+		&info.finalized_hash,
+		info.finalized_number,
+		finalized_slot,
+	).map_err(|e| ConsensusError::ClientImport(format!("{:?}", e)))?;
+
+	Ok(())
+}
+
 /// Produce a BABE block-import object to be used later on in the construction of
 /// an import-queue.
 ///
@@ -1001,7 +1021,8 @@ pub fn block_import<B, E, Block: BlockT<Hash=H256>, I, RA, PRA>(
 	api: Arc<PRA>,
 ) -> ClientResult<(BabeBlockImport<B, E, Block, I, RA, PRA>, BabeLink<Block>)> where
 	B: Backend<Block, Blake2Hasher>,
-	E: CallExecutor<Block, Blake2Hasher>,
+	E: CallExecutor<Block, Blake2Hasher> + Send + Sync,
+	RA: Send + Sync,
 {
 	let epoch_changes = aux_schema::load_epoch_changes(&*client)?;
 	let link = BabeLink {
@@ -1010,6 +1031,14 @@ pub fn block_import<B, E, Block: BlockT<Hash=H256>, I, RA, PRA>(
 		config: config.clone(),
 	};
 
+	// NOTE: this isn't entirely necessary, but since we didn't use to prune the
+	// epoch tree it is useful as a migration, so that nodes prune long trees on
+	// startup rather than waiting until importing the next epoch change block.
+	prune_finalized(
+		&client,
+		&mut epoch_changes.lock(),
+	)?;
+
 	let import = BabeBlockImport::new(
 		client,
 		api,