Remove babe's blocks weight on revert

bin/node/cli/src/command.rs

@@ -157,10 +157,9 @@ pub fn run() -> Result<()> {
 			let runner = cli.create_runner(cmd)?;
 			runner.async_run(|config| {
 				let PartialComponents { client, task_manager, backend, .. } = new_partial(&config)?;
-				let client_clone = client.clone();
-				let revert_aux = Box::new(move |blocks| {
-					sc_consensus_babe::revert(client_clone, blocks)?;
-					// TODO: grandpa
+				let revert_aux = Box::new(|client, backend, blocks| {
+					sc_consensus_babe::revert(client, backend, blocks)?;
+					// TODO: grandpa revert
 					Ok(())
 				});
 

client/cli/src/commands/revert_cmd.rs

@@ -44,15 +44,16 @@ pub struct RevertCmd {
 }
 
 /// Revert handler for auxiliary data (e.g. consensus).
-type AuxRevertHandler<B> = Box<dyn FnOnce(NumberFor<B>) -> error::Result<()>>;
+type AuxRevertHandler<C, BA, B> =
+	Box<dyn FnOnce(Arc<C>, Arc<BA>, NumberFor<B>) -> error::Result<()>>;
 
 impl RevertCmd {
 	/// Run the revert command
 	pub async fn run<B, BA, C>(
 		&self,
 		client: Arc<C>,
 		backend: Arc<BA>,
-		aux_revert: Option<AuxRevertHandler<B>>,
+		aux_revert: Option<AuxRevertHandler<C, BA, B>>,
 	) -> error::Result<()>
 	where
 		B: BlockT,
@@ -62,7 +63,7 @@ impl RevertCmd {
 	{
 		let blocks = self.num.parse()?;
 		if let Some(aux_revert) = aux_revert {
-			aux_revert(blocks)?;
+			aux_revert(client.clone(), backend.clone(), blocks)?;
 		}
 		revert_chain(client, backend, blocks)?;
 

client/consensus/babe/src/lib.rs

@@ -92,8 +92,8 @@ use retain_mut::RetainMut;
 use schnorrkel::SignatureError;
 
 use sc_client_api::{
-	backend::AuxStore, AuxDataOperations, BlockchainEvents, FinalityNotification, PreCommitActions,
-	ProvideUncles, UsageProvider,
+	backend::AuxStore, AuxDataOperations, Backend as BackendT, BlockchainEvents,
+	FinalityNotification, PreCommitActions, ProvideUncles, UsageProvider,
 };
 use sc_consensus::{
 	block_import::{
@@ -113,7 +113,9 @@ use sc_telemetry::{telemetry, TelemetryHandle, CONSENSUS_DEBUG, CONSENSUS_TRACE}
 use sp_api::{ApiExt, ProvideRuntimeApi};
 use sp_application_crypto::AppKey;
 use sp_block_builder::BlockBuilder as BlockBuilderApi;
-use sp_blockchain::{Error as ClientError, HeaderBackend, HeaderMetadata, Result as ClientResult};
+use sp_blockchain::{
+	Backend as _, Error as ClientError, HeaderBackend, HeaderMetadata, Result as ClientResult,
+};
 use sp_consensus::{
 	BlockOrigin, CacheKeyId, CanAuthorWith, Environment, Error as ConsensusError, Proposer,
 	SelectChain,
@@ -1832,7 +1834,11 @@ where
 }
 
 /// Reverts aux data.
-pub fn revert<Block, Client>(client: Arc<Client>, blocks: NumberFor<Block>) -> ClientResult<()>
+pub fn revert<Block, Client, Backend>(
+	client: Arc<Client>,
+	backend: Arc<Backend>,
+	blocks: NumberFor<Block>,
+) -> ClientResult<()>
 where
 	Block: BlockT,
 	Client: AuxStore
@@ -1841,6 +1847,7 @@ where
 		+ ProvideRuntimeApi<Block>
 		+ UsageProvider<Block>,
 	Client::Api: BabeApi<Block>,
+	Backend: BackendT<Block>,
 {
 	let best_number = client.info().best_number;
 	let finalized = client.info().finalized_number;
@@ -1849,7 +1856,12 @@ where
 	let number = best_number.saturating_sub(revertible);
 	let hash = client
 		.block_hash_from_id(&BlockId::Number(number))?
-		.ok_or(ClientError::Backend(format!("Hash lookup failure for block number: {}", number)))?;
+		.ok_or(ClientError::Backend(format!(
+			"Unexpected hash lookup failure for block number: {}",
+			number
+		)))?;
+
+	// Revert epoch changes tree.
 
 	let config = Config::get(&*client)?;
 	let epoch_changes =
@@ -1863,7 +1875,31 @@ where
 		epoch_changes.revert(descendent_query(&*client), hash, number);
 	}
 
+	// Remove block weights added after the revert point.
+
+	let mut weight_keys = HashSet::with_capacity(revertible.saturated_into());
+	let leaves = backend.blockchain().leaves()?.into_iter().filter(|&leaf| {
+		sp_blockchain::tree_route(&*client, hash, leaf)
+			.map(|route| route.retracted().is_empty())
+			.unwrap_or_default()
+	});
+	for leaf in leaves {
+		let mut hash = leaf;
+		// Insert parent after parent until we don't hit an already processed
+		// branch or we reach a direct child of the rollback point.
+		while weight_keys.insert(aux_schema::block_weight_key(hash)) {
+			let meta = client.header_metadata(hash)?;
+			if meta.number == number + One::one() {
+				// We've reached a child of the revert point, stop here.
+				break
+			}
+			hash = client.header_metadata(hash)?.parent;
+		}
+	}
+	let weight_keys: Vec<_> = weight_keys.iter().map(|val| val.as_slice()).collect();
+
+	// Write epoch changes and remove weights in one shot.
 	aux_schema::write_epoch_changes::<Block, _, _>(&epoch_changes, |values| {
-		client.insert_aux(values, &[])
+		client.insert_aux(values, weight_keys.iter())
 	})
 }

client/consensus/babe/src/tests.rs

@@ -736,13 +736,14 @@ fn importing_block_one_sets_genesis_epoch() {
 }
 
 #[test]
-fn revert_prunes_epoch_changes_tree() {
+fn revert_prunes_epoch_changes_and_removes_weights() {
 	let mut net = BabeTestNet::new(1);
 
 	let peer = net.peer(0);
 	let data = peer.data.as_ref().expect("babe link set up during initialization");
 
 	let client = peer.client().as_client();
+	let backend = peer.client().as_backend();
 	let mut block_import = data.block_import.lock().take().expect("import set up during init");
 	let epoch_changes = data.link.epoch_changes.clone();
 
@@ -762,26 +763,25 @@ fn revert_prunes_epoch_changes_tree() {
 	// One branch starts before the revert point (epoch data should be maintained).
 	// One branch starts after the revert point (epoch data should be removed).
 	//
-	//                        *----------------- F(#13)                < fork #2
+	//                        *----------------- F(#13) --#18                  < fork #2
 	//                       /
-	// A(#1) ---- B(#7) ----#8----+-----#12----- C(#13) ---- D(#19)    < canon
+	// A(#1) ---- B(#7) ----#8----+-----#12----- C(#13) ---- D(#19) ------#21  < canon
 	//   \                        ^       \
-	//    \                    revert      *---- G(#13) -----H(#19)    < fork #3
+	//    \                    revert      *---- G(#13) ---- H(#19) ---#20     < fork #3
 	//     \                   to #10
-	//      *-----E(#7)                                                < fork #1
-
+	//      *-----E(#7)---#11                                          < fork #1
 	let canon = propose_and_import_blocks_wrap(BlockId::Number(0), 21);
 	let fork1 = propose_and_import_blocks_wrap(BlockId::Hash(canon[0]), 10);
 	let fork2 = propose_and_import_blocks_wrap(BlockId::Hash(canon[7]), 10);
-	let _fork3 = propose_and_import_blocks_wrap(BlockId::Hash(canon[11]), 10);
+	let fork3 = propose_and_import_blocks_wrap(BlockId::Hash(canon[11]), 8);
 
 	// We should be tracking a total of 9 epochs in the fork tree
 	assert_eq!(epoch_changes.shared_data().tree().iter().count(), 8);
 	// And only one root
 	assert_eq!(epoch_changes.shared_data().tree().roots().count(), 1);
 
-	// Revert to block #10 (best(22) - 12)
-	revert(client.clone(), 12).expect("revert should work for the baked test scenario");
+	// Revert canon chain to block #10 (best(21) - 11)
+	revert(client.clone(), backend, 11).expect("revert should work for baked test scenario");
 
 	// Load and check epoch changes.
 
@@ -804,6 +804,17 @@ fn revert_prunes_epoch_changes_tree() {
 	];
 
 	assert_eq!(actual_nodes, expected_nodes);
+
+	let weight_data_check = |hashes: &[Hash], expected: bool| {
+		hashes.iter().all(|hash| {
+			aux_schema::load_block_weight(&*client, hash).unwrap().is_some() == expected
+		})
+	};
+	assert!(weight_data_check(&canon[..10], true));
+	assert!(weight_data_check(&canon[10..], false));
+	assert!(weight_data_check(&fork1, true));
+	assert!(weight_data_check(&fork2, true));
+	assert!(weight_data_check(&fork3, false));
 }
 
 #[test]