Fixes TotalValueLocked out of sync in nomination pools (#3052)

The `TotalLockedValue` storage value in nomination pools pallet may get
out of sync if the staking pallet does implicit withdrawal of unlocking
chunks belonging to a bonded pool stash. This fix is based on a new
method in the `OnStakingUpdate` traits, `on_withdraw`, which allows the
nomination pools pallet to adjust the `TotalLockedValue` every time
there is an implicit or explicit withdrawal from a bonded pool's stash.

This PR also adds a migration that checks and updates the on-chain TVL
if it got out of sync due to the bug this PR fixes.

**Changes to `trait OnStakingUpdate`**

In order for staking to notify the nomination pools pallet that chunks
where withdrew, we add a new method, `on_withdraw` to the
`OnStakingUpdate` trait. The nomination pools pallet filters the
withdraws that are related to bonded pool accounts and updates the
`TotalValueLocked` accordingly.

**Others**
- Adds try-state checks to the EPM/staking e2e tests
- Adds tests for auto withdrawing in the context of nomination pools

**To-do**
- [x] check if we need a migration to fix the current `TotalValueLocked`
(run try-runtime)
- [x] migrations to fix the current on-chain TVL value 

 ✅ **Kusama**:
```
TotalValueLocked: 99.4559 kKSM
TotalValueLocked (calculated) 99.4559 kKSM
```
⚠️ **Westend**:
```
TotalValueLocked: 18.4060 kWND
TotalValueLocked (calculated) 18.4050 kWND
```
**Polkadot**: TVL not released yet.

Closes https://github.com/paritytech/polkadot-sdk/issues/3055

---------

Co-authored-by: command-bot <>
Co-authored-by: Ross Bulat <ross@parity.io>
Co-authored-by: Dónal Murray <donal.murray@parity.io>

substrate/frame/nomination-pools/src/migration.rs

@@ -56,6 +56,59 @@ pub mod versioned {
 	>;
 }
 
+pub mod unversioned {
+	use super::*;
+
+	/// Checks and updates `TotalValueLocked` if out of sync.
+	pub struct TotalValueLockedSync<T>(sp_std::marker::PhantomData<T>);
+	impl<T: Config> OnRuntimeUpgrade for TotalValueLockedSync<T> {
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
+			Ok(Vec::new())
+		}
+
+		fn on_runtime_upgrade() -> Weight {
+			let migrated = BondedPools::<T>::count();
+
+			// recalcuate the `TotalValueLocked` to compare with the current on-chain TVL which may
+			// be out of sync.
+			let tvl: BalanceOf<T> = helpers::calculate_tvl_by_total_stake::<T>();
+			let onchain_tvl = TotalValueLocked::<T>::get();
+
+			let writes = if tvl != onchain_tvl {
+				TotalValueLocked::<T>::set(tvl);
+
+				log!(
+					info,
+					"on-chain TVL was out of sync, update. Old: {:?}, new: {:?}",
+					onchain_tvl,
+					tvl
+				);
+
+				// writes: onchain version + set total value locked.
+				2
+			} else {
+				log!(info, "on-chain TVL was OK: {:?}", tvl);
+
+				// writes: onchain version write.
+				1
+			};
+
+			// reads: migrated * (BondedPools +  Staking::total_stake) + count + onchain
+			// version
+			//
+			// writes: current version + (maybe) TVL
+			T::DbWeight::get()
+				.reads_writes(migrated.saturating_mul(2).saturating_add(2).into(), writes)
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_: Vec<u8>) -> Result<(), TryRuntimeError> {
+			Ok(())
+		}
+	}
+}
+
 pub mod v8 {
 	use super::{v7::V7BondedPoolInner, *};
 
@@ -146,6 +199,7 @@ pub(crate) mod v7 {
 	}
 
 	impl<T: Config> V7BondedPool<T> {
+		#[allow(dead_code)]
 		fn bonded_account(&self) -> T::AccountId {
 			Pallet::<T>::create_bonded_account(self.id)
 		}
@@ -157,26 +211,12 @@ pub(crate) mod v7 {
 		CountedStorageMap<Pallet<T>, Twox64Concat, PoolId, V7BondedPoolInner<T>>;
 
 	pub struct VersionUncheckedMigrateV6ToV7<T>(sp_std::marker::PhantomData<T>);
-	impl<T: Config> VersionUncheckedMigrateV6ToV7<T> {
-		fn calculate_tvl_by_total_stake() -> BalanceOf<T> {
-			BondedPools::<T>::iter()
-				.map(|(id, inner)| {
-					T::Staking::total_stake(
-						&V7BondedPool { id, inner: inner.clone() }.bonded_account(),
-					)
-					.unwrap_or_default()
-				})
-				.reduce(|acc, total_balance| acc + total_balance)
-				.unwrap_or_default()
-		}
-	}
-
 	impl<T: Config> OnRuntimeUpgrade for VersionUncheckedMigrateV6ToV7<T> {
 		fn on_runtime_upgrade() -> Weight {
 			let migrated = BondedPools::<T>::count();
 			// The TVL should be the sum of all the funds that are actively staked and in the
 			// unbonding process of the account of each pool.
-			let tvl: BalanceOf<T> = Self::calculate_tvl_by_total_stake();
+			let tvl: BalanceOf<T> = helpers::calculate_tvl_by_total_stake::<T>();
 
 			TotalValueLocked::<T>::set(tvl);
 
@@ -198,7 +238,7 @@ pub(crate) mod v7 {
 		fn post_upgrade(_data: Vec<u8>) -> Result<(), TryRuntimeError> {
 			// check that the `TotalValueLocked` written is actually the sum of `total_stake` of the
 			// `BondedPools``
-			let tvl: BalanceOf<T> = Self::calculate_tvl_by_total_stake();
+			let tvl: BalanceOf<T> = helpers::calculate_tvl_by_total_stake::<T>();
 			ensure!(
 				TotalValueLocked::<T>::get() == tvl,
 				"TVL written is not equal to `Staking::total_stake` of all `BondedPools`."
@@ -977,3 +1017,17 @@ pub mod v1 {
 		}
 	}
 }
+
+mod helpers {
+	use super::*;
+
+	pub(crate) fn calculate_tvl_by_total_stake<T: Config>() -> BalanceOf<T> {
+		BondedPools::<T>::iter()
+			.map(|(id, inner)| {
+				T::Staking::total_stake(&BondedPool { id, inner: inner.clone() }.bonded_account())
+					.unwrap_or_default()
+			})
+			.reduce(|acc, total_balance| acc + total_balance)
+			.unwrap_or_default()
+	}
+}