Add batching to fast-unstake pallet

frame/fast-unstake/src/benchmarking.rs

@@ -112,7 +112,7 @@ fn on_idle_full_block<T: Config>() {
 }
 
 benchmarks! {
-	// on_idle, we we don't check anyone, but fully unbond them.
+	// on_idle, we don't check anyone, but fully unbond them.
 	on_idle_unstake {
 		ErasToCheckPerBlock::<T>::put(1);
 		for who in create_unexposed_nominators::<T>() {
@@ -140,7 +140,7 @@ benchmarks! {
 	verify {
 		assert!(matches!(
 			fast_unstake_events::<T>().last(),
-			Some(Event::Terminated)
+			Some(Event::BatchFinished)
 		));
 	}
 
@@ -178,7 +178,7 @@ benchmarks! {
 		assert!(stashes.iter().all(|(s, _)| request.stashes.iter().find(|(ss, _)| ss == s).is_some()));
 		assert!(matches!(
 			fast_unstake_events::<T>().last(),
-			Some(Event::Checking { .. })
+			Some(Event::BatchChecked { .. })
 		));
 	}
 

frame/fast-unstake/src/lib.rs

@@ -60,6 +60,7 @@ mod tests;
 // NOTE: enable benchmarking in tests as well.
 #[cfg(feature = "runtime-benchmarks")]
 mod benchmarking;
+pub mod migrations;
 pub mod types;
 pub mod weights;
 
@@ -83,7 +84,7 @@ pub mod pallet {
 	use frame_election_provider_support::ElectionProviderBase;
 	use frame_support::{
 		pallet_prelude::*,
-		traits::{Defensive, ReservableCurrency},
+		traits::{Defensive, ReservableCurrency, StorageVersion},
 	};
 	use frame_system::{pallet_prelude::*, RawOrigin};
 	use pallet_staking::Pallet as Staking;
@@ -105,7 +106,10 @@ pub mod pallet {
 		}
 	}
 
+	const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+
 	#[pallet::pallet]
+	#[pallet::storage_version(STORAGE_VERSION)]
 	pub struct Pallet<T>(_);
 
 	#[pallet::config]
@@ -167,12 +171,12 @@ pub mod pallet {
 		/// An internal error happened. Operations will be paused now.
 		InternalError,
 		/// A batch was partially checked for the given eras, but the process did not finish.
-		Checking { eras: Vec<EraIndex> },
+		BatchChecked { eras: Vec<EraIndex> },
 		/// A batch was terminated.
 		///
 		/// This is always follows by a number of `Unstaked` or `Slashed` events, marking the end
 		/// of the batch. A new batch will be created upon next block.
-		Terminated,
+		BatchFinished,
 	}
 
 	#[pallet::error]
@@ -218,6 +222,7 @@ pub mod pallet {
 						return T::DbWeight::get().reads_writes(1, 1)
 					},
 				};
+
 				// insert them back into the queue.
 				Queue::<T>::insert(stash, deposit);
 				Head::<T>::kill();
@@ -401,8 +406,8 @@ pub mod pallet {
 
 			log!(
 				debug,
-				"checking {:?}, eras_to_check_per_block = {:?}, remaining_weight = {:?}",
-				stashes,
+				"checking {:?} stashes, eras_to_check_per_block = {:?}, remaining_weight = {:?}",
+				stashes.len(),
 				eras_to_check_per_block,
 				remaining_weight
 			);
@@ -480,7 +485,7 @@ pub mod pallet {
 			if unchecked_eras_to_check.is_empty() {
 				// `stash` is not exposed in any era now -- we can let go of them now.
 				stashes.into_iter().for_each(|(stash, deposit)| unstake_stash(stash, deposit));
-				Self::deposit_event(Event::<T>::Terminated);
+				Self::deposit_event(Event::<T>::BatchFinished);
 				<T as Config>::WeightInfo::on_idle_unstake()
 			} else {
 				// eras checked so far.
@@ -508,10 +513,10 @@ pub mod pallet {
 				match checked.try_extend(unchecked_eras_to_check.clone().into_iter()) {
 					Ok(_) =>
 						if stashes.is_empty() {
-							Self::deposit_event(Event::<T>::Terminated);
+							Self::deposit_event(Event::<T>::BatchFinished);
 						} else {
 							Head::<T>::put(UnstakeRequest { stashes, checked });
-							Self::deposit_event(Event::<T>::Checking {
+							Self::deposit_event(Event::<T>::BatchChecked {
 								eras: unchecked_eras_to_check,
 							});
 						},

frame/fast-unstake/src/migrations.rs

@@ -0,0 +1,74 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pub mod v1 {
+	use crate::{types::BalanceOf, *};
+	use frame_support::{
+		storage::unhashed,
+		traits::{Defensive, Get, GetStorageVersion, OnRuntimeUpgrade},
+		weights::Weight,
+	};
+	use sp_staking::EraIndex;
+
+	pub struct MigrateToV1<T>(sp_std::marker::PhantomData<T>);
+	impl<T: Config> OnRuntimeUpgrade for MigrateToV1<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let current = Pallet::<T>::current_storage_version();
+			let onchain = Pallet::<T>::on_chain_storage_version();
+
+			log!(
+				info,
+				"Running migration with current storage version {:?} / onchain {:?}",
+				current,
+				onchain
+			);
+
+			if current == 1 && onchain == 0 {
+				// if a head exists, then we put them back into the queue.
+				if Head::<T>::exists() {
+					if let Some((stash, _, deposit)) =
+						unhashed::take::<(T::AccountId, Vec<EraIndex>, BalanceOf<T>)>(
+							&Head::<T>::hashed_key(),
+						)
+						.defensive()
+					{
+						Queue::<T>::insert(stash, deposit);
+						current.put::<Pallet<T>>();
+					} else {
+						// not much we can do here -- head is already deleted.
+					}
+					T::DbWeight::get().reads_writes(2, 3)
+				} else {
+					T::DbWeight::get().reads(2)
+				}
+			} else {
+				log!(info, "Migration did not executed. This probably should be removed");
+				T::DbWeight::get().reads(1)
+			}
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
+			assert_eq!(Pallet::<T>::on_chain_storage_version(), 0);
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
+			assert_eq!(Pallet::<T>::on_chain_storage_version(), 1);
+		}
+	}
+}

frame/fast-unstake/src/mock.rs

@@ -317,6 +317,7 @@ impl ExtBuilder {
 			// because we read this value as a measure of how many validators we have.
 			pallet_staking::ValidatorCount::<Runtime>::put(VALIDATORS_PER_ERA as u32);
 		});
+
 		ext
 	}