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/cli/src/command.rs

@@ -435,7 +435,19 @@ pub fn run() -> Result<()> {
 
 			Ok(runner.async_run(|mut config| {
 				let (client, backend, _, task_manager) = service::new_chain_ops(&mut config, None)?;
-				Ok((cmd.run(client, backend, None).map_err(Error::SubstrateCli), task_manager))
+				let aux_revert = Box::new(|client, backend, blocks| {
+					service::revert_backend(client, backend, blocks, config).map_err(|err| {
+						match err {
+							service::Error::Blockchain(err) => err.into(),
+							// Generic application-specific error.
+							err => sc_cli::Error::Application(err.into()),
+						}
+					})
+				});
+				Ok((
+					cmd.run(client, backend, Some(aux_revert)).map_err(Error::SubstrateCli),
+					task_manager,
+				))
 			})?)
 		},
 		Some(Subcommand::PvfPrepareWorker(cmd)) => {

polkadot/node/client/src/lib.rs

@@ -26,7 +26,7 @@ use polkadot_primitives::{
 use sc_client_api::{AuxStore, Backend as BackendT, BlockchainEvents, KeyIterator, UsageProvider};
 use sc_executor::NativeElseWasmExecutor;
 use sp_api::{CallApiAt, Encode, NumberFor, ProvideRuntimeApi};
-use sp_blockchain::HeaderBackend;
+use sp_blockchain::{HeaderBackend, HeaderMetadata};
 use sp_consensus::BlockStatus;
 use sp_core::Pair;
 use sp_keyring::Sr25519Keyring;
@@ -173,6 +173,7 @@ pub trait AbstractClient<Block, Backend>:
 	+ CallApiAt<Block, StateBackend = Backend::State>
 	+ AuxStore
 	+ UsageProvider<Block>
+	+ HeaderMetadata<Block, Error = sp_blockchain::Error>
 where
 	Block: BlockT,
 	Backend: BackendT<Block>,
@@ -194,7 +195,8 @@ where
 		+ Sized
 		+ Send
 		+ Sync
-		+ CallApiAt<Block, StateBackend = Backend::State>,
+		+ CallApiAt<Block, StateBackend = Backend::State>
+		+ HeaderMetadata<Block, Error = sp_blockchain::Error>,
 	Client::Api: RuntimeApiCollection<StateBackend = Backend::State>,
 {
 }

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)
+}

polkadot/node/service/src/lib.rs

@@ -50,6 +50,8 @@ use {
 	sp_trie::PrefixedMemoryDB,
 };
 
+use polkadot_node_subsystem_util::database::Database;
+
 pub use sp_core::traits::SpawnNamed;
 #[cfg(feature = "full-node")]
 pub use {
@@ -58,7 +60,7 @@ pub use {
 	relay_chain_selection::SelectRelayChain,
 	sc_client_api::AuxStore,
 	sp_authority_discovery::AuthorityDiscoveryApi,
-	sp_blockchain::HeaderBackend,
+	sp_blockchain::{HeaderBackend, HeaderMetadata},
 	sp_consensus_babe::BabeApi,
 };
 
@@ -94,7 +96,7 @@ pub use polkadot_client::{
 	AbstractClient, Client, ClientHandle, ExecuteWithClient, FullBackend, FullClient,
 	RuntimeApiCollection,
 };
-pub use polkadot_primitives::v2::{Block, BlockId, CollatorPair, Hash, Id as ParaId};
+pub use polkadot_primitives::v2::{Block, BlockId, BlockNumber, CollatorPair, Hash, Id as ParaId};
 pub use sc_client_api::{Backend, CallExecutor, ExecutionStrategy};
 pub use sc_consensus::{BlockImport, LongestChain};
 use sc_executor::NativeElseWasmExecutor;
@@ -285,6 +287,36 @@ impl IdentifyVariant for Box<dyn ChainSpec> {
 	}
 }
 
+#[cfg(feature = "full-node")]
+fn open_database(db_source: &DatabaseSource) -> Result<Arc<dyn Database>, Error> {
+	let parachains_db = match db_source {
+		DatabaseSource::RocksDb { path, .. } => parachains_db::open_creating_rocksdb(
+			path.clone(),
+			parachains_db::CacheSizes::default(),
+		)?,
+		DatabaseSource::ParityDb { path, .. } => parachains_db::open_creating_paritydb(
+			path.parent().ok_or(Error::DatabasePathRequired)?.into(),
+			parachains_db::CacheSizes::default(),
+		)?,
+		DatabaseSource::Auto { paritydb_path, rocksdb_path, .. } =>
+			if paritydb_path.is_dir() && paritydb_path.exists() {
+				parachains_db::open_creating_paritydb(
+					paritydb_path.parent().ok_or(Error::DatabasePathRequired)?.into(),
+					parachains_db::CacheSizes::default(),
+				)?
+			} else {
+				parachains_db::open_creating_rocksdb(
+					rocksdb_path.clone(),
+					parachains_db::CacheSizes::default(),
+				)?
+			},
+		DatabaseSource::Custom { .. } => {
+			unimplemented!("No polkadot subsystem db for custom source.");
+		},
+	};
+	Ok(parachains_db)
+}
+
 /// Initialize the `Jeager` collector. The destination must listen
 /// on the given address and port for `UDP` packets.
 #[cfg(any(test, feature = "full-node"))]
@@ -866,39 +898,15 @@ where
 		);
 	}
 
-	let parachains_db = match &config.database {
-		DatabaseSource::RocksDb { path, .. } => crate::parachains_db::open_creating_rocksdb(
-			path.clone(),
-			crate::parachains_db::CacheSizes::default(),
-		)?,
-		DatabaseSource::ParityDb { path, .. } => crate::parachains_db::open_creating_paritydb(
-			path.parent().ok_or(Error::DatabasePathRequired)?.into(),
-			crate::parachains_db::CacheSizes::default(),
-		)?,
-		DatabaseSource::Auto { paritydb_path, rocksdb_path, .. } =>
-			if paritydb_path.is_dir() && paritydb_path.exists() {
-				crate::parachains_db::open_creating_paritydb(
-					paritydb_path.parent().ok_or(Error::DatabasePathRequired)?.into(),
-					crate::parachains_db::CacheSizes::default(),
-				)?
-			} else {
-				crate::parachains_db::open_creating_rocksdb(
-					rocksdb_path.clone(),
-					crate::parachains_db::CacheSizes::default(),
-				)?
-			},
-		DatabaseSource::Custom { .. } => {
-			unimplemented!("No polkadot subsystem db for custom source.");
-		},
-	};
+	let parachains_db = open_database(&config.database)?;
 
 	let availability_config = AvailabilityConfig {
-		col_data: crate::parachains_db::REAL_COLUMNS.col_availability_data,
-		col_meta: crate::parachains_db::REAL_COLUMNS.col_availability_meta,
+		col_data: parachains_db::REAL_COLUMNS.col_availability_data,
+		col_meta: parachains_db::REAL_COLUMNS.col_availability_meta,
 	};
 
 	let approval_voting_config = ApprovalVotingConfig {
-		col_data: crate::parachains_db::REAL_COLUMNS.col_approval_data,
+		col_data: parachains_db::REAL_COLUMNS.col_approval_data,
 		slot_duration_millis: slot_duration.as_millis() as u64,
 	};
 
@@ -915,12 +923,12 @@ where
 	};
 
 	let chain_selection_config = ChainSelectionConfig {
-		col_data: crate::parachains_db::REAL_COLUMNS.col_chain_selection_data,
+		col_data: parachains_db::REAL_COLUMNS.col_chain_selection_data,
 		stagnant_check_interval: chain_selection_subsystem::StagnantCheckInterval::never(),
 	};
 
 	let dispute_coordinator_config = DisputeCoordinatorConfig {
-		col_data: crate::parachains_db::REAL_COLUMNS.col_dispute_coordinator_data,
+		col_data: parachains_db::REAL_COLUMNS.col_dispute_coordinator_data,
 	};
 
 	let rpc_handlers = service::spawn_tasks(service::SpawnTasksParams {
@@ -1394,3 +1402,69 @@ pub fn build_full(
 	#[cfg(not(feature = "polkadot-native"))]
 	Err(Error::NoRuntime)
 }
+
+struct RevertConsensus {
+	blocks: BlockNumber,
+	backend: Arc<FullBackend>,
+}
+
+impl ExecuteWithClient for RevertConsensus {
+	type Output = sp_blockchain::Result<()>;
+
+	fn execute_with_client<Client, Api, Backend>(self, client: Arc<Client>) -> Self::Output
+	where
+		<Api as sp_api::ApiExt<Block>>::StateBackend: sp_api::StateBackend<BlakeTwo256>,
+		Backend: sc_client_api::Backend<Block> + 'static,
+		Backend::State: sp_api::StateBackend<BlakeTwo256>,
+		Api: polkadot_client::RuntimeApiCollection<StateBackend = Backend::State>,
+		Client: AbstractClient<Block, Backend, Api = Api> + 'static,
+	{
+		babe::revert(client.clone(), self.backend, self.blocks)?;
+		grandpa::revert(client, self.blocks)?;
+		Ok(())
+	}
+}
+
+/// Reverts the node state down to at most the last finalized block.
+///
+/// In particular this reverts:
+/// - `ChainSelectionSubsystem` data in the parachains-db.
+/// - Low level Babe and Grandpa consensus data.
+#[cfg(feature = "full-node")]
+pub fn revert_backend(
+	client: Arc<Client>,
+	backend: Arc<FullBackend>,
+	blocks: BlockNumber,
+	config: Configuration,
+) -> Result<(), Error> {
+	let best_number = client.info().best_number;
+	let finalized = client.info().finalized_number;
+	let revertible = blocks.min(best_number - finalized);
+
+	let number = best_number - revertible;
+	let hash = client.block_hash_from_id(&BlockId::Number(number))?.ok_or(
+		sp_blockchain::Error::Backend(format!(
+			"Unexpected hash lookup failure for block number: {}",
+			number
+		)),
+	)?;
+
+	let parachains_db = open_database(&config.database)
+		.map_err(|err| sp_blockchain::Error::Backend(err.to_string()))?;
+
+	let config = chain_selection_subsystem::Config {
+		col_data: parachains_db::REAL_COLUMNS.col_chain_selection_data,
+		stagnant_check_interval: chain_selection_subsystem::StagnantCheckInterval::never(),
+	};
+
+	let chain_selection =
+		chain_selection_subsystem::ChainSelectionSubsystem::new(config, parachains_db);
+
+	chain_selection
+		.revert(hash)
+		.map_err(|err| sp_blockchain::Error::Backend(err.to_string()))?;
+
+	client.execute_with(RevertConsensus { blocks, backend })?;
+
+	Ok(())
+}

polkadot/node/service/src/relay_chain_selection.rs

@@ -501,7 +501,7 @@ where
 				match rx.await.map_err(Error::DetermineUndisputedChainCanceled) {
 					// If request succeded we will receive (block number, block hash).
 					Ok((subchain_number, subchain_head)) => {
-						// The the total lag accounting for disputes.
+						// The total lag accounting for disputes.
 						let lag_disputes = initial_leaf_number.saturating_sub(subchain_number);
 						self.metrics.note_disputes_finality_lag(lag_disputes);
 						(lag_disputes, subchain_head)