Fix fast-unstake for accounts with slashing

frame/fast-unstake/src/lib.rs

@@ -81,6 +81,7 @@ macro_rules! log {
 pub mod pallet {
 	use super::*;
 	use crate::types::*;
+	use sp_std::collections::btree_set::BTreeSet;
 	use frame_support::{
 		pallet_prelude::*,
 		traits::{Defensive, ReservableCurrency, StorageVersion},
@@ -429,9 +430,9 @@ pub mod pallet {
 				}
 			};
 
-			let check_stash = |stash, deposit, eras_checked: &mut u32| {
+			let check_stash = |stash, deposit, eras_checked: &mut BTreeSet<EraIndex>| {
 				let is_exposed = unchecked_eras_to_check.iter().any(|e| {
-					eras_checked.saturating_inc();
+					eras_checked.insert(*e);
 					T::Staking::is_exposed_in_era(&stash, e)
 				});
 
@@ -452,7 +453,7 @@ pub mod pallet {
 				<T as Config>::WeightInfo::on_idle_unstake()
 			} else {
 				// eras checked so far.
-				let mut eras_checked = 0u32;
+				let mut eras_checked = BTreeSet::<EraIndex>::new();
 
 				let pre_length = stashes.len();
 				let stashes: BoundedVec<(T::AccountId, BalanceOf<T>), T::BatchSize> = stashes
@@ -468,7 +469,7 @@ pub mod pallet {
 				log!(
 					debug,
 					"checked {:?} eras, pre stashes: {:?}, post: {:?}",
-					eras_checked,
+					eras_checked.len(),
 					pre_length,
 					post_length,
 				);
@@ -489,7 +490,7 @@ pub mod pallet {
 					},
 				}
 
-				<T as Config>::WeightInfo::on_idle_check(validator_count * eras_checked)
+				<T as Config>::WeightInfo::on_idle_check(validator_count * eras_checked.len() as u32)
 			}
 		}
 	}

frame/staking/src/pallet/impls.rs

@@ -1569,7 +1569,7 @@ impl<T: Config> StakingInterface for Pallet<T> {
 	}
 
 	fn force_unstake(who: Self::AccountId) -> sp_runtime::DispatchResult {
-		let num_slashing_spans = Self::slashing_spans(&who).iter().count() as u32;
+		let num_slashing_spans = Self::slashing_spans(&who).map_or(0, |s| s.iter().count() as u32);
 		Self::force_unstake(RawOrigin::Root.into(), who.clone(), num_slashing_spans)
 	}
 

frame/staking/src/pallet/mod.rs

@@ -982,7 +982,7 @@ pub mod pallet {
 			// `BondingDuration` to proceed with the unbonding.
 			let maybe_withdraw_weight = {
 				if unlocking == T::MaxUnlockingChunks::get() as usize {
-					let real_num_slashing_spans = Self::slashing_spans(&controller).iter().count();
+					let real_num_slashing_spans = Self::slashing_spans(&controller).map_or(0, |s| s.iter().count());
 					Some(Self::do_withdraw_unbonded(&controller, real_num_slashing_spans as u32)?)
 				} else {
 					None

frame/staking/src/tests.rs

@@ -5725,3 +5725,76 @@ fn scale_validator_count_errors() {
 		);
 	})
 }
+
+mod staking_interface {
+	use frame_support::storage::with_storage_layer;
+	use sp_staking::StakingInterface;
+
+	use super::*;
+
+	#[test]
+	fn force_unstake_with_slash_works() {
+		ExtBuilder::default().build_and_execute(|| {
+			// without slash
+			let _ = with_storage_layer::<(), _, _>(|| {
+				// bond an account, can unstake
+				assert_eq!(Staking::bonded(&11), Some(10));
+				assert_ok!(<Staking as StakingInterface>::force_unstake(11));
+				Err(DispatchError::from("revert"))
+			});
+
+			// bond again and add a slash, still can unstake.
+			assert_eq!(Staking::bonded(&11), Some(10));
+			add_slash(&11);
+			assert_ok!(<Staking as StakingInterface>::force_unstake(11));
+		});
+	}
+
+	#[test]
+	fn withdraw_unbonded_with_slash_works() {
+		ExtBuilder::default().build_and_execute(|| {
+			let _ = with_storage_layer::<(), _, _>(|| {
+				assert_ok!(Staking::unbond(RuntimeOrigin::signed(10), 1));
+				Err(DispatchError::from("revert"))
+			});
+
+			add_slash(&10);
+			let mut current_era = 0;
+			assert_ok!(Staking::unbond(RuntimeOrigin::signed(10), 1));
+			let max_unlocking_chunks = <<Test as Config>::MaxUnlockingChunks as Get<u32>>::get();
+
+			// Fill up all unlocking chunks.
+			for i in 0..max_unlocking_chunks - 1 {
+				current_era = i as u32;
+				mock::start_active_era(current_era);
+				assert_ok!(Staking::unbond(RuntimeOrigin::signed(10), 1));
+			}
+
+			current_era += 1;
+			mock::start_active_era(current_era);
+
+			// This chunk is locked at `current_era` through `current_era + 2` (because
+			// `BondingDuration` == 3).
+			assert_ok!(Staking::unbond(RuntimeOrigin::signed(10), 1));
+			assert_eq!(
+				Staking::ledger(&10).map(|l| l.unlocking.len()).unwrap(),
+				<<Test as Config>::MaxUnlockingChunks as Get<u32>>::get() as usize
+			);
+
+			// even though the number of unlocked chunks is the same as `MaxUnlockingChunks`,
+			// unbonding works as expected.
+			for i in current_era..(current_era + max_unlocking_chunks) - 1 {
+				// There is only 1 chunk per era, so we need to be in a new era to create a chunk.
+				current_era = i as u32;
+				mock::start_active_era(current_era);
+				assert_ok!(Staking::unbond(RuntimeOrigin::signed(10), 1));
+			}
+
+			// only slots within last `BondingDuration` are filled.
+			assert_eq!(
+				Staking::ledger(&10).map(|l| l.unlocking.len()).unwrap(),
+				<<Test as Config>::BondingDuration>::get() as usize
+			);
+		})
+	}
+}