[AHM/Staking] Allow own stake to be collected in all pages (#10051)

TLDR: Thanks to a number of reports from Kusama validator, we have been
investigating an issue where the self stake of some validators is not
present in `ErasStakersPaged` and `ErasStakersOverview`, and
consequently not earning staking rewards.

After investigation, it was clear that the `polkadot-staking-miner` is
doing the right thing, and the self-stake is indeed submitted to the
chain as a part of the staking election result.

The root cause in how `pallet-staking-async` ingests the staking
election result. This code **made a (wrong) assumption that self-stake
is only present in the first page of a multi-page election. This PR
fixes this issue.**

---

- [x] (nice to have) revisit try-state checks of the pallet,
`ErasStakersPaged`, `ErasStakersOverview` and `ErasTotalStake` are
triple checked for correctness
- [ ] (nice to have) fix related displays in PJS:
polkadot-js/apps#11930
- [ ] backport to 2507 and 2509
- Westend 
- [ ] Find a similar issue in westend for monitoring
(paritytech/devops#4402)
  - [ ] Upgrade westend to a branch/release containing this fix 
- [ ] Upgrade westend with this fix, observe test examples are resolved
- [ ] Bump backported version in `fellowship/runtimes`
- [ ] Release 1.9.3 for Kusama with this fix
- [ ] Ensure known test cases are fixed
- One example: `F2WyUUFXLYnBg6acv7t2KFzH6D7CyNcvC4mRCwUdsHTUB4t` (in
election round 46, likely repeating in subsequent ones too)
- This PR will be part of Polkadot release 2.0.0, the Polkadot AHM.
- [ ] Forum post with further explanation once PR is merged.

---------

Co-authored-by: cmd[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: kianenigma

prdoc/pr_10051.prdoc

@@ -0,0 +1,8 @@
+title: '[AHM/Staking] Allow own stake to be collected in all pages'
+doc:
+- audience: Runtime Dev
+  description: |-
+    TLDR: Thanks to a number of reports from Kusama validator, we have been investigating an issue where the self stake of some validators is not present in `ErasStakersPaged` and `ErasStakersOverview`. After investigation, it was clear that the `polkadot-staking-miner` is doing the right thing, and the self-stake is indeed submitted to the chain as a part of the election result. The root cause in how `pallet-staking-async` ingests the paginated election result. This code **made a (wrong) assumption that self-stake is only present in the first page of a multi-page election. This PR fixes this issue.**
+crates:
+- name: pallet-staking-async
+  bump: patch

substrate/frame/staking-async/src/benchmarking.rs

@@ -29,7 +29,10 @@ pub use frame_benchmarking::{
 };
 use frame_election_provider_support::SortedListProvider;
 use frame_support::{
-	assert_ok, pallet_prelude::*, storage::bounded_vec::BoundedVec, traits::TryCollect,
+	assert_ok,
+	pallet_prelude::*,
+	storage::bounded_vec::BoundedVec,
+	traits::{fungible::Inspect, TryCollect},
 };
 use frame_system::RawOrigin;
 use pallet_staking_async_rc_client as rc_client;
@@ -1226,14 +1229,15 @@ mod benchmarks {
 			.map(|validator_index| account::<T::AccountId>("validator", validator_index, SEED))
 			.for_each(|validator| {
 				let exposure = sp_staking::Exposure::<T::AccountId, BalanceOf<T>> {
-					own: BalanceOf::<T>::max_value(),
-					total: BalanceOf::<T>::max_value(),
+					own: T::Currency::minimum_balance(),
+					total: T::Currency::minimum_balance() *
+						(exposed_nominators_per_validator + 1).into(),
 					others: (0..exposed_nominators_per_validator)
 						.map(|n| {
 							let nominator = account::<T::AccountId>("nominator", n, SEED);
 							IndividualExposure {
 								who: nominator,
-								value: BalanceOf::<T>::max_value(),
+								value: T::Currency::minimum_balance(),
 							}
 						})
 						.collect::<Vec<_>>(),

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

@@ -1949,78 +1949,62 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Invariants:
-	/// * ActiveEra is Some.
-	/// * For each paged era exposed validator, check if the exposure total is sane (exposure.total
-	/// = exposure.own + exposure.own).
-	/// * Paged exposures metadata (`ErasStakersOverview`) matches the paged exposures state.
+	/// Nothing to do if ActiveEra is not set.
+	/// For each page in `ErasStakersPaged`, `page_total` must be set.
+	/// For each metadata:
+	/// 	* page_count is correct
+	/// 	* nominator_count is correct
+	/// 	* total is own + sum of pages
+	/// `ErasTotalStake`` must be correct
 	fn check_paged_exposures() -> Result<(), TryRuntimeError> {
-		use alloc::collections::btree_map::BTreeMap;
-		use sp_staking::PagedExposureMetadata;
-
-		// Sanity check for the paged exposure of the active era.
-		let mut exposures: BTreeMap<T::AccountId, PagedExposureMetadata<BalanceOf<T>>> =
-			BTreeMap::new();
-		// If the pallet is not initialized, we return immediately from pallet's do_try_state() and
-		// we don't call this method. Otherwise, Eras::do_try_state enforces that both ActiveEra
-		// and CurrentEra are Some. Thus, we should never hit this error.
-		let era = ActiveEra::<T>::get()
-			.ok_or(TryRuntimeError::Other("ActiveEra must be set when checking paged exposures"))?
-			.index;
-
-		let accumulator_default = PagedExposureMetadata {
-			total: Zero::zero(),
-			own: Zero::zero(),
-			nominator_count: 0,
-			page_count: 0,
-		};
-
-		ErasStakersPaged::<T>::iter_prefix((era,))
-			.map(|((validator, _page), expo)| {
-				ensure!(
-					expo.page_total ==
-						expo.others.iter().map(|e| e.value).fold(Zero::zero(), |acc, x| acc + x),
-					"wrong total exposure for the page.",
-				);
-
-				let metadata = exposures.get(&validator).unwrap_or(&accumulator_default);
-				exposures.insert(
-					validator,
-					PagedExposureMetadata {
-						total: metadata.total + expo.page_total,
-						own: metadata.own,
-						nominator_count: metadata.nominator_count + expo.others.len() as u32,
-						page_count: metadata.page_count + 1,
-					},
-				);
-
-				Ok(())
+		let Some(era) = ActiveEra::<T>::get().map(|a| a.index) else { return Ok(()) };
+		let overview_and_pages = ErasStakersOverview::<T>::iter_prefix(era)
+			.map(|(validator, metadata)| {
+				let pages = ErasStakersPaged::<T>::iter_prefix((era, validator))
+					.map(|(_idx, page)| page)
+					.collect::<Vec<_>>();
+				(metadata, pages)
 			})
-			.collect::<Result<(), TryRuntimeError>>()?;
+			.collect::<Vec<_>>();
 
-		exposures
-			.iter()
-			.map(|(validator, metadata)| {
-				let actual_overview = ErasStakersOverview::<T>::get(era, validator);
+		ensure!(
+			overview_and_pages.iter().flat_map(|(_m, pages)| pages).all(|page| {
+				let expected = page
+					.others
+					.iter()
+					.map(|e| e.value)
+					.fold(BalanceOf::<T>::zero(), |acc, x| acc + x);
+				page.page_total == expected
+			}),
+			"found wrong page_total"
+		);
 
-				ensure!(actual_overview.is_some(), "No overview found for a paged exposure");
-				let actual_overview = actual_overview.unwrap();
+		ensure!(
+			overview_and_pages.iter().all(|(metadata, pages)| {
+				let page_count_good = metadata.page_count == pages.len() as u32;
+				let nominator_count_good = metadata.nominator_count ==
+					pages.iter().map(|p| p.others.len() as u32).fold(0u32, |acc, x| acc + x);
+				let total_good = metadata.total ==
+					metadata.own +
+						pages
+							.iter()
+							.fold(BalanceOf::<T>::zero(), |acc, page| acc + page.page_total);
+
+				page_count_good && nominator_count_good && total_good
+			}),
+			"found bad metadata"
+		);
 
-				ensure!(
-					actual_overview.total == metadata.total + actual_overview.own,
-					"Exposure metadata does not have correct total exposed stake."
-				);
-				ensure!(
-					actual_overview.nominator_count == metadata.nominator_count,
-					"Exposure metadata does not have correct count of nominators."
-				);
-				ensure!(
-					actual_overview.page_count == metadata.page_count,
-					"Exposure metadata does not have correct count of pages."
-				);
+		ensure!(
+			overview_and_pages
+				.iter()
+				.map(|(metadata, _pages)| metadata.total)
+				.fold(BalanceOf::<T>::zero(), |acc, x| acc + x) ==
+				ErasTotalStake::<T>::get(era),
+			"found bad eras total stake"
+		);
 
-				Ok(())
-			})
-			.collect::<Result<(), TryRuntimeError>>()
+		Ok(())
 	}
 
 	/// Ensures offence pipeline and slashing is in a healthy state.

substrate/frame/staking-async/src/session_rotation.rs

@@ -127,16 +127,19 @@ impl<T: Config> Eras<T> {
 
 	/// Get exposure for a validator at a given era and page.
 	///
+	/// This is mainly used for rewards and slashing. Validator's self-stake is only returned in
+	/// page 0.
+	///
 	/// This builds a paged exposure from `PagedExposureMetadata` and `ExposurePage` of the
-	/// validator. For older non-paged exposure, it returns the clipped exposure directly.
+	/// validator.
 	pub(crate) fn get_paged_exposure(
 		era: EraIndex,
 		validator: &T::AccountId,
 		page: Page,
 	) -> Option<PagedExposure<T::AccountId, BalanceOf<T>>> {
 		let overview = <ErasStakersOverview<T>>::get(&era, validator)?;
 
-		// validator stake is added only in page zero
+		// validator stake is added only in page zero.
 		let validator_stake = if page == 0 { overview.own } else { Zero::zero() };
 
 		// since overview is present, paged exposure will always be present except when a
@@ -230,57 +233,96 @@ impl<T: Config> Eras<T> {
 		mut exposure: Exposure<T::AccountId, BalanceOf<T>>,
 	) {
 		let page_size = T::MaxExposurePageSize::get().defensive_max(1);
+		if cfg!(debug_assertions) {
+			// sanitize the exposure in case some test data is wrong.
+			let expected_total = exposure
+				.others
+				.iter()
+				.map(|ie| ie.value)
+				.fold::<BalanceOf<T>, _>(Default::default(), |acc, x| acc + x)
+				.saturating_add(exposure.own);
+			debug_assert_eq!(expected_total, exposure.total, "exposure total must equal own + sum(others) for (era: {:?}, validator: {:?}, exposure: {:?})", era, validator, exposure);
+		}
 
-		if let Some(stored_overview) = ErasStakersOverview::<T>::get(era, &validator) {
-			let last_page_idx = stored_overview.page_count.saturating_sub(1);
-
+		if let Some(overview) = ErasStakersOverview::<T>::get(era, &validator) {
+			// collect some info from the un-touched overview for later use.
+			let last_page_idx = overview.page_count.saturating_sub(1);
 			let mut last_page =
 				ErasStakersPaged::<T>::get((era, validator, last_page_idx)).unwrap_or_default();
 			let last_page_empty_slots =
 				T::MaxExposurePageSize::get().saturating_sub(last_page.others.len() as u32);
 
-			// splits the exposure so that `exposures_append` will fit within the last exposure
-			// page, up to the max exposure page size. The remaining individual exposures in
-			// `exposure` will be added to new pages.
-			let exposures_append = exposure.split_others(last_page_empty_slots);
-
-			ErasStakersOverview::<T>::mutate(era, &validator, |stored| {
-				// new metadata is updated based on 3 different set of exposures: the
-				// current one, the exposure split to be "fitted" into the current last page and
-				// the exposure set that will be appended from the new page onwards.
-				let new_metadata =
-					stored.defensive_unwrap_or_default().update_with::<T::MaxExposurePageSize>(
-						[&exposures_append, &exposure]
-							.iter()
-							.fold(Default::default(), |total, expo| {
-								total.saturating_add(expo.total.saturating_sub(expo.own))
-							}),
-						[&exposures_append, &exposure]
-							.iter()
-							.fold(Default::default(), |count, expo| {
-								count.saturating_add(expo.others.len() as u32)
-							}),
+			// update nominator-count, page-count, and total stake in overview (done in
+			// `update_with`).
+			let new_stake_added = exposure.total;
+			let new_nominators_added = exposure.others.len() as u32;
+			let mut updated_overview = overview
+				.update_with::<T::MaxExposurePageSize>(new_stake_added, new_nominators_added);
+
+			// update own stake, if applicable.
+			match (updated_overview.own.is_zero(), exposure.own.is_zero()) {
+				(true, false) => {
+					// first time we see own exposure -- good.
+					// note: `total` is already updated above.
+					updated_overview.own = exposure.own;
+				},
+				(true, true) | (false, true) => {
+					// no new own exposure is added, nothing to do
+				},
+				(false, false) => {
+					debug_assert!(
+						false,
+						"validator own stake already set in overview for (era: {:?}, validator: {:?}, current overview: {:?}, new exposure: {:?})",
+						era,
+						validator,
+						updated_overview,
+						exposure,
 					);
-				*stored = new_metadata.into();
-			});
+					defensive!("duplicate validator self stake in election");
+				},
+			};
+
+			ErasStakersOverview::<T>::insert(era, &validator, updated_overview);
+			// we are done updating the overview now, `updated_overview` should not be used anymore.
+			// We've updated:
+			// * nominator count
+			// * total stake
+			// * own stake (if applicable)
+			// * page count
+			//
+			// next step:
+			// * new-keys or updates in `ErasStakersPaged`
+			//
+			// we don't need the information about own stake anymore -- drop it.
+			exposure.total = exposure.total.saturating_sub(exposure.own);
+			exposure.own = Zero::zero();
+
+			// splits the exposure so that `append_to_last_page` will fit within the last exposure
+			// page, up to the max exposure page size. The remaining individual exposures in
+			// `put_in_new_pages` will be added to new pages.
+			let append_to_last_page = exposure.split_others(last_page_empty_slots);
+			let put_in_new_pages = exposure;
+
+			// handle last page first.
 
 			// fill up last page with exposures.
-			last_page.page_total = last_page
-				.page_total
-				.saturating_add(exposures_append.total)
-				.saturating_sub(exposures_append.own);
-			last_page.others.extend(exposures_append.others);
+			last_page.page_total = last_page.page_total.saturating_add(append_to_last_page.total);
+			last_page.others.extend(append_to_last_page.others);
 			ErasStakersPaged::<T>::insert((era, &validator, last_page_idx), last_page);
 
 			// now handle the remaining exposures and append the exposure pages. The metadata update
 			// has been already handled above.
-			let (_, exposure_pages) = exposure.into_pages(page_size);
-
-			exposure_pages.into_iter().enumerate().for_each(|(idx, paged_exposure)| {
-				let append_at =
-					(last_page_idx.saturating_add(1).saturating_add(idx as u32)) as Page;
-				<ErasStakersPaged<T>>::insert((era, &validator, append_at), paged_exposure);
-			});
+			let (_unused_metadata, put_in_new_pages_chunks) =
+				put_in_new_pages.into_pages(page_size);
+
+			put_in_new_pages_chunks
+				.into_iter()
+				.enumerate()
+				.for_each(|(idx, paged_exposure)| {
+					let append_at =
+						(last_page_idx.saturating_add(1).saturating_add(idx as u32)) as Page;
+					<ErasStakersPaged<T>>::insert((era, &validator, append_at), paged_exposure);
+				});
 		} else {
 			// expected page count is the number of nominators divided by the page size, rounded up.
 			let expected_page_count = exposure
@@ -1003,16 +1045,17 @@ impl<T: Config> EraElectionPlanner<T> {
 		exposures.into_iter().for_each(|(stash, exposure)| {
 			log!(
 				trace,
-				"stored exposure for stash {:?} and {:?} backers",
+				"storing exposure for stash {:?} with {:?} own-stake and {:?} backers",
 				stash,
+				exposure.own,
 				exposure.others.len()
 			);
 			// build elected stash.
 			elected_stashes_page.push(stash.clone());
-			// accumulate total stake.
+			// accumulate total stake and backer count for bookkeeping.
 			total_stake_page = total_stake_page.saturating_add(exposure.total);
-			// set or update staker exposure for this era.
 			total_backers += exposure.others.len() as u32;
+			// set or update staker exposure for this era.
 			Eras::<T>::upsert_exposure(new_planned_era, &stash, exposure);
 		});
 
@@ -1025,6 +1068,8 @@ impl<T: Config> EraElectionPlanner<T> {
 		Eras::<T>::add_total_stake(new_planned_era, total_stake_page);
 
 		// collect or update the pref of all winners.
+		// TODO: rather inefficient, we can do this once at the last page across all entries in
+		// `ElectableStashes`.
 		for stash in &elected_stashes {
 			let pref = Validators::<T>::get(stash);
 			Eras::<T>::set_validator_prefs(new_planned_era, stash, pref);