Add sorted voter list preserve mechanism to multi-block staking pallet

prdoc/pr_6520.prdoc

@@ -0,0 +1,10 @@
+title: Implements a mechanism to perserve the bags-list ID order
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      Refactors the bags list pallet to support a mode which perserves the order of the elements in the list.
+
+crates:
+ - name: pallet-bags-list
+   bump: major

substrate/frame/bags-list/src/lib.rs

@@ -68,6 +68,15 @@
 //! - iteration over the top* N items by score, where the precise ordering of items doesn't
 //!   particularly matter.
 //!
+//! ### Preserve order mode
+//!
+//! The bags list pallet supports a mode which preserves the order of the elements in the list.
+//! While the implementor of [`Config::PreserveOrder`] returns `true`, no rebags or moves within a
+//! bag are allowed. Which in practice means that:
+//! - calling [`Pallet::rebag`], [`Pallet::put_in_front_of`] and [`Pallet::put_in_front_of_other`]
+//!   will return a [`Error::MustPreserveOrder`];
+//! - calling [`SortedListProvider::on_update`] is a noop.
+//!
 //! ### Further Details
 //!
 //! - items are kept in bags, which are delineated by their range of score (See
@@ -127,8 +136,12 @@ extern crate alloc;
 use alloc::boxed::Box;
 use codec::FullCodec;
 use frame_election_provider_support::{ScoreProvider, SortedListProvider};
+use frame_support::{ensure, traits::Get};
 use frame_system::ensure_signed;
-use sp_runtime::traits::{AtLeast32BitUnsigned, Bounded, StaticLookup};
+use sp_runtime::{
+	traits::{AtLeast32BitUnsigned, Bounded, StaticLookup},
+	DispatchError,
+};
 
 #[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
 use sp_runtime::TryRuntimeError;
@@ -247,6 +260,9 @@ pub mod pallet {
 			+ TypeInfo
 			+ FullCodec
 			+ MaxEncodedLen;
+
+		/// Something that signals whether the order of the element in the bags list may change.
+		type PreserveOrder: Get<bool>;
 	}
 
 	/// A single node, within some bag.
@@ -277,6 +293,8 @@ pub mod pallet {
 	pub enum Error<T, I = ()> {
 		/// A error in the list interface implementation.
 		List(ListError),
+		/// A request that does not preserve the list's order was requested out of time.
+		MustPreserveOrder,
 	}
 
 	impl<T, I> From<ListError> for Error<T, I> {
@@ -301,6 +319,8 @@ pub mod pallet {
 		#[pallet::weight(T::WeightInfo::rebag_non_terminal().max(T::WeightInfo::rebag_terminal()))]
 		pub fn rebag(origin: OriginFor<T>, dislocated: AccountIdLookupOf<T>) -> DispatchResult {
 			ensure_signed(origin)?;
+			ensure!(!T::PreserveOrder::get(), Error::<T, I>::MustPreserveOrder);
+
 			let dislocated = T::Lookup::lookup(dislocated)?;
 			let current_score = T::ScoreProvider::score(&dislocated);
 			let _ = Pallet::<T, I>::do_rebag(&dislocated, current_score)
@@ -324,6 +344,8 @@ pub mod pallet {
 			origin: OriginFor<T>,
 			lighter: AccountIdLookupOf<T>,
 		) -> DispatchResult {
+			ensure!(!T::PreserveOrder::get(), Error::<T, I>::MustPreserveOrder);
+
 			let heavier = ensure_signed(origin)?;
 			let lighter = T::Lookup::lookup(lighter)?;
 			List::<T, I>::put_in_front_of(&lighter, &heavier)
@@ -341,6 +363,8 @@ pub mod pallet {
 			heavier: AccountIdLookupOf<T>,
 			lighter: AccountIdLookupOf<T>,
 		) -> DispatchResult {
+			ensure!(!T::PreserveOrder::get(), Error::<T, I>::MustPreserveOrder);
+
 			let _ = ensure_signed(origin)?;
 			let lighter = T::Lookup::lookup(lighter)?;
 			let heavier = T::Lookup::lookup(heavier)?;
@@ -431,7 +455,13 @@ impl<T: Config<I>, I: 'static> SortedListProvider<T::AccountId> for Pallet<T, I>
 	}
 
 	fn on_update(id: &T::AccountId, new_score: T::Score) -> Result<(), ListError> {
-		Pallet::<T, I>::do_rebag(id, new_score).map(|_| ())
+		if T::PreserveOrder::get() {
+			// lock is set, on_update is a noop.
+			ensure!(list::Node::<T, I>::get(&id).is_some(), ListError::NodeNotFound);
+			Ok(())
+		} else {
+			Pallet::<T, I>::do_rebag(id, new_score).map(|_| ())
+		}
 	}
 
 	fn on_remove(id: &T::AccountId) -> Result<(), ListError> {

substrate/frame/bags-list/src/list/mod.rs

@@ -840,7 +840,7 @@ impl<T: Config<I>, I: 'static> Node<T, I> {
 		crate::ListNodes::<T, I>::insert(self.id.clone(), self);
 	}
 
-	/// Update neighboring nodes to point to reach other.
+	/// Update neighboring nodes to point to each other.
 	///
 	/// Only updates storage for adjacent nodes, but not `self`; so the user may need to call
 	/// `self.put`.

substrate/frame/bags-list/src/mock.rs

@@ -56,6 +56,7 @@ impl frame_system::Config for Runtime {
 
 parameter_types! {
 	pub static BagThresholds: &'static [VoteWeight] = &[10, 20, 30, 40, 50, 60, 1_000, 2_000, 10_000];
+	pub static PreserveOrder: bool = false;
 }
 
 impl bags_list::Config for Runtime {
@@ -64,6 +65,7 @@ impl bags_list::Config for Runtime {
 	type BagThresholds = BagThresholds;
 	type ScoreProvider = StakingMock;
 	type Score = VoteWeight;
+	type PreserveOrder = PreserveOrder;
 }
 
 type Block = frame_system::mocking::MockBlock<Runtime>;

substrate/frame/bags-list/src/tests.rs

@@ -735,4 +735,67 @@ mod sorted_list_provider {
 			assert!(non_existent_ids.iter().all(|id| !BagsList::contains(id)));
 		})
 	}
+
+	#[test]
+	fn on_update_with_preserve_order_works() {
+		ExtBuilder::default().build_and_execute(|| {
+			// initial iter state.
+			let initial_iter = ListNodes::<Runtime>::iter().map(|(a, _)| a).collect::<Vec<_>>();
+			assert_eq!(initial_iter, vec![3, 4, 1, 2]);
+
+			assert_eq!(ListBags::<Runtime>::iter().count(), 2);
+			assert_eq!(List::<Runtime>::get_bags(), vec![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			assert_eq!(BagsList::get_score(&1), Ok(10));
+
+			// sets preserve order.
+			PreserveOrder::set(true);
+
+			// update 1's score in score provider and call on_update.
+			StakingMock::set_score_of(&1, 2_000);
+			assert_ok!(BagsList::on_update(&1, 2_000));
+
+			// preserve order is set, so the node's bag score and iter state remain the same,
+			// despite the updated score in the score provider for 1.
+			assert_eq!(BagsList::get_score(&1), Ok(10));
+			assert_eq!(
+				initial_iter,
+				ListNodes::<Runtime>::iter().map(|(a, _)| a).collect::<Vec<_>>()
+			);
+			assert_eq!(List::<Runtime>::get_bags(), vec![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			// unsets preserve order.
+			PreserveOrder::set(false);
+			assert_ok!(BagsList::rebag(RuntimeOrigin::signed(0), 1));
+
+			// now the rebag worked as expected.
+			assert_eq!(BagsList::get_score(&1), Ok(2_000));
+			assert_eq!(List::<Runtime>::get_bags(), vec![(1_000, vec![2, 3, 4]), (2_000, vec![1])]);
+		})
+	}
+
+	#[test]
+	fn calls_with_preserve_order_fails() {
+		use crate::pallet::Error;
+
+		ExtBuilder::default().build_and_execute(|| {
+			// sets preserve order.
+			PreserveOrder::set(true);
+
+			assert_noop!(
+				BagsList::rebag(RuntimeOrigin::signed(0), 1),
+				Error::<Runtime>::MustPreserveOrder
+			);
+
+			assert_noop!(
+				BagsList::put_in_front_of(RuntimeOrigin::signed(0), 1),
+				Error::<Runtime>::MustPreserveOrder
+			);
+
+			assert_noop!(
+				BagsList::put_in_front_of_other(RuntimeOrigin::signed(0), 1, 2),
+				Error::<Runtime>::MustPreserveOrder
+			);
+		})
+	}
 }

substrate/frame/nomination-pools/test-delegate-stake/src/mock.rs

@@ -117,6 +117,7 @@ impl pallet_bags_list::Config<VoterBagsListInstance> for Runtime {
 	type BagThresholds = BagThresholds;
 	type ScoreProvider = Staking;
 	type Score = VoteWeight;
+	type PreserveOrder = Staking;
 }
 
 pub struct BalanceToU256;

substrate/frame/nomination-pools/test-transfer-stake/src/mock.rs

@@ -109,6 +109,7 @@ impl pallet_bags_list::Config<VoterBagsListInstance> for Runtime {
 	type BagThresholds = BagThresholds;
 	type ScoreProvider = Staking;
 	type Score = VoteWeight;
+	type PreserveOrder = Staking;
 }
 
 pub struct BalanceToU256;

substrate/frame/staking/src/mock.rs

@@ -220,6 +220,7 @@ impl pallet_bags_list::Config<VoterBagsListInstance> for Test {
 	type ScoreProvider = Staking;
 	type BagThresholds = BagThresholds;
 	type Score = VoteWeight;
+	type PreserveOrder = Staking;
 }
 
 // multi-page types and controller.

substrate/frame/staking/src/pallet/impls.rs

@@ -1033,9 +1033,15 @@ impl<T: Config> Pallet<T> {
 		let mut sorted_voters = match VoterSnapshotStatus::<T>::get() {
 			// start the snapshot procssing from the beginning.
 			SnapshotStatus::Waiting => T::VoterList::iter(),
-			// snapshot continues, start from the last iterated voter in the list.
-			SnapshotStatus::Ongoing(account_id) => T::VoterList::iter_from(&account_id)
-				.defensive_unwrap_or(Box::new(vec![].into_iter())),
+			// snapshot continues.
+			SnapshotStatus::Ongoing(account_id) => {
+				let iter = T::VoterList::iter_from(&account_id)
+					.defensive_unwrap_or(Box::new(vec![].into_iter()));
+
+				// add cursor account to the beginning of the iter, as it hasn't been processed
+				// yet.
+				Box::new(vec![account_id].into_iter().chain(iter))
+			},
 			// all voters have been consumed already, return an empty iterator.
 			SnapshotStatus::Consumed => Box::new(vec![].into_iter()),
 		};
@@ -1112,24 +1118,21 @@ impl<T: Config> Pallet<T> {
 			}
 		}
 
-		// update the voter snapshot status.
+		// update cursor if there are remaining accounts in the voter list. The cursor must be the
+		// next *non-processed* voter if it exists.
 		VoterSnapshotStatus::<T>::mutate(|status| {
 			match (page, status.clone()) {
 				// last page, reset status for next round.
 				(0, _) => *status = SnapshotStatus::Waiting,
 
+				// progress, update cursor.
 				(_, SnapshotStatus::Waiting) | (_, SnapshotStatus::Ongoing(_)) => {
-					let maybe_last = all_voters.last().map(|(x, _, _)| x).cloned();
-
-					if let Some(ref last) = maybe_last {
-						if maybe_last == T::VoterList::iter().last() {
-							// all voters in the voter list have been consumed.
-							*status = SnapshotStatus::Consumed;
-						} else {
-							*status = SnapshotStatus::Ongoing(last.clone());
-						}
+					if let Some(cursor) = sorted_voters.next() {
+						// cursor is the next non-processed element in the iter.
+						*status = SnapshotStatus::Ongoing(cursor);
 					} else {
-						debug_assert!(*status == SnapshotStatus::Consumed);
+						// all voters in the list consumed.
+						*status = SnapshotStatus::Consumed;
 					}
 				},
 				// do nothing.
@@ -1236,6 +1239,8 @@ impl<T: Config> Pallet<T> {
 	/// `Nominators` or `VoterList` outside of this function is almost certainly
 	/// wrong.
 	pub fn do_remove_nominator(who: &T::AccountId) -> bool {
+		Self::ensure_voter_cursor(who);
+
 		let outcome = if Nominators::<T>::contains_key(who) {
 			Nominators::<T>::remove(who);
 			let _ = T::VoterList::on_remove(who).defensive();
@@ -1248,6 +1253,7 @@ impl<T: Config> Pallet<T> {
 			Nominators::<T>::count() + Validators::<T>::count(),
 			T::VoterList::count()
 		);
+		debug_assert!(VoterSnapshotStatus::<T>::get() != SnapshotStatus::Ongoing(who.clone()));
 
 		outcome
 	}
@@ -1281,8 +1287,11 @@ impl<T: Config> Pallet<T> {
 	/// `Validators` or `VoterList` outside of this function is almost certainly
 	/// wrong.
 	pub fn do_remove_validator(who: &T::AccountId) -> bool {
+		Self::ensure_voter_cursor(who);
+
 		let outcome = if Validators::<T>::contains_key(who) {
 			Validators::<T>::remove(who);
+
 			let _ = T::VoterList::on_remove(who).defensive();
 			true
 		} else {
@@ -1293,6 +1302,7 @@ impl<T: Config> Pallet<T> {
 			Nominators::<T>::count() + Validators::<T>::count(),
 			T::VoterList::count()
 		);
+		debug_assert!(VoterSnapshotStatus::<T>::get() != SnapshotStatus::Ongoing(who.clone()));
 
 		outcome
 	}
@@ -1727,6 +1737,20 @@ where
 	}
 }
 
+/// Implements the lock to be used by the the sorted list providers implemented by the bags list
+/// pallet.
+///
+/// If the voter snapshot creation is in progress return `true`, in which case the bags list ID
+/// re-ordering should be disabled.
+impl<T: Config> Get<bool> for Pallet<T> {
+	fn get() -> bool {
+		match VoterSnapshotStatus::<T>::get() {
+			SnapshotStatus::Ongoing(_) => true,
+			_ => false,
+		}
+	}
+}
+
 impl<T: Config> ScoreProvider<T::AccountId> for Pallet<T> {
 	type Score = VoteWeight;
 
@@ -1808,6 +1832,7 @@ impl<T: Config> SortedListProvider<T::AccountId> for UseValidatorsMap<T> {
 		// nothing to do upon regenerate.
 		0
 	}
+
 	#[cfg(feature = "try-runtime")]
 	fn try_state() -> Result<(), TryRuntimeError> {
 		Ok(())

substrate/frame/staking/src/pallet/mod.rs

@@ -917,11 +917,17 @@ pub mod pallet {
 			mode: Forcing,
 		},
 		/// Report of a controller batch deprecation.
-		ControllerBatchDeprecated { failures: u32 },
+		ControllerBatchDeprecated {
+			failures: u32,
+		},
 		/// Validator has been disabled.
-		ValidatorDisabled { stash: T::AccountId },
+		ValidatorDisabled {
+			stash: T::AccountId,
+		},
 		/// Validator has been re-enabled.
-		ValidatorReenabled { stash: T::AccountId },
+		ValidatorReenabled {
+			stash: T::AccountId,
+		},
 	}
 
 	#[pallet::error]
@@ -1233,6 +1239,21 @@ pub mod pallet {
 		pub fn current_planned_session() -> SessionIndex {
 			CurrentPlannedSession::<T>::get()
 		}
+
+		/// Ensures the snapshot voter cursor is *not* the same as `who`. Updates to next in the
+		/// voter list interator if required.
+		pub(crate) fn ensure_voter_cursor(who: &T::AccountId) {
+			if VoterSnapshotStatus::<T>::get() == SnapshotStatus::Ongoing(who.clone()) {
+				let next =
+					T::VoterList::iter_from(who).unwrap_or(Box::new(vec![].into_iter())).next();
+
+				match next {
+					Some(new_cursor) =>
+						VoterSnapshotStatus::<T>::set(SnapshotStatus::Ongoing(new_cursor)),
+					None => VoterSnapshotStatus::<T>::set(SnapshotStatus::Consumed),
+				}
+			}
+		}
 	}
 
 	#[pallet::call]