Add Ability to Add/Remove Invulnerable Collators (#2596)

* add add and remove invulnerable dispatchables

* add migration

* fix benchmarking code

* add weights

* add migration to runtimes

* clippy

* pass SafeCallFilter

* make try-runtime work

* typos

Co-authored-by: Ankan <10196091+Ank4n@users.noreply.github.com>

* better insert and added test

* fix try-runtime update

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <git@kchr.de>

* Update pallets/collator-selection/src/migration.rs

* check events in test

* Update pallets/collator-selection/src/migration.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* just dispatchresult

* only sp_std for try-runtime

---------

Co-authored-by: Ankan <10196091+Ank4n@users.noreply.github.com>
Co-authored-by: Bastian Köcher <git@kchr.de>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

pallets/collator-selection/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
-authors = ["Anonymous"]
-description = "Simple staking pallet with a fixed stake."
+authors = ["Parity Technologies <admin@parity.io>"]
+description = "Simple pallet to select collators for a parachain."
 edition = "2021"
 homepage = "https://substrate.io"
 license = "Apache-2.0"

pallets/collator-selection/src/benchmarking.rs

@@ -110,17 +110,61 @@ benchmarks! {
 	where_clause { where T: pallet_authorship::Config + session::Config }
 
 	set_invulnerables {
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
 		let b in 1 .. T::MaxInvulnerables::get();
 		let new_invulnerables = register_validators::<T>(b);
+		let mut sorted_new_invulnerables = new_invulnerables.clone();
+		sorted_new_invulnerables.sort();
+	}: {
+		assert_ok!(
+			// call the function with the unsorted list
+			<CollatorSelection<T>>::set_invulnerables(origin, new_invulnerables.clone())
+		);
+	}
+	verify {
+		// assert that it comes out sorted
+		assert_last_event::<T>(Event::NewInvulnerables{invulnerables: sorted_new_invulnerables}.into());
+	}
+
+	add_invulnerable {
 		let origin =
 			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+		// we're going to add one, so need one less than max set as invulnerables to start
+		let b in 1 .. T::MaxInvulnerables::get() - 1;
+		let mut invulnerables = register_validators::<T>(b);
+		invulnerables.sort();
+		let invulnerables: frame_support::BoundedVec<_, T::MaxInvulnerables> = frame_support::BoundedVec::try_from(invulnerables).unwrap();
+		<Invulnerables<T>>::put(invulnerables);
+
+		// now let's set up a new one to add
+		let (new, keys) = validator::<T>(b + 1);
+		<session::Pallet<T>>::set_keys(RawOrigin::Signed(new.clone()).into(), keys, Vec::new()).unwrap();
 	}: {
 		assert_ok!(
-			<CollatorSelection<T>>::set_invulnerables(origin, new_invulnerables.clone())
+			<CollatorSelection<T>>::add_invulnerable(origin, new.clone())
+		);
+	}
+	verify {
+		assert_last_event::<T>(Event::InvulnerableAdded{account_id: new}.into());
+	}
+
+	remove_invulnerable {
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+		let b in 1 .. T::MaxInvulnerables::get();
+		let mut invulnerables = register_validators::<T>(b);
+		invulnerables.sort();
+		let invulnerables: frame_support::BoundedVec<_, T::MaxInvulnerables> = frame_support::BoundedVec::try_from(invulnerables).unwrap();
+		<Invulnerables<T>>::put(invulnerables);
+		let to_remove = <Invulnerables<T>>::get().first().unwrap().clone();
+	}: {
+		assert_ok!(
+			<CollatorSelection<T>>::remove_invulnerable(origin, to_remove.clone())
 		);
 	}
 	verify {
-		assert_last_event::<T>(Event::NewInvulnerables{invulnerables: new_invulnerables}.into());
+		assert_last_event::<T>(Event::InvulnerableRemoved{account_id: to_remove}.into());
 	}
 
 	set_desired_candidates {

pallets/collator-selection/src/lib.rs

@@ -70,8 +70,11 @@ mod tests;
 
 #[cfg(feature = "runtime-benchmarks")]
 mod benchmarking;
+pub mod migration;
 pub mod weights;
 
+const LOG_TARGET: &str = "runtime::collator-selection";
+
 #[frame_support::pallet]
 pub mod pallet {
 	pub use crate::weights::WeightInfo;
@@ -95,6 +98,9 @@ pub mod pallet {
 	use sp_runtime::traits::Convert;
 	use sp_staking::SessionIndex;
 
+	/// The current storage version.
+	const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+
 	type BalanceOf<T> =
 		<<T as Config>::Currency as Currency<<T as SystemConfig>::AccountId>>::Balance;
 
@@ -165,9 +171,10 @@ pub mod pallet {
 	}
 
 	#[pallet::pallet]
+	#[pallet::storage_version(STORAGE_VERSION)]
 	pub struct Pallet<T>(_);
 
-	/// The invulnerable, fixed collators.
+	/// The invulnerable, permissioned collators. This list must be sorted.
 	#[pallet::storage]
 	#[pallet::getter(fn invulnerables)]
 	pub type Invulnerables<T: Config> =
@@ -222,14 +229,16 @@ pub mod pallet {
 				"duplicate invulnerables in genesis."
 			);
 
-			let bounded_invulnerables =
+			let mut bounded_invulnerables =
 				BoundedVec::<_, T::MaxInvulnerables>::try_from(self.invulnerables.clone())
 					.expect("genesis invulnerables are more than T::MaxInvulnerables");
 			assert!(
 				T::MaxCandidates::get() >= self.desired_candidates,
 				"genesis desired_candidates are more than T::MaxCandidates",
 			);
 
+			bounded_invulnerables.sort();
+
 			<DesiredCandidates<T>>::put(self.desired_candidates);
 			<CandidacyBond<T>>::put(self.candidacy_bond);
 			<Invulnerables<T>>::put(bounded_invulnerables);
@@ -239,35 +248,41 @@ pub mod pallet {
 	#[pallet::event]
 	#[pallet::generate_deposit(pub(super) fn deposit_event)]
 	pub enum Event<T: Config> {
+		/// New Invulnerables were set.
 		NewInvulnerables { invulnerables: Vec<T::AccountId> },
+		/// A new Invulnerable was added.
+		InvulnerableAdded { account_id: T::AccountId },
+		/// An Invulnerable was removed.
+		InvulnerableRemoved { account_id: T::AccountId },
+		/// The number of desired candidates was set.
 		NewDesiredCandidates { desired_candidates: u32 },
+		/// The candidacy bond was set.
 		NewCandidacyBond { bond_amount: BalanceOf<T> },
+		/// A new candidate joined.
 		CandidateAdded { account_id: T::AccountId, deposit: BalanceOf<T> },
+		/// A candidate was removed.
 		CandidateRemoved { account_id: T::AccountId },
 	}
 
-	// Errors inform users that something went wrong.
 	#[pallet::error]
 	pub enum Error<T> {
-		/// Too many candidates
+		/// The pallet has too many candidates.
 		TooManyCandidates,
-		/// Too few candidates
+		/// Leaving would result in too few candidates.
 		TooFewCandidates,
-		/// Unknown error
-		Unknown,
-		/// Permission issue
-		Permission,
-		/// User is already a candidate
+		/// Account is already a candidate.
 		AlreadyCandidate,
-		/// User is not a candidate
+		/// Account is not a candidate.
 		NotCandidate,
-		/// Too many invulnerables
+		/// There are too many Invulnerables.
 		TooManyInvulnerables,
-		/// User is already an Invulnerable
+		/// Account is already an Invulnerable.
 		AlreadyInvulnerable,
-		/// Account has no associated validator ID
+		/// Account is not an Invulnerable.
+		NotInvulnerable,
+		/// Account has no associated validator ID.
 		NoAssociatedValidatorId,
-		/// Validator ID is not yet registered
+		/// Validator ID is not yet registered.
 		ValidatorNotRegistered,
 	}
 
@@ -284,7 +299,7 @@ pub mod pallet {
 			new: Vec<T::AccountId>,
 		) -> DispatchResultWithPostInfo {
 			T::UpdateOrigin::ensure_origin(origin)?;
-			let bounded_invulnerables = BoundedVec::<_, T::MaxInvulnerables>::try_from(new)
+			let mut bounded_invulnerables = BoundedVec::<_, T::MaxInvulnerables>::try_from(new)
 				.map_err(|_| Error::<T>::TooManyInvulnerables)?;
 
 			// check if the invulnerables have associated validator keys before they are set
@@ -297,6 +312,9 @@ pub mod pallet {
 				);
 			}
 
+			// Invulnerables must be sorted for removal.
+			bounded_invulnerables.sort();
+
 			<Invulnerables<T>>::put(&bounded_invulnerables);
 			Self::deposit_event(Event::NewInvulnerables {
 				invulnerables: bounded_invulnerables.to_vec(),
@@ -398,6 +416,56 @@ pub mod pallet {
 
 			Ok(Some(T::WeightInfo::leave_intent(current_count as u32)).into())
 		}
+
+		/// Add a new account `who` to the list of `Invulnerables` collators.
+		///
+		/// The origin for this call must be the `UpdateOrigin`.
+		#[pallet::call_index(5)]
+		#[pallet::weight(T::WeightInfo::add_invulnerable(T::MaxInvulnerables::get() - 1))]
+		pub fn add_invulnerable(origin: OriginFor<T>, who: T::AccountId) -> DispatchResult {
+			T::UpdateOrigin::ensure_origin(origin)?;
+
+			// ensure `who` has registered a validator key
+			let validator_key = T::ValidatorIdOf::convert(who.clone())
+				.ok_or(Error::<T>::NoAssociatedValidatorId)?;
+			ensure!(
+				T::ValidatorRegistration::is_registered(&validator_key),
+				Error::<T>::ValidatorNotRegistered
+			);
+
+			<Invulnerables<T>>::try_mutate(|invulnerables| -> DispatchResult {
+				match invulnerables.binary_search(&who) {
+					Ok(_) => return Err(Error::<T>::AlreadyInvulnerable)?,
+					Err(pos) => invulnerables
+						.try_insert(pos, who.clone())
+						.map_err(|_| Error::<T>::TooManyInvulnerables)?,
+				}
+				Ok(())
+			})?;
+
+			Self::deposit_event(Event::InvulnerableAdded { account_id: who });
+			Ok(())
+		}
+
+		/// Remove an account `who` from the list of `Invulnerables` collators. `Invulnerables` must
+		/// be sorted.
+		///
+		/// The origin for this call must be the `UpdateOrigin`.
+		#[pallet::call_index(6)]
+		#[pallet::weight(T::WeightInfo::remove_invulnerable(T::MaxInvulnerables::get()))]
+		pub fn remove_invulnerable(origin: OriginFor<T>, who: T::AccountId) -> DispatchResult {
+			T::UpdateOrigin::ensure_origin(origin)?;
+
+			<Invulnerables<T>>::try_mutate(|invulnerables| -> DispatchResult {
+				let pos =
+					invulnerables.binary_search(&who).map_err(|_| Error::<T>::NotInvulnerable)?;
+				invulnerables.remove(pos);
+				Ok(())
+			})?;
+
+			Self::deposit_event(Event::InvulnerableRemoved { account_id: who });
+			Ok(())
+		}
 	}
 
 	impl<T: Config> Pallet<T> {

pallets/collator-selection/src/migration.rs

@@ -0,0 +1,91 @@
+// Copyright Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! A module that is responsible for migration of storage for Collator Selection.
+
+use super::*;
+use frame_support::{log, traits::OnRuntimeUpgrade};
+
+/// Version 1 Migration
+/// This migration ensures that any existing `Invulnerables` storage lists are sorted.
+pub mod v1 {
+	use super::*;
+	use frame_support::pallet_prelude::*;
+	#[cfg(feature = "try-runtime")]
+	use sp_std::prelude::*;
+
+	pub struct MigrateToV1<T>(sp_std::marker::PhantomData<T>);
+	impl<T: Config> OnRuntimeUpgrade for MigrateToV1<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let onchain_version = Pallet::<T>::on_chain_storage_version();
+			if onchain_version == 0 {
+				let invulnerables_len = Invulnerables::<T>::get().to_vec().len();
+				<Invulnerables<T>>::mutate(|invulnerables| {
+					invulnerables.sort();
+				});
+
+				StorageVersion::new(1).put::<Pallet<T>>();
+				log::info!(
+					target: LOG_TARGET,
+					"Sorted {} Invulnerables, upgraded storage to version 1",
+					invulnerables_len,
+				);
+				// Similar complexity to `set_invulnerables` (put storage value)
+				// Plus 1 read for length, 1 read for `onchain_version`, 1 write to put version
+				T::WeightInfo::set_invulnerables(invulnerables_len as u32)
+					.saturating_add(T::DbWeight::get().reads_writes(2, 1))
+			} else {
+				log::info!(
+					target: LOG_TARGET,
+					"Migration did not execute. This probably should be removed"
+				);
+				T::DbWeight::get().reads(1)
+			}
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::DispatchError> {
+			let number_of_invulnerables = Invulnerables::<T>::get().to_vec().len();
+			Ok((number_of_invulnerables as u32).encode())
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(number_of_invulnerables: Vec<u8>) -> Result<(), sp_runtime::DispatchError> {
+			let stored_invulnerables = Invulnerables::<T>::get().to_vec();
+			let mut sorted_invulnerables = stored_invulnerables.clone();
+			sorted_invulnerables.sort();
+			assert_eq!(
+				stored_invulnerables, sorted_invulnerables,
+				"after migration, the stored invulnerables should be sorted"
+			);
+
+			let number_of_invulnerables: u32 = Decode::decode(
+				&mut number_of_invulnerables.as_slice(),
+			)
+			.expect("the state parameter should be something that was generated by pre_upgrade");
+			let stored_invulnerables_len = stored_invulnerables.len() as u32;
+			assert_eq!(
+				number_of_invulnerables, stored_invulnerables_len,
+				"after migration, there should be the same number of invulnerables"
+			);
+
+			let onchain_version = Pallet::<T>::on_chain_storage_version();
+			frame_support::ensure!(onchain_version >= 1, "must_upgrade");
+
+			Ok(())
+		}
+	}
+}

pallets/collator-selection/src/mock.rs

@@ -195,7 +195,7 @@ parameter_types! {
 pub struct IsRegistered;
 impl ValidatorRegistration<u64> for IsRegistered {
 	fn is_registered(id: &u64) -> bool {
-		*id != 7u64
+		*id != 42u64
 	}
 }
 
@@ -217,7 +217,7 @@ impl Config for Test {
 pub fn new_test_ext() -> sp_io::TestExternalities {
 	sp_tracing::try_init_simple();
 	let mut t = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
-	let invulnerables = vec![1, 2];
+	let invulnerables = vec![2, 1]; // unsorted
 
 	let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100)];
 	let keys = balances