Implements try_state hook in elections and EPM pallets (#13979)

* Adds try_state hook to elections pallets

* Addresses PR review comments

Co-authored-by: Liam Aharon <liam.aharon@hotmail.com>

* remove unecessary println

* ensures try-runtime does not mutate storage

* Addresses PR comments

* Fixes snapshot invariant checks; simplifies test infra

---------

Co-authored-by: Liam Aharon <liam.aharon@hotmail.com>
Co-authored-by: parity-processbot <>

frame/election-provider-multi-phase/src/lib.rs

@@ -893,6 +893,11 @@ pub mod pallet {
 			// configure this pallet.
 			assert!(T::SignedMaxSubmissions::get() >= T::SignedMaxRefunds::get());
 		}
+
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: T::BlockNumber) -> Result<(), &'static str> {
+			Self::do_try_state()
+		}
 	}
 
 	#[pallet::call]
@@ -1264,6 +1269,8 @@ pub mod pallet {
 	pub type CurrentPhase<T: Config> = StorageValue<_, Phase<T::BlockNumber>, ValueQuery>;
 
 	/// Current best solution, signed or unsigned, queued to be returned upon `elect`.
+	///
+	/// Always sorted by score.
 	#[pallet::storage]
 	#[pallet::getter(fn queued_solution)]
 	pub type QueuedSolution<T: Config> =
@@ -1582,6 +1589,89 @@ impl<T: Config> Pallet<T> {
 	}
 }
 
+#[cfg(feature = "try-runtime")]
+impl<T: Config> Pallet<T> {
+	fn do_try_state() -> Result<(), &'static str> {
+		Self::try_state_snapshot()?;
+		Self::try_state_signed_submissions_map()?;
+		Self::try_state_phase_off()
+	}
+
+	// [`Snapshot`] state check. Invariants:
+	// - [`DesiredTargets`] exists if and only if [`Snapshot`] is present.
+	// - [`SnapshotMetadata`] exist if and only if [`Snapshot`] is present.
+	fn try_state_snapshot() -> Result<(), &'static str> {
+		if <Snapshot<T>>::exists() &&
+			<SnapshotMetadata<T>>::exists() &&
+			<DesiredTargets<T>>::exists()
+		{
+			Ok(())
+		} else if !<Snapshot<T>>::exists() &&
+			!<SnapshotMetadata<T>>::exists() &&
+			!<DesiredTargets<T>>::exists()
+		{
+			Ok(())
+		} else {
+			Err("If snapshot exists, metadata and desired targets should be set too. Otherwise, none should be set.")
+		}
+	}
+
+	// [`SignedSubmissionsMap`] state check. Invariants:
+	// - All [`SignedSubmissionIndices`] are present in [`SignedSubmissionsMap`], and no more;
+	// - [`SignedSubmissionNextIndex`] is not present in [`SignedSubmissionsMap`];
+	// - [`SignedSubmissionIndices`] is sorted by election score.
+	fn try_state_signed_submissions_map() -> Result<(), &'static str> {
+		let mut last_score: ElectionScore = Default::default();
+		let indices = <SignedSubmissionIndices<T>>::get();
+
+		for (i, indice) in indices.iter().enumerate() {
+			let submission = <SignedSubmissionsMap<T>>::get(indice.2);
+			if submission.is_none() {
+				return Err("All signed submissions indices must be part of the submissions map")
+			}
+
+			if i == 0 {
+				last_score = indice.0
+			} else {
+				if last_score.strict_threshold_better(indice.0, Perbill::zero()) {
+					return Err("Signed submission indices vector must be ordered by election score")
+				}
+				last_score = indice.0;
+			}
+		}
+
+		if <SignedSubmissionsMap<T>>::iter().nth(indices.len()).is_some() {
+			return Err("Signed submissions map length should be the same as the indices vec length")
+		}
+
+		match <SignedSubmissionNextIndex<T>>::get() {
+			0 => Ok(()),
+			next =>
+				if <SignedSubmissionsMap<T>>::get(next).is_some() {
+					return Err(
+						"The next submissions index should not be in the submissions maps already",
+					)
+				} else {
+					Ok(())
+				},
+		}
+	}
+
+	// [`Phase::Off`] state check. Invariants:
+	// - If phase is `Phase::Off`, [`Snapshot`] must be none.
+	fn try_state_phase_off() -> Result<(), &'static str> {
+		match Self::current_phase().is_off() {
+			false => Ok(()),
+			true =>
+				if <Snapshot<T>>::get().is_some() {
+					Err("Snapshot must be none when in Phase::Off")
+				} else {
+					Ok(())
+				},
+		}
+	}
+}
+
 impl<T: Config> ElectionProviderBase for Pallet<T> {
 	type AccountId = T::AccountId;
 	type BlockNumber = T::BlockNumber;
@@ -1654,6 +1744,11 @@ mod feasibility_check {
 				MultiPhase::feasibility_check(solution, COMPUTE),
 				FeasibilityError::SnapshotUnavailable
 			);
+
+			// kill also `SnapshotMetadata` and `DesiredTargets` for the storage state to be
+			// consistent for the try_state checks to pass.
+			<SnapshotMetadata<Runtime>>::kill();
+			<DesiredTargets<Runtime>>::kill();
 		})
 	}
 

frame/election-provider-multi-phase/src/mock.rs

@@ -640,7 +640,17 @@ impl ExtBuilder {
 	}
 
 	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
-		self.build().execute_with(test)
+		sp_tracing::try_init_simple();
+
+		let mut ext = self.build();
+		ext.execute_with(test);
+
+		#[cfg(feature = "try-runtime")]
+		ext.execute_with(|| {
+			assert_ok!(<MultiPhase as frame_support::traits::Hooks<u64>>::try_state(
+				System::block_number()
+			));
+		});
 	}
 }
 

frame/election-provider-multi-phase/src/signed.rs

@@ -587,6 +587,11 @@ mod tests {
 				MultiPhase::submit(RuntimeOrigin::signed(10), Box::new(solution)),
 				Error::<Runtime>::PreDispatchEarlySubmission,
 			);
+
+			// make sure invariants hold true and post-test try state checks to pass.
+			<crate::Snapshot<Runtime>>::kill();
+			<crate::SnapshotMetadata<Runtime>>::kill();
+			<crate::DesiredTargets<Runtime>>::kill();
 		})
 	}
 

frame/elections-phragmen/src/lib.rs

@@ -325,6 +325,11 @@ pub mod pallet {
 				T::MaxVotesPerVoter::get(),
 			);
 		}
+
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: T::BlockNumber) -> Result<(), &'static str> {
+			Self::do_try_state()
+		}
 	}
 
 	#[pallet::call]
@@ -893,7 +898,7 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Get the members' account ids.
-	fn members_ids() -> Vec<T::AccountId> {
+	pub(crate) fn members_ids() -> Vec<T::AccountId> {
 		Self::members().into_iter().map(|m| m.who).collect::<Vec<T::AccountId>>()
 	}
 
@@ -1192,6 +1197,104 @@ impl<T: Config> ContainsLengthBound for Pallet<T> {
 	}
 }
 
+#[cfg(any(feature = "try-runtime", test))]
+impl<T: Config> Pallet<T> {
+	fn do_try_state() -> Result<(), &'static str> {
+		Self::try_state_members()?;
+		Self::try_state_runners_up()?;
+		Self::try_state_candidates()?;
+		Self::try_state_candidates_runners_up_disjoint()?;
+		Self::try_state_members_disjoint()?;
+		Self::try_state_members_approval_stake()
+	}
+
+	/// [`Members`] state checks. Invariants:
+	///  - Members are always sorted based on account ID.
+	fn try_state_members() -> Result<(), &'static str> {
+		let mut members = Members::<T>::get().clone();
+		members.sort_by_key(|m| m.who.clone());
+
+		if Members::<T>::get() == members {
+			Ok(())
+		} else {
+			Err("try_state checks: Members must be always sorted by account ID")
+		}
+	}
+
+	// [`RunnersUp`] state checks. Invariants:
+	//  - Elements are sorted based on weight (worst to best).
+	fn try_state_runners_up() -> Result<(), &'static str> {
+		let mut sorted = RunnersUp::<T>::get();
+		// worst stake first
+		sorted.sort_by(|a, b| a.stake.cmp(&b.stake));
+
+		if RunnersUp::<T>::get() == sorted {
+			Ok(())
+		} else {
+			Err("try_state checks: Runners Up must always be sorted by stake (worst to best)")
+		}
+	}
+
+	// [`Candidates`] state checks. Invariants:
+	//  - Always sorted based on account ID.
+	fn try_state_candidates() -> Result<(), &'static str> {
+		let mut candidates = Candidates::<T>::get().clone();
+		candidates.sort_by_key(|(c, _)| c.clone());
+
+		if Candidates::<T>::get() == candidates {
+			Ok(())
+		} else {
+			Err("try_state checks: Candidates must be always sorted by account ID")
+		}
+	}
+	// [`Candidates`] and [`RunnersUp`] state checks. Invariants:
+	//  - Candidates and runners-ups sets are disjoint.
+	fn try_state_candidates_runners_up_disjoint() -> Result<(), &'static str> {
+		match Self::intersects(&Self::candidates_ids(), &Self::runners_up_ids()) {
+			true => Err("Candidates and runners up sets should always be disjoint"),
+			false => Ok(()),
+		}
+	}
+
+	// [`Members`], [`Candidates`] and [`RunnersUp`] state checks. Invariants:
+	//  - Members and candidates sets are disjoint;
+	//  - Members and runners-ups sets are disjoint.
+	fn try_state_members_disjoint() -> Result<(), &'static str> {
+		match Self::intersects(&Pallet::<T>::members_ids(), &Self::candidates_ids()) &&
+			Self::intersects(&Pallet::<T>::members_ids(), &Self::runners_up_ids())
+		{
+			true => Err("Members set should be disjoint from candidates and runners-up sets"),
+			false => Ok(()),
+		}
+	}
+
+	// [`Members`], [`RunnersUp`] and approval stake state checks. Invariants:
+	// - Selected members should have approval stake;
+	// - Selected RunnersUp should have approval stake.
+	fn try_state_members_approval_stake() -> Result<(), &'static str> {
+		match Members::<T>::get()
+			.iter()
+			.chain(RunnersUp::<T>::get().iter())
+			.all(|s| s.stake != BalanceOf::<T>::zero())
+		{
+			true => Ok(()),
+			false => Err("Members and RunnersUp must have approval stake"),
+		}
+	}
+
+	fn intersects<P: PartialEq>(a: &[P], b: &[P]) -> bool {
+		a.iter().any(|e| b.contains(e))
+	}
+
+	fn candidates_ids() -> Vec<T::AccountId> {
+		Pallet::<T>::candidates().iter().map(|(x, _)| x).cloned().collect::<Vec<_>>()
+	}
+
+	fn runners_up_ids() -> Vec<T::AccountId> {
+		Pallet::<T>::runners_up().into_iter().map(|r| r.who).collect::<Vec<_>>()
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -1418,7 +1521,13 @@ mod tests {
 			.into();
 			ext.execute_with(pre_conditions);
 			ext.execute_with(test);
-			ext.execute_with(post_conditions)
+
+			#[cfg(feature = "try-runtime")]
+			ext.execute_with(|| {
+				assert_ok!(<Elections as frame_support::traits::Hooks<u64>>::try_state(
+					System::block_number()
+				));
+			});
 		}
 	}
 
@@ -1475,54 +1584,13 @@ mod tests {
 			.unwrap_or_default()
 	}
 
-	fn intersects<T: PartialEq>(a: &[T], b: &[T]) -> bool {
-		a.iter().any(|e| b.contains(e))
-	}
-
-	fn ensure_members_sorted() {
-		let mut members = Elections::members().clone();
-		members.sort_by_key(|m| m.who);
-		assert_eq!(Elections::members(), members);
-	}
-
-	fn ensure_candidates_sorted() {
-		let mut candidates = Elections::candidates().clone();
-		candidates.sort_by_key(|(c, _)| *c);
-		assert_eq!(Elections::candidates(), candidates);
-	}
-
 	fn locked_stake_of(who: &u64) -> u64 {
 		Voting::<Test>::get(who).stake
 	}
 
-	fn ensure_members_has_approval_stake() {
-		// we filter members that have no approval state. This means that even we have more seats
-		// than candidates, we will never ever chose a member with no votes.
-		assert!(Elections::members()
-			.iter()
-			.chain(Elections::runners_up().iter())
-			.all(|s| s.stake != u64::zero()));
-	}
-
-	fn ensure_member_candidates_runners_up_disjoint() {
-		// members, candidates and runners-up must always be disjoint sets.
-		assert!(!intersects(&members_ids(), &candidate_ids()));
-		assert!(!intersects(&members_ids(), &runners_up_ids()));
-		assert!(!intersects(&candidate_ids(), &runners_up_ids()));
-	}
-
 	fn pre_conditions() {
 		System::set_block_number(1);
-		ensure_members_sorted();
-		ensure_candidates_sorted();
-		ensure_member_candidates_runners_up_disjoint();
-	}
-
-	fn post_conditions() {
-		ensure_members_sorted();
-		ensure_candidates_sorted();
-		ensure_member_candidates_runners_up_disjoint();
-		ensure_members_has_approval_stake();
+		Elections::do_try_state().unwrap();
 	}
 
 	fn submit_candidacy(origin: RuntimeOrigin) -> sp_runtime::DispatchResult {