client/beefy: add some bounds on enqueued votes

client/beefy/src/lib.rs

@@ -1,3 +1,4 @@
+#![feature(map_try_insert)]
 // This file is part of Substrate.
 
 // Copyright (C) 2021-2022 Parity Technologies (UK) Ltd.
@@ -24,7 +25,7 @@ use sc_consensus::BlockImport;
 use sc_network::ProtocolName;
 use sc_network_common::service::NetworkRequest;
 use sc_network_gossip::Network as GossipNetwork;
-use sp_api::{NumberFor, ProvideRuntimeApi};
+use sp_api::{BlockT, HeaderT, NumberFor, ProvideRuntimeApi};
 use sp_blockchain::HeaderBackend;
 use sp_consensus::{Error as ConsensusError, SyncOracle};
 use sp_keystore::SyncCryptoStorePtr;
@@ -202,6 +203,8 @@ where
 	R: ProvideRuntimeApi<B>,
 	R::Api: BeefyApi<B> + MmrApi<B, MmrRootHash, NumberFor<B>>,
 	N: GossipNetwork<B> + NetworkRequest + SyncOracle + Send + Sync + 'static,
+	u64: From<NumberFor<B>>,
+	<<B as BlockT>::Header as HeaderT>::Number: From<u64>,
 {
 	let BeefyParams {
 		client,

client/beefy/src/worker.rs

@@ -35,22 +35,23 @@ use sc_network_common::{
 };
 use sc_network_gossip::GossipEngine;
 
-use sp_api::{BlockId, ProvideRuntimeApi};
-use sp_arithmetic::traits::{AtLeast32Bit, Saturating};
+use sp_api::{BlockId, BlockT, HeaderT, ProvideRuntimeApi};
+use sp_arithmetic::traits::AtLeast32Bit;
 use sp_blockchain::Backend as BlockchainBackend;
 use sp_consensus::SyncOracle;
 use sp_mmr_primitives::MmrApi;
 use sp_runtime::{
 	generic::OpaqueDigestItemId,
-	traits::{Block, Header, NumberFor},
-	SaturatedConversion,
+	traits::{Block, NumberFor},
+	BoundedBTreeMap, SaturatedConversion,
 };
 
 use beefy_primitives::{
 	crypto::{AuthorityId, Signature},
 	BeefyApi, Commitment, ConsensusLog, MmrRootHash, Payload, PayloadProvider, SignedCommitment,
 	ValidatorSet, VersionedFinalityProof, VoteMessage, BEEFY_ENGINE_ID, GENESIS_AUTHORITY_SET_ID,
 };
+use sp_runtime::traits::{CheckedConversion, ConstU32};
 
 use crate::{
 	communication::{
@@ -66,6 +67,11 @@ use crate::{
 	BeefyVoterLinks, Client, KnownPeers,
 };
 
+/// Depicts the bound for the number of pending votes
+const MAX_PENDING_VOTES: u32 = 50;
+/// Depicts the bound for the number of pending justifications
+const MAX_PENDING_JUSTIFICATIONS: u32 = 50;
+
 enum RoundAction {
 	Drop,
 	Process,
@@ -238,9 +244,17 @@ pub(crate) struct BeefyWorker<B: Block, BE, C, P, R, N> {
 	/// Best block a BEEFY voting round has been concluded for.
 	best_beefy_block: Option<NumberFor<B>>,
 	/// Buffer holding votes for future processing.
-	pending_votes: BTreeMap<NumberFor<B>, Vec<VoteMessage<NumberFor<B>, AuthorityId, Signature>>>,
+	pending_votes: BoundedBTreeMap<
+		NumberFor<B>,
+		Vec<VoteMessage<NumberFor<B>, AuthorityId, Signature>>,
+		ConstU32<MAX_PENDING_VOTES>,
+	>,
 	/// Buffer holding justifications for future processing.
-	pending_justifications: BTreeMap<NumberFor<B>, BeefyVersionedFinalityProof<B>>,
+	pending_justifications: BoundedBTreeMap<
+		NumberFor<B>,
+		BeefyVersionedFinalityProof<B>,
+		ConstU32<MAX_PENDING_JUSTIFICATIONS>,
+	>,
 	/// Chooses which incoming votes to accept and which votes to generate.
 	voting_oracle: VoterOracle<B>,
 }
@@ -254,6 +268,8 @@ where
 	R: ProvideRuntimeApi<B>,
 	R::Api: BeefyApi<B> + MmrApi<B, MmrRootHash, NumberFor<B>>,
 	N: NetworkEventStream + NetworkRequest + SyncOracle + Send + Sync + Clone + 'static,
+	u64: From<NumberFor<B>>,
+	<<B as BlockT>::Header as HeaderT>::Number: From<u64>,
 {
 	/// Return a new BEEFY worker instance.
 	///
@@ -298,8 +314,8 @@ where
 			metrics,
 			best_grandpa_block_header: last_finalized_header,
 			best_beefy_block: None,
-			pending_votes: BTreeMap::new(),
-			pending_justifications: BTreeMap::new(),
+			pending_votes: BoundedBTreeMap::new(),
+			pending_justifications: BoundedBTreeMap::new(),
 			voting_oracle: VoterOracle::new(min_block_delta),
 		}
 	}
@@ -409,7 +425,19 @@ where
 			)?,
 			RoundAction::Enqueue => {
 				debug!(target: "beefy", "🥩 Buffer vote for round: {:?}.", block_num);
-				self.pending_votes.entry(block_num).or_default().push(vote)
+				if self.pending_votes.remove(&block_num).is_some() {
+					let mut vec_of_votes = self.pending_votes.remove(&block_num).unwrap();
+					vec_of_votes.push(vote);
+					if self.pending_votes.try_insert(block_num, vec_of_votes).is_err() {
+						warn!(target: "beefy", "🥩 Buffer vote dropped for round: {:?}.", block_num)
+					}
+				} else {
+					let mut vec_of_votes = vec![];
+					vec_of_votes.push(vote);
+					if self.pending_votes.try_insert(block_num, vec_of_votes).is_err() {
+						warn!(target: "beefy", "🥩 Buffer vote dropped for round: {:?}.", block_num)
+					}
+				}
 			},
 			RoundAction::Drop => (),
 		};
@@ -435,7 +463,9 @@ where
 			},
 			RoundAction::Enqueue => {
 				debug!(target: "beefy", "🥩 Buffer justification for round: {:?}.", block_num);
-				self.pending_justifications.entry(block_num).or_insert(justification);
+				if self.pending_justifications.try_insert(block_num, justification).is_err() {
+					error!(target: "beefy", "🥩 Buffer justification dropped for round: {:?}.", block_num);
+				}
 			},
 			RoundAction::Drop => (),
 		};
@@ -529,18 +559,41 @@ where
 		let best_grandpa = *self.best_grandpa_block_header.number();
 		let _ph = PhantomData::<B>::default();
 
-		fn to_process_for<B: Block, T>(
-			pending: &mut BTreeMap<NumberFor<B>, T>,
+		fn to_process_for<B: Block, T, const U: u32>(
+			pending: &mut BoundedBTreeMap<NumberFor<B>, T, ConstU32<U>>,
 			(start, end): (NumberFor<B>, NumberFor<B>),
 			_: PhantomData<B>,
-		) -> BTreeMap<NumberFor<B>, T> {
-			// These are still pending.
-			let still_pending = pending.split_off(&end.saturating_add(1u32.into()));
-			// These can be processed.
-			let to_handle = pending.split_off(&start);
-			// The rest can be dropped.
-			*pending = still_pending;
-			// Return ones to process.
+		) -> BTreeMap<NumberFor<B>, T>
+		where
+			u64: From<NumberFor<B>>,
+			<<B as BlockT>::Header as HeaderT>::Number: From<u64>,
+		{
+			let mut to_handle = BTreeMap::new();
+
+			let mut still_pending = BTreeMap::new();
+
+			let end: u64 = end.into();
+			let start: u64 = start.into();
+			let still_pending_range = end.saturating_add(1u32.into())..U.into();
+
+			for i in still_pending_range {
+				let value_to_be_removed = pending.remove(&i.into());
+				if value_to_be_removed.is_some() {
+					still_pending.insert(i.into(), value_to_be_removed.unwrap());
+				}
+			}
+
+			let to_handle_range = start..=end;
+
+			for i in to_handle_range {
+				let value_to_be_removed = pending.remove(&i.into());
+				if value_to_be_removed.is_some() {
+					to_handle.insert(i.into(), value_to_be_removed.unwrap());
+				}
+			}
+
+			*pending = BoundedBTreeMap::checked_from(still_pending).expect("Should not fail");
+
 			to_handle
 		}
 		// Interval of blocks for which we can process justifications and votes right now.
@@ -1441,7 +1494,6 @@ pub(crate) mod tests {
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
 		let mut net = BeefyTestNet::new(1);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1);
-
 		fn new_vote(
 			block_number: NumberFor<Block>,
 		) -> VoteMessage<NumberFor<Block>, AuthorityId, Signature> {