Revert approval voting

node/core/approval-voting/src/approval_db/v1/mod.rs

@@ -95,6 +95,9 @@ impl Backend for DbBackend {
 						stored_block_range.encode(),
 					);
 				},
+				BackendWriteOp::DeleteStoredBlockRange => {
+					tx.delete(self.config.col_data, &STORED_BLOCKS_KEY);
+				},
 				BackendWriteOp::WriteBlocksAtHeight(h, blocks) => {
 					tx.put_vec(self.config.col_data, &blocks_at_height_key(h), blocks.encode());
 				},

node/core/approval-voting/src/backend.rs

@@ -37,6 +37,7 @@ pub enum BackendWriteOp {
 	WriteBlocksAtHeight(BlockNumber, Vec<Hash>),
 	WriteBlockEntry(BlockEntry),
 	WriteCandidateEntry(CandidateEntry),
+	DeleteStoredBlockRange,
 	DeleteBlocksAtHeight(BlockNumber),
 	DeleteBlockEntry(Hash),
 	DeleteCandidateEntry(CandidateHash),
@@ -63,16 +64,23 @@ pub trait Backend {
 		I: IntoIterator<Item = BackendWriteOp>;
 }
 
+// Status of block range in the `OverlayedBackend`.
+#[derive(PartialEq)]
+enum BlockRangeStatus {
+	NotModified,
+	Deleted,
+	Inserted(StoredBlockRange),
+}
+
 /// An in-memory overlay over the backend.
 ///
 /// This maintains read-only access to the underlying backend, but can be
 /// converted into a set of write operations which will, when written to
 /// the underlying backend, give the same view as the state of the overlay.
 pub struct OverlayedBackend<'a, B: 'a> {
 	inner: &'a B,
-
-	// `None` means unchanged
-	stored_block_range: Option<StoredBlockRange>,
+	// `Some(None)` means deleted. Missing (`None`) means query inner.
+	stored_block_range: BlockRangeStatus,
 	// `None` means 'deleted', missing means query inner.
 	blocks_at_height: HashMap<BlockNumber, Option<Vec<Hash>>>,
 	// `None` means 'deleted', missing means query inner.
@@ -85,7 +93,7 @@ impl<'a, B: 'a + Backend> OverlayedBackend<'a, B> {
 	pub fn new(backend: &'a B) -> Self {
 		OverlayedBackend {
 			inner: backend,
-			stored_block_range: None,
+			stored_block_range: BlockRangeStatus::NotModified,
 			blocks_at_height: HashMap::new(),
 			block_entries: HashMap::new(),
 			candidate_entries: HashMap::new(),
@@ -96,7 +104,7 @@ impl<'a, B: 'a + Backend> OverlayedBackend<'a, B> {
 		self.block_entries.is_empty() &&
 			self.candidate_entries.is_empty() &&
 			self.blocks_at_height.is_empty() &&
-			self.stored_block_range.is_none()
+			self.stored_block_range == BlockRangeStatus::NotModified
 	}
 
 	pub fn load_all_blocks(&self) -> SubsystemResult<Vec<Hash>> {
@@ -111,11 +119,11 @@ impl<'a, B: 'a + Backend> OverlayedBackend<'a, B> {
 	}
 
 	pub fn load_stored_blocks(&self) -> SubsystemResult<Option<StoredBlockRange>> {
-		if let Some(val) = self.stored_block_range.clone() {
-			return Ok(Some(val))
+		match self.stored_block_range {
+			BlockRangeStatus::Inserted(ref value) => Ok(Some(value.clone())),
+			BlockRangeStatus::Deleted => Ok(None),
+			BlockRangeStatus::NotModified => self.inner.load_stored_blocks(),
 		}
-
-		self.inner.load_stored_blocks()
 	}
 
 	pub fn load_blocks_at_height(&self, height: &BlockNumber) -> SubsystemResult<Vec<Hash>> {
@@ -145,10 +153,12 @@ impl<'a, B: 'a + Backend> OverlayedBackend<'a, B> {
 		self.inner.load_candidate_entry(candidate_hash)
 	}
 
-	// The assumption is that stored block range is only None on initialization.
-	// Therefore, there is no need to delete_stored_block_range.
 	pub fn write_stored_block_range(&mut self, range: StoredBlockRange) {
-		self.stored_block_range = Some(range);
+		self.stored_block_range = BlockRangeStatus::Inserted(range);
+	}
+
+	pub fn delete_stored_block_range(&mut self) {
+		self.stored_block_range = BlockRangeStatus::Deleted;
 	}
 
 	pub fn write_blocks_at_height(&mut self, height: BlockNumber, blocks: Vec<Hash>) {
@@ -193,8 +203,13 @@ impl<'a, B: 'a + Backend> OverlayedBackend<'a, B> {
 			None => BackendWriteOp::DeleteCandidateEntry(h),
 		});
 
-		self.stored_block_range
-			.map(|v| BackendWriteOp::WriteStoredBlockRange(v))
+		let stored_block_range_ops = match self.stored_block_range {
+			BlockRangeStatus::Inserted(val) => Some(BackendWriteOp::WriteStoredBlockRange(val)),
+			BlockRangeStatus::Deleted => Some(BackendWriteOp::DeleteStoredBlockRange),
+			BlockRangeStatus::NotModified => None,
+		};
+
+		stored_block_range_ops
 			.into_iter()
 			.chain(blocks_at_height_ops)
 			.chain(block_entry_ops)

node/core/approval-voting/src/lib.rs

@@ -340,6 +340,19 @@ impl ApprovalVotingSubsystem {
 			metrics,
 		}
 	}
+
+	/// Revert to the block corresponding to the specified `hash`.
+	/// The operation is not allowed for blocks older than the last finalized one.
+	pub fn revert_to(&self, hash: Hash) -> Result<(), SubsystemError> {
+		let config = approval_db::v1::Config { col_data: self.db_config.col_data };
+		let mut backend = approval_db::v1::DbBackend::new(self.db.clone(), config);
+		let mut overlay = OverlayedBackend::new(&backend);
+
+		ops::revert_to(&mut overlay, hash)?;
+
+		let ops = overlay.into_write_ops();
+		backend.write(ops)
+	}
 }
 
 impl<Context> overseer::Subsystem<Context, SubsystemError> for ApprovalVotingSubsystem

node/core/approval-voting/src/ops.rs

@@ -17,7 +17,7 @@
 //! Middleware interface that leverages low-level database operations
 //! to provide a clean API for processing block and candidate imports.
 
-use polkadot_node_subsystem::SubsystemResult;
+use polkadot_node_subsystem::{SubsystemError, SubsystemResult};
 
 use bitvec::order::Lsb0 as BitOrderLsb0;
 use polkadot_primitives::v2::{BlockNumber, CandidateHash, CandidateReceipt, GroupIndex, Hash};
@@ -311,3 +311,92 @@ pub fn force_approve(
 
 	Ok(approved_hashes)
 }
+
+/// Revert to the block corresponding to the specified `hash`.
+/// The operation is not allowed for blocks older than the last finalized one.
+pub fn revert_to(
+	overlay: &mut OverlayedBackend<'_, impl Backend>,
+	hash: Hash,
+) -> SubsystemResult<()> {
+	let mut stored_range = overlay.load_stored_blocks()?.ok_or_else(|| {
+		SubsystemError::Context("no available blocks to infer revert point height".to_string())
+	})?;
+
+	let (children, children_height) = match overlay.load_block_entry(&hash)? {
+		Some(mut entry) => {
+			let children_height = entry.block_number() + 1;
+			let children = std::mem::take(&mut entry.children);
+			// Write revert point block entry without the children.
+			overlay.write_block_entry(entry);
+			(children, children_height)
+		},
+		None => {
+			let children_height = stored_range.0;
+			let children = overlay.load_blocks_at_height(&children_height)?;
+
+			let child_entry = children
+				.first()
+				.and_then(|hash| overlay.load_block_entry(hash).ok())
+				.flatten()
+				.ok_or_else(|| {
+					SubsystemError::Context("lookup failure for first block".to_string())
+				})?;
+
+			// The parent is expected to be the revert point
+			if child_entry.parent_hash() != hash {
+				return Err(SubsystemError::Context(
+					"revert below last finalized block or corrupted storage".to_string(),
+				))
+			}
+
+			(children, children_height)
+		},
+	};
+
+	let mut stack: Vec<_> = children.into_iter().map(|h| (h, children_height)).collect();
+	let mut range_end = stored_range.1;
+
+	while let Some((hash, number)) = stack.pop() {
+		let mut blocks_at_height = overlay.load_blocks_at_height(&number)?;
+		blocks_at_height.retain(|h| h != &hash);
+
+		// Check if we need to update the range top
+		if blocks_at_height.is_empty() && number < range_end {
+			range_end = number;
+		}
+
+		overlay.write_blocks_at_height(number, blocks_at_height);
+
+		if let Some(entry) = overlay.load_block_entry(&hash)? {
+			overlay.delete_block_entry(&hash);
+
+			// Cleanup the candidate entries by removing any reference to the
+			// removed block. If for a candidate entry the block block_assignments
+			// drops to zero then we remove the entry.
+			for (_, candidate_hash) in entry.candidates() {
+				if let Some(mut candidate_entry) = overlay.load_candidate_entry(candidate_hash)? {
+					candidate_entry.block_assignments.remove(&hash);
+					if candidate_entry.block_assignments.is_empty() {
+						overlay.delete_candidate_entry(candidate_hash);
+					} else {
+						overlay.write_candidate_entry(candidate_entry);
+					}
+				}
+			}
+
+			stack.extend(entry.children.into_iter().map(|h| (h, number + 1)));
+		}
+	}
+
+	// Check if our modifications to the dag has reduced the range top
+	if range_end != stored_range.1 {
+		if stored_range.0 < range_end {
+			stored_range.1 = range_end;
+			overlay.write_stored_block_range(stored_range);
+		} else {
+			overlay.delete_stored_block_range();
+		}
+	}
+
+	Ok(())
+}

node/core/approval-voting/src/tests.rs

@@ -306,6 +306,9 @@ impl Backend for TestStoreInner {
 				BackendWriteOp::WriteStoredBlockRange(stored_block_range) => {
 					self.stored_block_range = Some(stored_block_range);
 				},
+				BackendWriteOp::DeleteStoredBlockRange => {
+					self.stored_block_range = None;
+				},
 				BackendWriteOp::WriteBlocksAtHeight(h, blocks) => {
 					self.blocks_at_height.insert(h, blocks);
 				},

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

@@ -317,10 +317,10 @@ impl ChainSelectionSubsystem {
 	}
 
 	/// 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);
+	/// The operation is not allowed for blocks older than the last finalized one.
+	pub fn revert_to(&self, hash: Hash) -> Result<(), Error> {
+		let config = db_backend::v1::Config { col_data: self.config.col_data };
+		let mut backend = db_backend::v1::DbBackend::new(self.db.clone(), config);
 
 		let ops = tree::revert_to(&backend, hash)?.into_write_ops();
 

node/service/src/lib.rs

@@ -37,7 +37,9 @@ use {
 	beefy_gadget::notification::{BeefyBestBlockSender, BeefySignedCommitmentSender},
 	grandpa::{self, FinalityProofProvider as GrandpaFinalityProofProvider},
 	gum::info,
-	polkadot_node_core_approval_voting::Config as ApprovalVotingConfig,
+	polkadot_node_core_approval_voting::{
+		self as approval_voting_subsystem, Config as ApprovalVotingConfig,
+	},
 	polkadot_node_core_av_store::Config as AvailabilityConfig,
 	polkadot_node_core_av_store::Error as AvailabilityError,
 	polkadot_node_core_candidate_validation::Config as CandidateValidationConfig,
@@ -1403,32 +1405,11 @@ pub fn build_full(
 	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.
+/// - `ApprovalVotingSubsystem` data in the parachains-db;
+/// - `ChainSelectionSubsystem` data in the parachains-db;
 /// - Low level Babe and Grandpa consensus data.
 #[cfg(feature = "full-node")]
 pub fn revert_backend(
@@ -1441,6 +1422,10 @@ pub fn revert_backend(
 	let finalized = client.info().finalized_number;
 	let revertible = blocks.min(best_number - finalized);
 
+	if revertible == 0 {
+		return Ok(())
+	}
+
 	let number = best_number - revertible;
 	let hash = client.block_hash_from_id(&BlockId::Number(number))?.ok_or(
 		sp_blockchain::Error::Backend(format!(
@@ -1452,19 +1437,66 @@ pub fn revert_backend(
 	let parachains_db = open_database(&config.database)
 		.map_err(|err| sp_blockchain::Error::Backend(err.to_string()))?;
 
+	revert_approval_voting(parachains_db.clone(), hash)?;
+	revert_chain_selection(parachains_db, hash)?;
+	// Revert Substrate consensus related components
+	client.execute_with(RevertConsensus { blocks, backend })?;
+
+	Ok(())
+}
+
+fn revert_chain_selection(db: Arc<dyn Database>, hash: Hash) -> sp_blockchain::Result<()> {
 	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);
+	let chain_selection = chain_selection_subsystem::ChainSelectionSubsystem::new(config, db);
 
 	chain_selection
-		.revert(hash)
-		.map_err(|err| sp_blockchain::Error::Backend(err.to_string()))?;
+		.revert_to(hash)
+		.map_err(|err| sp_blockchain::Error::Backend(err.to_string()))
+}
 
-	client.execute_with(RevertConsensus { blocks, backend })?;
+fn revert_approval_voting(db: Arc<dyn Database>, hash: Hash) -> sp_blockchain::Result<()> {
+	let config = approval_voting_subsystem::Config {
+		col_data: parachains_db::REAL_COLUMNS.col_approval_data,
+		slot_duration_millis: Default::default(),
+	};
 
-	Ok(())
+	let approval_voting = approval_voting_subsystem::ApprovalVotingSubsystem::with_config(
+		config,
+		db,
+		Arc::new(sc_keystore::LocalKeystore::in_memory()),
+		Box::new(consensus_common::NoNetwork),
+		approval_voting_subsystem::Metrics::default(),
+	);
+
+	approval_voting
+		.revert_to(hash)
+		.map_err(|err| sp_blockchain::Error::Backend(err.to_string()))
+}
+
+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,
+	{
+		// Revert consensus-related components.
+		// The operations are not correlated, thus call order is not relevant.
+		babe::revert(client.clone(), self.backend, self.blocks)?;
+		grandpa::revert(client, self.blocks)?;
+		Ok(())
+	}
 }