Chain-selection subsystem data revert (#5350)

* Chain-selection subsystem data revert

* Cargo fmt

* Better code comments

* Remove unwraps

* Document public method

* Remove duplicated 'ChainSelectionSubsystem' impl block

* Fix typos

* Nitpicks

* Revert returns a service Error

* Removed superflous error handling

* Apply suggestions from code review

* Rename tree 'revert' to 'revert_to'

Co-authored-by: Sebastian Kunert <skunert49@gmail.com>

polkadot/node/core/chain-selection/src/lib.rs

@@ -315,6 +315,17 @@ impl ChainSelectionSubsystem {
 	pub fn new(config: Config, db: Arc<dyn Database>) -> Self {
 		ChainSelectionSubsystem { config, db }
 	}
+
+	/// Revert to the block corresponding to the specified `hash`.
+	/// The revert is not allowed for blocks older than the last finalized one.
+	pub fn revert(&self, hash: Hash) -> Result<(), Error> {
+		let backend_config = db_backend::v1::Config { col_data: self.config.col_data };
+		let mut backend = db_backend::v1::DbBackend::new(self.db.clone(), backend_config);
+
+		let ops = tree::revert_to(&backend, hash)?.into_write_ops();
+
+		backend.write(ops)
+	}
 }
 
 impl<Context> overseer::Subsystem<Context, SubsystemError> for ChainSelectionSubsystem
@@ -323,9 +334,9 @@ where
 	Context: overseer::SubsystemContext<Message = ChainSelectionMessage>,
 {
 	fn start(self, ctx: Context) -> SpawnedSubsystem {
-		let backend = crate::db_backend::v1::DbBackend::new(
+		let backend = db_backend::v1::DbBackend::new(
 			self.db,
-			crate::db_backend::v1::Config { col_data: self.config.col_data },
+			db_backend::v1::Config { col_data: self.config.col_data },
 		);
 
 		SpawnedSubsystem {
@@ -412,7 +423,7 @@ where
 							let _ = tx.send(leaves);
 						}
 						ChainSelectionMessage::BestLeafContaining(required, tx) => {
-							let best_containing = crate::backend::find_best_leaf_containing(
+							let best_containing = backend::find_best_leaf_containing(
 								&*backend,
 								required,
 							)?;
@@ -549,7 +560,7 @@ async fn handle_active_leaf(
 		};
 
 		let reversion_logs = extract_reversion_logs(&header);
-		crate::tree::import_block(
+		tree::import_block(
 			&mut overlay,
 			hash,
 			header.number,
@@ -612,8 +623,7 @@ fn handle_finalized_block(
 	finalized_hash: Hash,
 	finalized_number: BlockNumber,
 ) -> Result<(), Error> {
-	let ops =
-		crate::tree::finalize_block(&*backend, finalized_hash, finalized_number)?.into_write_ops();
+	let ops = tree::finalize_block(&*backend, finalized_hash, finalized_number)?.into_write_ops();
 
 	backend.write(ops)
 }
@@ -623,7 +633,7 @@ fn handle_approved_block(backend: &mut impl Backend, approved_block: Hash) -> Re
 	let ops = {
 		let mut overlay = OverlayedBackend::new(&*backend);
 
-		crate::tree::approve_block(&mut overlay, approved_block)?;
+		tree::approve_block(&mut overlay, approved_block)?;
 
 		overlay.into_write_ops()
 	};
@@ -633,7 +643,7 @@ fn handle_approved_block(backend: &mut impl Backend, approved_block: Hash) -> Re
 
 fn detect_stagnant(backend: &mut impl Backend, now: Timestamp) -> Result<(), Error> {
 	let ops = {
-		let overlay = crate::tree::detect_stagnant(&*backend, now)?;
+		let overlay = tree::detect_stagnant(&*backend, now)?;
 
 		overlay.into_write_ops()
 	};

polkadot/node/core/chain-selection/src/tree.rs

@@ -24,6 +24,7 @@
 //! and as the finalized block advances, orphaned sub-trees are entirely pruned.
 
 use polkadot_node_primitives::BlockWeight;
+use polkadot_node_subsystem::ChainApiError;
 use polkadot_primitives::v2::{BlockNumber, Hash};
 
 use std::collections::HashMap;
@@ -86,7 +87,7 @@ impl ViabilityUpdate {
 
 // Propagate viability update to descendants of the given block. This writes
 // the `base` entry as well as all descendants. If the parent of the block
-// entry is not viable, this wlil not affect any descendants.
+// entry is not viable, this will not affect any descendants.
 //
 // If the block entry provided is self-unviable, then it's assumed that an
 // unviability update needs to be propagated to descendants.
@@ -561,3 +562,102 @@ pub(super) fn detect_stagnant<'a, B: 'a + Backend>(
 
 	Ok(backend)
 }
+
+/// Revert the tree to the block relative to `hash`.
+///
+/// This accepts a fresh backend and returns an overlay on top of it representing
+/// all changes made.
+pub(super) fn revert_to<'a, B: Backend + 'a>(
+	backend: &'a B,
+	hash: Hash,
+) -> Result<OverlayedBackend<'a, B>, Error> {
+	let first_number = backend.load_first_block_number()?.unwrap_or_default();
+
+	let mut backend = OverlayedBackend::new(backend);
+
+	let mut entry = match backend.load_block_entry(&hash)? {
+		Some(entry) => entry,
+		None => {
+			// May be a revert to the last finalized block. If this is the case,
+			// then revert to this block should be handled specially since no
+			// information about finalized blocks is persisted within the tree.
+			//
+			// We use part of the information contained in the finalized block
+			// children (that are expected to be in the tree) to construct a
+			// dummy block entry for the last finalized block. This will be
+			// wiped as soon as the next block is finalized.
+
+			let blocks = backend.load_blocks_by_number(first_number)?;
+
+			let block = blocks
+				.first()
+				.and_then(|hash| backend.load_block_entry(hash).ok())
+				.flatten()
+				.ok_or_else(|| {
+					ChainApiError::from(format!(
+						"Lookup failure for block at height {}",
+						first_number
+					))
+				})?;
+
+			// The parent is expected to be the last finalized block.
+			if block.parent_hash != hash {
+				return Err(ChainApiError::from("Can't revert below last finalized block").into())
+			}
+
+			// The weight is set to the one of the first child. Even though this is
+			// not accurate, it does the job. The reason is that the revert point is
+			// the last finalized block, i.e. this is the best and only choice.
+			let block_number = first_number.saturating_sub(1);
+			let viability = ViabilityCriteria {
+				explicitly_reverted: false,
+				approval: Approval::Approved,
+				earliest_unviable_ancestor: None,
+			};
+			let entry = BlockEntry {
+				block_hash: hash,
+				block_number,
+				parent_hash: Hash::default(),
+				children: blocks,
+				viability,
+				weight: block.weight,
+			};
+			// This becomes the first entry according to the block number.
+			backend.write_blocks_by_number(block_number, vec![hash]);
+			entry
+		},
+	};
+
+	let mut stack: Vec<_> = std::mem::take(&mut entry.children)
+		.into_iter()
+		.map(|h| (h, entry.block_number + 1))
+		.collect();
+
+	// Write revert point block entry without the children.
+	backend.write_block_entry(entry.clone());
+
+	let mut viable_leaves = backend.load_leaves()?;
+
+	viable_leaves.insert(LeafEntry {
+		block_hash: hash,
+		block_number: entry.block_number,
+		weight: entry.weight,
+	});
+
+	while let Some((hash, number)) = stack.pop() {
+		let entry = backend.load_block_entry(&hash)?;
+		backend.delete_block_entry(&hash);
+
+		viable_leaves.remove(&hash);
+
+		let mut blocks_at_height = backend.load_blocks_by_number(number)?;
+		blocks_at_height.retain(|h| h != &hash);
+		backend.write_blocks_by_number(number, blocks_at_height);
+
+		stack.extend(entry.into_iter().flat_map(|e| e.children).map(|h| (h, number + 1)));
+	}
+
+	backend.write_leaves(viable_leaves);
+
+	Ok(backend)
+}