From f84b89710b0e6da655a1b40a491a8184780489e7 Mon Sep 17 00:00:00 2001
From: Ankan <10196091+Ank4n@users.noreply.github.com>
Date: Sat, 4 Nov 2023 16:41:51 +0100
Subject: [PATCH] Speed up nominator state checks in staking pallet (#2153)

Should help https://github.com/paritytech/polkadot-sdk/issues/234.
Related to https://github.com/paritytech/polkadot-sdk/issues/2020 and
https://github.com/paritytech/polkadot-sdk/issues/2108.

Refactors and improves running time for try runtime checks for staking
pallet.

Tested on westend on my M2 pro: running time drops from 90 seconds to 7
seconds.
---
 substrate/frame/staking/src/pallet/impls.rs | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index bb16fe56d51e..9c36c94b87b4 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -1864,6 +1864,13 @@ impl<T: Config> Pallet<T> {
 		// a check per nominator to ensure their entire stake is correctly distributed. Will only
 		// kick-in if the nomination was submitted before the current era.
 		let era = Self::active_era().unwrap().index;
+
+		// cache era exposures to avoid too many db reads.
+		let era_exposures = T::SessionInterface::validators()
+			.iter()
+			.map(|v| Self::eras_stakers(era, v))
+			.collect::<Vec<_>>();
+
 		<Nominators<T>>::iter()
 			.filter_map(
 				|(nominator, nomination)| {
@@ -1878,9 +1885,8 @@ impl<T: Config> Pallet<T> {
 				// must be bonded.
 				Self::ensure_is_stash(&nominator)?;
 				let mut sum = BalanceOf::<T>::zero();
-				T::SessionInterface::validators()
+				era_exposures
 					.iter()
-					.map(|v| Self::eras_stakers(era, v))
 					.map(|e| -> Result<(), TryRuntimeError> {
 						let individual =
 							e.others.iter().filter(|e| e.who == nominator).collect::<Vec<_>>();
@@ -1896,6 +1902,14 @@ impl<T: Config> Pallet<T> {
 						Ok(())
 					})
 					.collect::<Result<Vec<_>, _>>()?;
+
+				// We take total instead of active as the nominator might have requested to unbond
+				// some of their stake that is still exposed in the current era.
+				if sum <= Self::ledger(Stash(nominator.clone()))?.total {
+					// This can happen when there is a slash in the current era so we only warn.
+					log!(warn, "nominator stake exceeds what is bonded.");
+				}
+
 				Ok(())
 			})
 			.collect::<Result<Vec<_>, _>>()?;