Aux data revert implementation

bin/node-sassafras/node/src/command.rs

@@ -96,7 +96,8 @@ pub fn run() -> sc_cli::Result<()> {
 			runner.async_run(|config| {
 				let PartialComponents { client, task_manager, backend, .. } =
 					service::new_partial(&config)?;
-				let aux_revert = Box::new(|client, _, blocks| {
+				let aux_revert = Box::new(|client, backend, blocks| {
+					sc_consensus_sassafras::revert(backend, blocks)?;
 					sc_finality_grandpa::revert(client, blocks)?;
 					Ok(())
 				});

client/consensus/sassafras/src/aux_schema.rs

@@ -17,14 +17,20 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 //! Schema for auxiliary data persistence.
+//!
+//! TODO-SASS-P2 : RENAME FROM aux_schema.rs => aux_data.rs
+
+use std::{collections::HashSet, sync::Arc};
 
 use scale_codec::{Decode, Encode};
 
 use sc_client_api::backend::AuxStore;
 use sc_consensus_epochs::{EpochChangesFor, SharedEpochChanges};
-use sp_blockchain::{Error as ClientError, Result as ClientResult};
+
+use sc_client_api::{blockchain::Backend as _, Backend as BackendT};
+use sp_blockchain::{Error as ClientError, HeaderBackend, HeaderMetadata, Result as ClientResult};
 use sp_consensus_sassafras::SassafrasBlockWeight;
-use sp_runtime::traits::Block as BlockT;
+use sp_runtime::traits::{Block as BlockT, NumberFor, SaturatedConversion, Zero};
 
 use crate::Epoch;
 
@@ -99,3 +105,68 @@ pub fn load_block_weight<H: Encode, B: AuxStore>(
 ) -> ClientResult<Option<SassafrasBlockWeight>> {
 	load_decode(backend, block_weight_key(block_hash).as_slice())
 }
+
+/// Reverts protocol aux data from the best block to at most the last finalized block.
+///
+/// Epoch-changes and block weights announced after the revert point are removed.
+pub fn revert<Block, Backend>(backend: Arc<Backend>, blocks: NumberFor<Block>) -> ClientResult<()>
+where
+	Block: BlockT,
+	Backend: BackendT<Block>,
+{
+	let blockchain = backend.blockchain();
+	let best_number = blockchain.info().best_number;
+	let finalized = blockchain.info().finalized_number;
+
+	let revertible = blocks.min(best_number - finalized);
+	if revertible == Zero::zero() {
+		return Ok(())
+	}
+
+	let revert_up_to_number = best_number - revertible;
+	let revert_up_to_hash = blockchain.hash(revert_up_to_number)?.ok_or(ClientError::Backend(
+		format!("Unexpected hash lookup failure for block number: {}", revert_up_to_number),
+	))?;
+
+	// Revert epoch changes tree.
+
+	// This config is only used on-genesis.
+	let epoch_changes = load_epoch_changes::<Block, _>(&*backend)?;
+	let mut epoch_changes = epoch_changes.shared_data();
+
+	if revert_up_to_number == Zero::zero() {
+		// Special case, no epoch changes data were present on genesis.
+		*epoch_changes = EpochChangesFor::<Block, Epoch>::new();
+	} else {
+		let descendent_query = sc_consensus_epochs::descendent_query(blockchain);
+		epoch_changes.revert(descendent_query, revert_up_to_hash, revert_up_to_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(blockchain, revert_up_to_hash, leaf)
+			.map(|route| route.retracted().is_empty())
+			.unwrap_or_default()
+	});
+
+	for mut hash in leaves {
+		loop {
+			let meta = blockchain.header_metadata(hash)?;
+			if meta.number <= revert_up_to_number || !weight_keys.insert(block_weight_key(hash)) {
+				// We've reached the revert point or an already processed branch, stop here.
+				break
+			}
+			hash = meta.parent;
+		}
+	}
+
+	let weight_keys: Vec<_> = weight_keys.iter().map(|val| val.as_slice()).collect();
+
+	// Write epoch changes and remove weights in one shot.
+	write_epoch_changes::<Block, _, _>(&epoch_changes, |values| {
+		AuxStore::insert_aux(&*backend, values, weight_keys.iter())
+	})
+}

client/consensus/sassafras/src/lib.rs

@@ -94,7 +94,9 @@ mod block_import;
 mod tests;
 mod verification;
 
+// Export core components.
 pub use authorship::{start_sassafras, SassafrasParams, SassafrasWorker};
+pub use aux_schema::revert;
 pub use block_import::{block_import, SassafrasBlockImport};
 pub use verification::SassafrasVerifier;
 

client/consensus/sassafras/src/tests.rs

@@ -22,11 +22,10 @@
 use super::*;
 
 use futures::executor::block_on;
-use std::{cell::RefCell, sync::Arc};
+use std::sync::Arc;
 
-use sc_block_builder::{BlockBuilder, BlockBuilderProvider};
-use sc_client_api::TransactionFor;
-use sc_consensus::{BlockImport, BoxBlockImport, BoxJustificationImport};
+use sc_block_builder::BlockBuilderProvider;
+use sc_consensus::{BlockImport, BoxJustificationImport};
 use sc_keystore::LocalKeystore;
 use sc_network_test::*;
 use sp_application_crypto::key_types::SASSAFRAS;
@@ -188,13 +187,8 @@ impl TestEnvironment {
 
 	// Propose and import a new Sassafras block on top of the given parent.
 	// This skips verification.
-	fn propose_and_import_block(
-		//<Transaction: Send + 'static>(
-		&mut self,
-		parent_id: BlockId,
-		slot: Option<Slot>,
-	) -> Hash {
-		let mut parent = self.client.header(&parent_id).unwrap().unwrap();
+	fn propose_and_import_block(&mut self, parent_id: BlockId, slot: Option<Slot>) -> Hash {
+		let parent = self.client.header(&parent_id).unwrap().unwrap();
 
 		let proposer = block_on(self.init(&parent)).unwrap();
 
@@ -269,12 +263,7 @@ impl TestEnvironment {
 	// Propose and import n valid Sassafras blocks that are built on top of the given parent.
 	// The proposer takes care of producing epoch change digests according to the epoch
 	// duration (which is set to 6 slots in the test runtime).
-	fn propose_and_import_blocks(
-		//)<Transaction: Send + 'static>(
-		&mut self,
-		mut parent_id: BlockId,
-		n: usize,
-	) -> Vec<Hash> {
+	fn propose_and_import_blocks(&mut self, mut parent_id: BlockId, n: usize) -> Vec<Hash> {
 		let mut hashes = Vec::with_capacity(n);
 
 		for _ in 0..n {
@@ -382,7 +371,7 @@ fn revert_prunes_epoch_changes_and_removes_weights() {
 	assert_eq!(epoch_changes.shared_data().tree().roots().count(), 1);
 
 	// Revert canon chain to block #10 (best(21) - 11)
-	// crate::revert(client.clone(), backend, 11).unwrap();
+	crate::aux_schema::revert(backend, 11).unwrap();
 
 	// Load and check epoch changes.