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 <>

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

@@ -881,6 +881,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]
@@ -1252,6 +1257,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> =
@@ -1570,6 +1577,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;
@@ -1642,6 +1732,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();
 		})
 	}
 

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

@@ -615,7 +615,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()
+			));
+		});
 	}
 }
 

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

@@ -554,6 +554,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();
 		})
 	}