Limit number of blocks per level (2nd attempt)

client/consensus/aura/src/import_queue.rs

@@ -17,7 +17,7 @@
 //! Parachain specific wrapper for the AuRa import queue.
 
 use codec::Codec;
-use sc_client_api::{backend::AuxStore, BlockOf, UsageProvider};
+use sc_client_api::{backend::AuxStore, Backend, BlockOf, UsageProvider};
 use sc_consensus::{import_queue::DefaultImportQueue, BlockImport};
 use sc_consensus_aura::AuraVerifier;
 use sc_consensus_slots::InherentDataProviderExt;
@@ -34,7 +34,7 @@ use std::{fmt::Debug, hash::Hash, sync::Arc};
 use substrate_prometheus_endpoint::Registry;
 
 /// Parameters of [`import_queue`].
-pub struct ImportQueueParams<'a, I, C, CIDP, S, CAW> {
+pub struct ImportQueueParams<'a, I, C, CIDP, S, CAW, BE> {
 	/// The block import to use.
 	pub block_import: I,
 	/// The client to interact with the chain.
@@ -49,10 +49,12 @@ pub struct ImportQueueParams<'a, I, C, CIDP, S, CAW> {
 	pub can_author_with: CAW,
 	/// The telemetry handle.
 	pub telemetry: Option<TelemetryHandle>,
+	/// The backend.
+	pub backend: Arc<BE>,
 }
 
 /// Start an import queue for the Aura consensus algorithm.
-pub fn import_queue<'a, P, Block, I, C, S, CAW, CIDP>(
+pub fn import_queue<'a, P, Block, I, C, S, CAW, CIDP, BE>(
 	ImportQueueParams {
 		block_import,
 		client,
@@ -61,7 +63,8 @@ pub fn import_queue<'a, P, Block, I, C, S, CAW, CIDP>(
 		registry,
 		can_author_with,
 		telemetry,
-	}: ImportQueueParams<'a, I, C, CIDP, S, CAW>,
+		backend,
+	}: ImportQueueParams<'a, I, C, CIDP, S, CAW, BE>,
 ) -> Result<DefaultImportQueue<Block, C>, sp_consensus::Error>
 where
 	Block: BlockT,
@@ -85,9 +88,12 @@ where
 	CAW: CanAuthorWith<Block> + Send + Sync + 'static,
 	CIDP: CreateInherentDataProviders<Block, ()> + Sync + Send + 'static,
 	CIDP::InherentDataProviders: InherentDataProviderExt + Send + Sync,
+	BE: Backend<Block> + 'static,
 {
+	use cumulus_client_consensus_common::{LeavesLevelLimit, ParachainBlockImport};
+
 	sc_consensus_aura::import_queue::<P, _, _, _, _, _, _>(sc_consensus_aura::ImportQueueParams {
-		block_import: cumulus_client_consensus_common::ParachainBlockImport::new(block_import),
+		block_import: ParachainBlockImport::new(block_import, backend, LeavesLevelLimit::Default),
 		justification_import: None,
 		client,
 		create_inherent_data_providers,

client/consensus/aura/src/lib.rs

@@ -24,12 +24,12 @@
 
 use codec::{Decode, Encode};
 use cumulus_client_consensus_common::{
-	ParachainBlockImport, ParachainCandidate, ParachainConsensus,
+	LeavesLevelLimit, ParachainBlockImport, ParachainCandidate, ParachainConsensus,
 };
 use cumulus_primitives_core::{relay_chain::v2::Hash as PHash, PersistedValidationData};
 
 use futures::lock::Mutex;
-use sc_client_api::{backend::AuxStore, BlockOf};
+use sc_client_api::{backend::AuxStore, Backend, BlockOf};
 use sc_consensus::BlockImport;
 use sc_consensus_slots::{BackoffAuthoringBlocksStrategy, SimpleSlotWorker, SlotInfo};
 use sc_telemetry::TelemetryHandle;
@@ -78,7 +78,7 @@ where
 	CIDP::InherentDataProviders: InherentDataProviderExt,
 {
 	/// Create a new boxed instance of AURA consensus.
-	pub fn build<P, Client, BI, SO, PF, BS, Error>(
+	pub fn build<P, Client, BI, SO, PF, BS, Error, BE>(
 		BuildAuraConsensusParams {
 			proposer_factory,
 			create_inherent_data_providers,
@@ -92,7 +92,8 @@ where
 			telemetry,
 			block_proposal_slot_portion,
 			max_block_proposal_slot_portion,
-		}: BuildAuraConsensusParams<PF, BI, CIDP, Client, BS, SO>,
+			backend,
+		}: BuildAuraConsensusParams<PF, BI, CIDP, Client, BS, SO, BE>,
 	) -> Box<dyn ParachainConsensus<B>>
 	where
 		Client:
@@ -113,11 +114,16 @@ where
 		P: Pair + Send + Sync,
 		P::Public: AppPublic + Hash + Member + Encode + Decode,
 		P::Signature: TryFrom<Vec<u8>> + Hash + Member + Encode + Decode,
+		BE: Backend<B> + 'static,
 	{
 		let worker = sc_consensus_aura::build_aura_worker::<P, _, _, _, _, _, _, _, _>(
 			BuildAuraWorkerParams {
 				client: para_client,
-				block_import: ParachainBlockImport::new(block_import),
+				block_import: ParachainBlockImport::new(
+					block_import,
+					backend,
+					LeavesLevelLimit::Default,
+				),
 				justification_sync_link: (),
 				proposer_factory,
 				sync_oracle,
@@ -219,7 +225,7 @@ where
 }
 
 /// Parameters of [`AuraConsensus::build`].
-pub struct BuildAuraConsensusParams<PF, BI, CIDP, Client, BS, SO> {
+pub struct BuildAuraConsensusParams<PF, BI, CIDP, Client, BS, SO, BE> {
 	pub proposer_factory: PF,
 	pub create_inherent_data_providers: CIDP,
 	pub block_import: BI,
@@ -232,4 +238,5 @@ pub struct BuildAuraConsensusParams<PF, BI, CIDP, Client, BS, SO> {
 	pub telemetry: Option<TelemetryHandle>,
 	pub block_proposal_slot_portion: SlotProportion,
 	pub max_block_proposal_slot_portion: Option<SlotProportion>,
+	pub backend: Arc<BE>,
 }

client/consensus/common/Cargo.toml

@@ -11,6 +11,7 @@ codec = { package = "parity-scale-codec", version = "3.0.0", features = [ "deriv
 dyn-clone = "1.0.9"
 futures = "0.3.23"
 tracing = "0.1.36"
+log = "0.4.17"
 
 # Substrate
 sc-client-api = { git = "https://github.com/paritytech/substrate", branch = "master" }

client/consensus/common/src/lib.rs

@@ -15,12 +15,24 @@
 // along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
 
 use polkadot_primitives::v2::{Hash as PHash, PersistedValidationData};
+
+use sc_client_api::{
+	blockchain::{Backend as _, HeaderBackend as _},
+	Backend,
+};
 use sc_consensus::BlockImport;
-use sp_runtime::traits::Block as BlockT;
+use sp_api::NumberFor;
+use sp_runtime::traits::{Block as BlockT, Header as HeaderT};
+
+use std::{
+	collections::{HashMap, HashSet},
+	sync::Arc,
+};
 
 mod parachain_consensus;
 #[cfg(test)]
 mod tests;
+
 pub use parachain_consensus::run_parachain_consensus;
 
 /// The result of [`ParachainConsensus::produce_candidate`].
@@ -69,46 +81,173 @@ impl<B: BlockT> ParachainConsensus<B> for Box<dyn ParachainConsensus<B> + Send +
 	}
 }
 
+/// Value good enough to be used with parachains using the current backend implementation
+/// that ships with Substrate. This value may change in the future.
+pub const MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT: usize = 32;
+
+/// Upper bound to the number of leaves allowed for each level of the blockchain.
+///
+/// If the limit is set and more leaves are detected on block import, then the older ones are
+/// dropped to make space for the fresh blocks.
+///
+/// In environments where blocks confirmations from the relay chain may be "slow", then
+/// setting an upper bound helps keeping the chain health by dropping old (presumably) stale
+/// leaves and prevents discarding new blocks because we've reached the backend max value.
+pub enum LeavesLevelLimit {
+	/// Limit set to `MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT`.
+	Default,
+	/// No explicit limit, however a limit may be implicitly imposed by the backend implementation.
+	None,
+	/// Custom value.
+	Some(usize),
+}
+
+type BlockHash<Block> = <<Block as BlockT>::Header as HeaderT>::Hash;
+
+// Support structure to constrain the number of leaves for each level.
+struct LeavesLevelMonitor<Block: BlockT, BE> {
+	// Max number of leaves for each level.
+	level_limit: usize,
+	// Monotonic counter used to keep track of block import age (bigger is younger).
+	import_counter: u64,
+	// Map between block hash and age.
+	leaves_cache: HashMap<BlockHash<Block>, u64>,
+	// Backend reference to remove leaves on level saturation.
+	backend: Arc<BE>,
+}
+
+// Threshold after which we are going to cleanup our internal leaves cache by removing hashes that
+// doesn't belong to leaf blocks anymore.
+// This is a farly arbitrary value that can be changed in the future without breaking anything.
+const CLEANUP_THRESHOLD: usize = 64;
+
+impl<Block: BlockT, BE: Backend<Block>> LeavesLevelMonitor<Block, BE> {
+	fn update(&mut self, hash: BlockHash<Block>, parent_hash: BlockHash<Block>) {
+		self.import_counter += 1;
+		self.leaves_cache.insert(hash, self.import_counter);
+		self.leaves_cache.remove(&parent_hash);
+	}
+
+	fn check(&mut self, number: NumberFor<Block>) {
+		let blockchain = self.backend.blockchain();
+		let mut leaves = blockchain.leaves().unwrap_or_default();
+
+		if self.leaves_cache.len().saturating_sub(leaves.len()) >= CLEANUP_THRESHOLD {
+			// Update the cache once in a while by using the leaves set returned by the backend.
+			// Using a temporary HashSet we reduce iterations from O(n^2) to O(2n)
+			let leaves_set: HashSet<_> = leaves.iter().collect();
+			self.leaves_cache.retain(|hash, _| leaves_set.contains(hash));
+		}
+
+		// First cheap check: the number of leaves at level `number` is always less than the total.
+		if leaves.len() < self.level_limit {
+			return
+		}
+
+		// Now focus on the leaves at the given height.
+		leaves.retain(|hash| {
+			blockchain.number(*hash).ok().flatten().map(|n| n == number).unwrap_or_default()
+		});
+
+		if leaves.len() < self.level_limit {
+			return
+		}
+
+		log::debug!(
+			target: "parachain",
+			"Detected leaves overflow at height {}, removing old blocks", number
+		);
+
+		let best = blockchain.info().best_hash;
+		let mut remove_count = leaves.len() - self.level_limit + 1;
+
+		// Sort leaves by import chronological order
+		leaves.sort_unstable_by(|a, b| self.leaves_cache.get(a).cmp(&self.leaves_cache.get(b)));
+
+		for hash in leaves.into_iter().filter(|hash| *hash != best) {
+			log::debug!(target: "parachain", "Removing block {}", hash);
+			if self.backend.remove_leaf_block(&hash).is_err() {
+				log::warn!(target: "parachain", "Unable to remove block {}, skipping it...", hash);
+				continue
+			}
+			self.leaves_cache.remove(&hash);
+			remove_count -= 1;
+			if remove_count == 0 {
+				break
+			}
+		}
+	}
+}
+
 /// Parachain specific block import.
 ///
 /// This is used to set `block_import_params.fork_choice` to `false` as long as the block origin is
 /// not `NetworkInitialSync`. The best block for parachains is determined by the relay chain. Meaning
 /// we will update the best block, as it is included by the relay-chain.
-pub struct ParachainBlockImport<I>(I);
+pub struct ParachainBlockImport<Block: BlockT, BI, BE> {
+	inner: BI,
+	leaves_monitor: Option<LeavesLevelMonitor<Block, BE>>,
+}
 
-impl<I> ParachainBlockImport<I> {
+impl<Block: BlockT, BI, BE> ParachainBlockImport<Block, BI, BE> {
 	/// Create a new instance.
-	pub fn new(inner: I) -> Self {
-		Self(inner)
+	pub fn new(inner: BI, backend: Arc<BE>, level_leaves_max: LeavesLevelLimit) -> Self {
+		let level_limit = match level_leaves_max {
+			LeavesLevelLimit::None => None,
+			LeavesLevelLimit::Some(limit) => Some(limit),
+			LeavesLevelLimit::Default => Some(MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT),
+		};
+		let leaves_monitor = level_limit.map(|limit| LeavesLevelMonitor {
+			level_limit: limit,
+			leaves_cache: HashMap::new(),
+			import_counter: 0,
+			backend,
+		});
+		Self { inner, leaves_monitor }
 	}
 }
 
 #[async_trait::async_trait]
-impl<Block, I> BlockImport<Block> for ParachainBlockImport<I>
+impl<Block, BI, BE> BlockImport<Block> for ParachainBlockImport<Block, BI, BE>
 where
 	Block: BlockT,
-	I: BlockImport<Block> + Send,
+	BI: BlockImport<Block> + Send,
+	BE: Backend<Block>,
 {
-	type Error = I::Error;
-	type Transaction = I::Transaction;
+	type Error = BI::Error;
+	type Transaction = BI::Transaction;
 
 	async fn check_block(
 		&mut self,
 		block: sc_consensus::BlockCheckParams<Block>,
 	) -> Result<sc_consensus::ImportResult, Self::Error> {
-		self.0.check_block(block).await
+		self.inner.check_block(block).await
 	}
 
 	async fn import_block(
 		&mut self,
-		mut block_import_params: sc_consensus::BlockImportParams<Block, Self::Transaction>,
+		mut params: sc_consensus::BlockImportParams<Block, Self::Transaction>,
 		cache: std::collections::HashMap<sp_consensus::CacheKeyId, Vec<u8>>,
 	) -> Result<sc_consensus::ImportResult, Self::Error> {
+		let hash = params.header.hash();
+		let parent = *params.header.parent_hash();
+
+		if let Some(ref mut monitor) = self.leaves_monitor {
+			monitor.check(*params.header.number());
+		}
+
 		// Best block is determined by the relay chain, or if we are doing the initial sync
 		// we import all blocks as new best.
-		block_import_params.fork_choice = Some(sc_consensus::ForkChoiceStrategy::Custom(
-			block_import_params.origin == sp_consensus::BlockOrigin::NetworkInitialSync,
+		params.fork_choice = Some(sc_consensus::ForkChoiceStrategy::Custom(
+			params.origin == sp_consensus::BlockOrigin::NetworkInitialSync,
 		));
-		self.0.import_block(block_import_params, cache).await
+
+		let res = self.inner.import_block(params, cache).await?;
+
+		if let Some(ref mut monitor) = self.leaves_monitor {
+			monitor.update(hash, parent);
+		}
+
+		Ok(res)
 	}
 }