Enforce MaxCandidates limit for staking storage items

pallets/parachain-staking/src/lib.rs

@@ -80,7 +80,7 @@ pub mod pallet {
 	use crate::delegation_requests::{
 		CancelledScheduledRequest, DelegationAction, ScheduledRequest,
 	};
-	use crate::{set::OrderedSet, traits::*, types::*, InflationInfo, Range, WeightInfo};
+	use crate::{set::BoundedOrderedSet, traits::*, types::*, InflationInfo, Range, WeightInfo};
 	use crate::{AutoCompoundConfig, AutoCompoundDelegations};
 	use frame_support::pallet_prelude::*;
 	use frame_support::traits::{
@@ -177,6 +177,9 @@ pub mod pallet {
 		type OnNewRound: OnNewRound;
 		/// Weight information for extrinsics in this pallet.
 		type WeightInfo: WeightInfo;
+		/// Maximum candidates
+		#[pallet::constant]
+		type MaxCandidates: Get<u32>;
 	}
 
 	#[pallet::error]
@@ -229,8 +232,8 @@ pub mod pallet {
 		TooLowCandidateAutoCompoundingDelegationCountToLeaveCandidates,
 		TooLowCandidateCountWeightHint,
 		TooLowCandidateCountWeightHintGoOffline,
-		TooLowCandidateCountWeightHintGoOnline,
-		TooLowCandidateCountWeightHintCandidateBondMore,
+		CandidateLimitReached,
+		CannotSetAboveMaxCandidates,
 	}
 
 	#[pallet::event]
@@ -579,7 +582,8 @@ pub mod pallet {
 	#[pallet::storage]
 	#[pallet::getter(fn selected_candidates)]
 	/// The collator candidates selected for the current round
-	type SelectedCandidates<T: Config> = StorageValue<_, Vec<T::AccountId>, ValueQuery>;
+	type SelectedCandidates<T: Config> =
+		StorageValue<_, BoundedVec<T::AccountId, T::MaxCandidates>, ValueQuery>;
 
 	#[pallet::storage]
 	#[pallet::getter(fn total)]
@@ -589,8 +593,11 @@ pub mod pallet {
 	#[pallet::storage]
 	#[pallet::getter(fn candidate_pool)]
 	/// The pool of collator candidates, each with their total backing stake
-	pub(crate) type CandidatePool<T: Config> =
-		StorageValue<_, OrderedSet<Bond<T::AccountId, BalanceOf<T>>>, ValueQuery>;
+	pub(crate) type CandidatePool<T: Config> = StorageValue<
+		_,
+		BoundedOrderedSet<Bond<T::AccountId, BalanceOf<T>>, T::MaxCandidates>,
+		ValueQuery,
+	>;
 
 	#[pallet::storage]
 	#[pallet::getter(fn at_stake)]
@@ -764,6 +771,11 @@ pub mod pallet {
 				"{:?}",
 				Error::<T>::CannotSetBelowMin
 			);
+			assert!(
+				self.num_selected_candidates <= T::MaxCandidates::get(),
+				"{:?}",
+				Error::<T>::CannotSetAboveMaxCandidates
+			);
 			<TotalSelected<T>>::put(self.num_selected_candidates);
 			// Choose top TotalSelected collator candidates
 			let (_, v_count, _, total_staked) = <Pallet<T>>::select_top_candidates(1u32);
@@ -886,6 +898,10 @@ pub mod pallet {
 				new >= T::MinSelectedCandidates::get(),
 				Error::<T>::CannotSetBelowMin
 			);
+			ensure!(
+				new <= T::MaxCandidates::get(),
+				Error::<T>::CannotSetAboveMaxCandidates
+			);
 			let old = <TotalSelected<T>>::get();
 			ensure!(old != new, Error::<T>::NoWritingSameValue);
 			ensure!(
@@ -970,13 +986,14 @@ pub mod pallet {
 				candidate_count >= old_count,
 				Error::<T>::TooLowCandidateCountWeightHintJoinCandidates
 			);
-			ensure!(
-				candidates.insert(Bond {
+			let maybe_inserted_candidate = candidates
+				.try_insert(Bond {
 					owner: acc.clone(),
-					amount: bond
-				}),
-				Error::<T>::CandidateExists
-			);
+					amount: bond,
+				})
+				.map_err(|_| Error::<T>::CandidateLimitReached)?;
+			ensure!(maybe_inserted_candidate, Error::<T>::CandidateExists);
+
 			ensure!(
 				Self::get_collator_stakable_free_balance(&acc) >= bond,
 				Error::<T>::InsufficientBalance,
@@ -1065,13 +1082,13 @@ pub mod pallet {
 				candidates.0.len() as u32 <= candidate_count,
 				Error::<T>::TooLowCandidateCountWeightHintCancelLeaveCandidates
 			);
-			ensure!(
-				candidates.insert(Bond {
+			let maybe_inserted_candidate = candidates
+				.try_insert(Bond {
 					owner: collator.clone(),
-					amount: state.total_counted
-				}),
-				Error::<T>::AlreadyActive
-			);
+					amount: state.total_counted,
+				})
+				.map_err(|_| Error::<T>::CandidateLimitReached)?;
+			ensure!(maybe_inserted_candidate, Error::<T>::AlreadyActive);
 			<CandidatePool<T>>::put(candidates);
 			<CandidateInfo<T>>::insert(&collator, state);
 			Self::deposit_event(Event::CancelledCandidateExit {
@@ -1447,16 +1464,20 @@ pub mod pallet {
 			);
 			state.go_online();
 
-			ensure!(
-				candidates.insert(Bond {
+			let maybe_inserted_candidate = candidates
+				.try_insert(Bond {
 					owner: collator.clone(),
-					amount: state.total_counted
-				}),
+					amount: state.total_counted,
+				})
+				.map_err(|_| Error::<T>::CandidateLimitReached)?;
+			ensure!(
+				maybe_inserted_candidate,
 				DispatchErrorWithPostInfo {
 					post_info: Some(actual_weight).into(),
 					error: <Error<T>>::AlreadyActive.into(),
 				},
 			);
+
 			<CandidatePool<T>>::put(candidates);
 			<CandidateInfo<T>>::insert(&collator, state);
 			Self::deposit_event(Event::CandidateBackOnline {
@@ -1619,10 +1640,14 @@ pub mod pallet {
 		pub(crate) fn update_active(candidate: T::AccountId, total: BalanceOf<T>) {
 			let mut candidates = <CandidatePool<T>>::get();
 			candidates.remove(&Bond::from_owner(candidate.clone()));
-			candidates.insert(Bond {
-				owner: candidate,
-				amount: total,
-			});
+			candidates
+				.try_insert(Bond {
+					owner: candidate,
+					amount: total,
+				})
+				.expect(
+					"the candidate is removed in previous step so the length cannot increase; qed",
+				);
 			<CandidatePool<T>>::put(candidates);
 		}
 
@@ -1853,27 +1878,31 @@ pub mod pallet {
 				return vec![];
 			}
 
-			let mut candidates = <CandidatePool<T>>::get().0;
+			let candidates = <CandidatePool<T>>::get().0;
 
 			// If the number of candidates is greater than top_n, select the candidates with higher
 			// amount. Otherwise, return all the candidates.
 			if candidates.len() > top_n {
 				// Partially sort candidates such that element at index `top_n - 1` is sorted, and
 				// all the elements in the range 0..top_n are the top n elements.
-				candidates.select_nth_unstable_by(top_n - 1, |a, b| {
-					// Order by amount, then owner. The owner is needed to ensure a stable order
-					// when two accounts have the same amount.
-					a.amount
-						.cmp(&b.amount)
-						.then_with(|| a.owner.cmp(&b.owner))
-						.reverse()
-				});
+				let sorted_candidates = candidates
+					.try_mutate(|inner| {
+						inner.select_nth_unstable_by(top_n - 1, |a, b| {
+							// Order by amount, then owner. The owner is needed to ensure a stable order
+							// when two accounts have the same amount.
+							a.amount
+								.cmp(&b.amount)
+								.then_with(|| a.owner.cmp(&b.owner))
+								.reverse()
+						});
+					})
+					.expect("sort cannot increase item count; qed");
 
-				let mut collators = candidates
+				let mut collators = sorted_candidates
 					.into_iter()
 					.take(top_n)
 					.map(|x| x.owner)
-					.collect::<Vec<T::AccountId>>();
+					.collect::<Vec<_>>();
 
 				// Sort collators by AccountId
 				collators.sort();
@@ -1882,10 +1911,7 @@ pub mod pallet {
 			} else {
 				// Return all candidates
 				// The candidates are already sorted by AccountId, so no need to sort again
-				candidates
-					.into_iter()
-					.map(|x| x.owner)
-					.collect::<Vec<T::AccountId>>()
+				candidates.into_iter().map(|x| x.owner).collect::<Vec<_>>()
 			}
 		}
 		/// Best as in most cumulatively supported in terms of stake
@@ -1967,7 +1993,10 @@ pub mod pallet {
 				});
 			}
 			// insert canonical collator set
-			<SelectedCandidates<T>>::put(collators);
+			<SelectedCandidates<T>>::put(
+				BoundedVec::try_from(collators)
+					.expect("subset of collators is always less than or equal to max candidates"),
+			);
 
 			let avg_delegator_count = delegation_count.checked_div(collator_count).unwrap_or(0);
 			let weight = T::WeightInfo::select_top_candidates(collator_count, avg_delegator_count);
@@ -2149,7 +2178,7 @@ pub mod pallet {
 
 	impl<T: Config> Get<Vec<T::AccountId>> for Pallet<T> {
 		fn get() -> Vec<T::AccountId> {
-			Self::selected_candidates()
+			Self::selected_candidates().into_inner()
 		}
 	}
 }

pallets/parachain-staking/src/mock.rs

@@ -120,6 +120,7 @@ parameter_types! {
 	pub const MinCandidateStk: u128 = 10;
 	pub const MinDelegatorStk: u128 = 5;
 	pub const MinDelegation: u128 = 3;
+	pub const MaxCandidates: u32 = 10;
 }
 impl Config for Test {
 	type RuntimeEvent = RuntimeEvent;
@@ -144,6 +145,7 @@ impl Config for Test {
 	type PayoutCollatorReward = ();
 	type OnNewRound = ();
 	type WeightInfo = ();
+	type MaxCandidates = MaxCandidates;
 }
 
 pub(crate) struct ExtBuilder {

pallets/parachain-staking/src/set.rs

@@ -15,11 +15,13 @@
 // along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
 
 /* TODO: use orml_utilities::OrderedSet without leaking substrate v2.0 dependencies*/
+
+use frame_support::traits::Get;
 use parity_scale_codec::{Decode, Encode};
 use scale_info::TypeInfo;
 #[cfg(feature = "std")]
 use serde::{Deserialize, Serialize};
-use sp_runtime::RuntimeDebug;
+use sp_runtime::{BoundedVec, RuntimeDebug};
 use sp_std::prelude::*;
 
 /// An ordered set backed by `Vec`
@@ -87,3 +89,78 @@ impl<T: Ord> From<Vec<T>> for OrderedSet<T> {
 		Self::from(v)
 	}
 }
+
+/// An ordered set backed by `BoundedVec`
+#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
+#[derive(RuntimeDebug, PartialEq, Eq, Encode, Decode, Clone, TypeInfo)]
+#[scale_info(skip_type_params(S))]
+pub struct BoundedOrderedSet<T, S: Get<u32>>(pub BoundedVec<T, S>);
+
+impl<T, S: Get<u32>> sp_std::default::Default for BoundedOrderedSet<T, S> {
+	fn default() -> Self {
+		BoundedOrderedSet(BoundedVec::default())
+	}
+}
+
+impl<T: Ord, S: Get<u32>> BoundedOrderedSet<T, S> {
+	/// Create a new empty set
+	pub fn new() -> Self {
+		Self(BoundedVec::default())
+	}
+
+	/// Create a set from a `Vec`.
+	/// `v` will be sorted and dedup first.
+	pub fn from(mut v: BoundedVec<T, S>) -> Self {
+		v.sort();
+		let v = v.try_mutate(|inner| inner.dedup()).expect(
+			"input is a valid BoundedVec and deduping can only reduce the number of entires; qed",
+		);
+		Self::from_sorted_set(v)
+	}
+
+	/// Create a set from a `Vec`.
+	/// Assume `v` is sorted and contain unique elements.
+	pub fn from_sorted_set(v: BoundedVec<T, S>) -> Self {
+		Self(v)
+	}
+
+	/// Insert an element.
+	/// Return true if insertion happened.
+	pub fn try_insert(&mut self, value: T) -> Result<bool, ()> {
+		match self.0.binary_search(&value) {
+			Ok(_) => Ok(false),
+			Err(loc) => self.0.try_insert(loc, value).map(|_| true).map_err(|_| ()),
+		}
+	}
+
+	/// Remove an element.
+	/// Return true if removal happened.
+	pub fn remove(&mut self, value: &T) -> bool {
+		match self.0.binary_search(value) {
+			Ok(loc) => {
+				self.0.remove(loc);
+				true
+			}
+			Err(_) => false,
+		}
+	}
+
+	/// Return if the set contains `value`
+	pub fn contains(&self, value: &T) -> bool {
+		self.0.binary_search(value).is_ok()
+	}
+
+	/// Clear the set
+	pub fn clear(mut self) {
+		let v = self.0.try_mutate(|inner| inner.clear()).expect(
+			"input is a valid BoundedVec and clearing can only reduce the number of entires; qed",
+		);
+		self.0 = v;
+	}
+}
+
+impl<T: Ord, S: Get<u32>> From<BoundedVec<T, S>> for BoundedOrderedSet<T, S> {
+	fn from(v: BoundedVec<T, S>) -> Self {
+		Self::from(v)
+	}
+}