Tvl pool staking

frame/nomination-pools/src/lib.rs

@@ -535,6 +535,39 @@ impl<T: Config> PoolMember<T> {
 		}
 	}
 
+	/// Total balance of the member, both active and unbonding.
+	/// Doesn't mutate state.
+	#[cfg(any(feature = "try-runtime", test))]
+	fn total_balance(&mut self) -> BalanceOf<T> {
+		if let Some(pool) = BondedPool::<T>::get(self.pool_id).defensive() {
+			let active_balance = pool.points_to_balance(self.active_points());
+
+			// get all unbonding balance, both already unlocked and yet to be unlocked.
+			match SubPoolsStorage::<T>::get(self.pool_id) {
+				Some(mut sub_pools) => {
+					let unbonding_balance = &self.unbonding_eras.iter().fold(
+						BalanceOf::<T>::zero(),
+						|accumulator, (era, unlocked_points)| {
+							if let Some(era_pool) = sub_pools.with_era.get_mut(era) {
+								accumulator.saturating_add(
+									era_pool.point_to_balance(unlocked_points.clone()),
+								)
+							} else {
+								accumulator.saturating_add(
+									sub_pools.no_era.point_to_balance(unlocked_points.clone()),
+								)
+							}
+						},
+					);
+					active_balance.saturating_add(unbonding_balance.clone())
+				},
+				None => active_balance,
+			}
+		} else {
+			Zero::zero()
+		}
+	}
+
 	/// Total points of this member, both active and unbonding.
 	fn total_points(&self) -> BalanceOf<T> {
 		self.active_points().saturating_add(self.unbonding_points())
@@ -963,9 +996,13 @@ impl<T: Config> BondedPool<T> {
 	}
 
 	/// Issue points to [`Self`] for `new_funds`.
+	/// Increase the [`TotalValueLocked`] by `new_funds`.
 	fn issue(&mut self, new_funds: BalanceOf<T>) -> BalanceOf<T> {
 		let points_to_issue = self.balance_to_point(new_funds);
 		self.points = self.points.saturating_add(points_to_issue);
+		TotalValueLocked::<T>::mutate(|tvl| {
+			tvl.saturating_accrue(new_funds);
+		});
 		points_to_issue
 	}
 
@@ -1183,9 +1220,8 @@ impl<T: Config> BondedPool<T> {
 
 	/// Bond exactly `amount` from `who`'s funds into this pool.
 	///
-	/// If the bond type is `Create`, `Staking::bond` is called, and `who`
-	/// is allowed to be killed. Otherwise, `Staking::bond_extra` is called and `who`
-	/// cannot be killed.
+	/// If the bond is [`BondType::Create`], [`Staking::bond`] is called, and `who` is allowed to be
+	/// killed. Otherwise, [`Staking::bond_extra`] is called and `who` cannot be killed.
 	///
 	/// Returns `Ok(points_issues)`, `Err` otherwise.
 	fn try_bond_funds(
@@ -1416,8 +1452,8 @@ impl<T: Config> UnbondPool<T> {
 		new_points
 	}
 
-	/// Dissolve some points from the unbonding pool, reducing the balance of the pool
-	/// proportionally.
+	/// Dissolve some points from the unbonding pool, reducing the balance of the pool and the
+	/// [`TotalValueLocked`] proportionally.
 	///
 	/// This is the opposite of `issue`.
 	///
@@ -1427,6 +1463,10 @@ impl<T: Config> UnbondPool<T> {
 		self.points = self.points.saturating_sub(points);
 		self.balance = self.balance.saturating_sub(balance_to_unbond);
 
+		TotalValueLocked::<T>::mutate(|tvl| {
+			tvl.saturating_reduce(balance_to_unbond);
+		});
+
 		balance_to_unbond
 	}
 }
@@ -1577,6 +1617,14 @@ pub mod pallet {
 		type MaxUnbonding: Get<u32>;
 	}
 
+	/// The sum of funds across all pools.
+	///
+	/// This might be higher but never lower than the actual sum of the currently unlocking and
+	/// bonded funds as this is only decreased if a user withdraws unlocked funds or a slash
+	/// happened.
+	#[pallet::storage]
+	pub type TotalValueLocked<T: Config> = StorageValue<_, BalanceOf<T>, ValueQuery>;
+
 	/// Minimum amount to bond to join a pool.
 	#[pallet::storage]
 	pub type MinJoinBond<T: Config> = StorageValue<_, BalanceOf<T>, ValueQuery>;
@@ -1795,9 +1843,9 @@ pub mod pallet {
 		CannotWithdrawAny,
 		/// The amount does not meet the minimum bond to either join or create a pool.
 		///
-		/// The depositor can never unbond to a value less than
-		/// `Pallet::depositor_min_bond`. The caller does not have nominating
-		/// permissions for the pool. Members can never unbond to a value below `MinJoinBond`.
+		/// The depositor can never unbond to a value less than `Pallet::depositor_min_bond`. The
+		/// caller does not have nominating permissions for the pool. Members can never unbond to a
+		/// value below `MinJoinBond`.
 		MinimumBondNotMet,
 		/// The transaction could not be executed due to overflow risk for the pool.
 		OverflowRisk,
@@ -2074,7 +2122,7 @@ pub mod pallet {
 
 		/// Call `withdraw_unbonded` for the pools account. This call can be made by any account.
 		///
-		/// This is useful if their are too many unlocking chunks to call `unbond`, and some
+		/// This is useful if there are too many unlocking chunks to call `unbond`, and some
 		/// can be cleared by withdrawing. In the case there are too many unlocking chunks, the user
 		/// would probably see an error like `NoMoreChunks` emitted from the staking system when
 		/// they attempt to unbond.
@@ -3077,6 +3125,7 @@ impl<T: Config> Pallet<T> {
 			bonded_pools == reward_pools,
 			"`BondedPools` and `RewardPools` must all have the EXACT SAME key-set."
 		);
+		// let mut expected_tvl = Zero::zero();
 
 		ensure!(
 			SubPoolsStorage::<T>::iter_keys().all(|k| bonded_pools.contains(&k)),
@@ -3169,8 +3218,12 @@ impl<T: Config> Pallet<T> {
 				"depositor must always have MinCreateBond stake in the pool, except for when the \
 				pool is being destroyed and the depositor is the last member",
 			);
+			//expected_tvl += T::Staking::total_stake(&bonded_pool.bonded_account())
+			//.expect("all pools must have total stake");
+
 			Ok(())
 		})?;
+
 		ensure!(
 			MaxPoolMembers::<T>::get().map_or(true, |max| all_members <= max),
 			Error::<T>::MaxPoolMembers
@@ -3269,12 +3322,18 @@ impl<T: Config> sp_staking::OnStakingUpdate<T::AccountId, BalanceOf<T>> for Pall
 		// anything here.
 		slashed_bonded: BalanceOf<T>,
 		slashed_unlocking: &BTreeMap<EraIndex, BalanceOf<T>>,
+		total_slashed: BalanceOf<T>,
 	) {
-		if let Some(pool_id) = ReversePoolIdLookup::<T>::get(pool_account) {
-			let mut sub_pools = match SubPoolsStorage::<T>::get(pool_id).defensive() {
+		if let Some(pool_id) = ReversePoolIdLookup::<T>::get(pool_account).defensive() {
+			// TODO: bug here: a slash toward a pool who's subpools were missing would not be
+			// tracked. write a test for it.
+			Self::deposit_event(Event::<T>::PoolSlashed { pool_id, balance: slashed_bonded });
+
+			let mut sub_pools = match SubPoolsStorage::<T>::get(pool_id) {
 				Some(sub_pools) => sub_pools,
 				None => return,
 			};
+
 			for (era, slashed_balance) in slashed_unlocking.iter() {
 				if let Some(pool) = sub_pools.with_era.get_mut(era) {
 					pool.balance = *slashed_balance;
@@ -3285,9 +3344,11 @@ impl<T: Config> sp_staking::OnStakingUpdate<T::AccountId, BalanceOf<T>> for Pall
 					});
 				}
 			}
-
-			Self::deposit_event(Event::<T>::PoolSlashed { pool_id, balance: slashed_bonded });
 			SubPoolsStorage::<T>::insert(pool_id, sub_pools);
+
+			TotalValueLocked::<T>::mutate(|tvl| {
+				tvl.saturating_reduce(total_slashed);
+			});
 		}
 	}
 }

frame/nomination-pools/src/migration.rs

@@ -724,4 +724,99 @@ pub mod v5 {
 			Ok(())
 		}
 	}
+
+	/// This migration accumulates and initializes the [`TotalValueLocked`] for all pools.
+	pub struct VersionUncheckedMigrateV5ToV6<T>(sp_std::marker::PhantomData<T>);
+	impl<T: Config> OnRuntimeUpgrade for VersionUncheckedMigrateV5ToV6<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let migrated = BondedPools::<T>::count();
+			let tvl: BalanceOf<T> = BondedPools::<T>::iter()
+				.map(|(id, inner)| {
+					T::Staking::total_stake(&BondedPool { id, inner }.bonded_account())
+						.unwrap_or_default()
+				})
+				.reduce(|acc, balance| acc + balance)
+				.unwrap_or_default();
+
+			TotalValueLocked::<T>::set(tvl);
+
+			log!(info, "Upgraded {} pools", migrated);
+
+			// reads: migrated * (BondedPools +  Staking::total_stake) + count + onchain
+			// version
+			//
+			// writes: current version + TVL
+			T::DbWeight::get().reads_writes(migrated.saturating_mul(2).saturating_add(2).into(), 2)
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
+			ensure!(
+				PoolMembers::<T>::iter_keys().count() == PoolMembers::<T>::iter_values().count(),
+				"There are undecodable PoolMembers in storage. This migration will not fix that."
+			);
+			ensure!(
+				BondedPools::<T>::iter_keys().count() == BondedPools::<T>::iter_values().count(),
+				"There are undecodable BondedPools in storage. This migration will not fix that."
+			);
+			ensure!(
+				SubPoolsStorage::<T>::iter_keys().count() ==
+					SubPoolsStorage::<T>::iter_values().count(),
+				"There are undecodable SubPools in storage. This migration will not fix that."
+			);
+			ensure!(
+				Metadata::<T>::iter_keys().count() == Metadata::<T>::iter_values().count(),
+				"There are undecodable Metadata in storage. This migration will not fix that."
+			);
+
+			Ok(Vec::new())
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_data: Vec<u8>) -> Result<(), TryRuntimeError> {
+			// ensure [`TotalValueLocked`] contains a value greater zero if any instances of
+			// BondedPools exist.
+			if !BondedPools::<T>::count().is_zero() {
+				ensure!(!TotalValueLocked::<T>::get().is_zero(), "tvl written incorrectly");
+			}
+
+			ensure!(
+				Pallet::<T>::on_chain_storage_version() >= 6,
+				"nomination-pools::migration::v6: wrong storage version"
+			);
+
+			// These should not have been touched - just in case.
+			ensure!(
+				PoolMembers::<T>::iter_keys().count() == PoolMembers::<T>::iter_values().count(),
+				"There are undecodable PoolMembers in storage."
+			);
+			ensure!(
+				BondedPools::<T>::iter_keys().count() == BondedPools::<T>::iter_values().count(),
+				"There are undecodable BondedPools in storage."
+			);
+			ensure!(
+				SubPoolsStorage::<T>::iter_keys().count() ==
+					SubPoolsStorage::<T>::iter_values().count(),
+				"There are undecodable SubPools in storage."
+			);
+			ensure!(
+				Metadata::<T>::iter_keys().count() == Metadata::<T>::iter_values().count(),
+				"There are undecodable Metadata in storage."
+			);
+
+			Ok(())
+		}
+	}
+
+	/// [`VersionUncheckedMigrateV5ToV6`] wrapped in a
+	/// [`frame_support::migrations::VersionedRuntimeUpgrade`], ensuring the migration is only
+	/// performed when on-chain version is 5.
+	#[cfg(feature = "experimental")]
+	pub type VersionCheckedMigrateV5ToV6<T, I> = frame_support::migrations::VersionedRuntimeUpgrade<
+		5,
+		6,
+		VersionUncheckedMigrateV5ToV6<T, I>,
+		crate::pallet::Pallet<T, I>,
+		<T as frame_system::Config>::DbWeight,
+	>;
 }

frame/nomination-pools/src/mock.rs

@@ -3,7 +3,7 @@ use crate::{self as pools};
 use frame_support::{assert_ok, parameter_types, PalletId};
 use frame_system::RawOrigin;
 use sp_runtime::{BuildStorage, FixedU128};
-use sp_staking::Stake;
+use sp_staking::{OnStakingUpdate, Stake};
 
 pub type BlockNumber = u64;
 pub type AccountId = u128;
@@ -28,7 +28,8 @@ parameter_types! {
 	pub static CurrentEra: EraIndex = 0;
 	pub static BondingDuration: EraIndex = 3;
 	pub storage BondedBalanceMap: BTreeMap<AccountId, Balance> = Default::default();
-	pub storage UnbondingBalanceMap: BTreeMap<AccountId, Balance> = Default::default();
+	// map from user, to a vec of era to amount being unlocked in that era.
+	pub storage UnbondingBalanceMap: BTreeMap<AccountId, Vec<(EraIndex, Balance)>> = Default::default();
 	#[derive(Clone, PartialEq)]
 	pub static MaxUnbonding: u32 = 8;
 	pub static StakingMinBond: Balance = 10;
@@ -42,6 +43,14 @@ impl StakingMock {
 		x.insert(who, bonded);
 		BondedBalanceMap::set(&x)
 	}
+
+	pub(crate) fn slash_to(pool_id: PoolId, amount: Balance) {
+		let acc = Pools::create_bonded_account(pool_id);
+		let bonded = BondedBalanceMap::get();
+		let pre_total = bonded.get(&acc).unwrap();
+		Self::set_bonded_balance(acc, pre_total - amount);
+		Pools::on_slash(&acc, pre_total - amount, &Default::default(), amount);
+	}
 }
 
 impl sp_staking::StakingInterface for StakingMock {
@@ -87,8 +96,11 @@ impl sp_staking::StakingInterface for StakingMock {
 		let mut x = BondedBalanceMap::get();
 		*x.get_mut(who).unwrap() = x.get_mut(who).unwrap().saturating_sub(amount);
 		BondedBalanceMap::set(&x);
+
+		let era = Self::current_era();
+		let unlocking_at = era + Self::bonding_duration();
 		let mut y = UnbondingBalanceMap::get();
-		*y.entry(*who).or_insert(Self::Balance::zero()) += amount;
+		y.entry(*who).or_insert(Default::default()).push((unlocking_at, amount));
 		UnbondingBalanceMap::set(&y);
 		Ok(())
 	}
@@ -98,11 +110,13 @@ impl sp_staking::StakingInterface for StakingMock {
 	}
 
 	fn withdraw_unbonded(who: Self::AccountId, _: u32) -> Result<bool, DispatchError> {
-		// Simulates removing unlocking chunks and only having the bonded balance locked
-		let mut x = UnbondingBalanceMap::get();
-		x.remove(&who);
-		UnbondingBalanceMap::set(&x);
+		let mut unbonding_map = UnbondingBalanceMap::get();
+		let staker_map = unbonding_map.get_mut(&who).ok_or("not a staker")?;
+
+		let current_era = Self::current_era();
+		staker_map.retain(|(unlocking_at, _amount)| *unlocking_at > current_era);
 
+		UnbondingBalanceMap::set(&unbonding_map);
 		Ok(UnbondingBalanceMap::get().is_empty() && BondedBalanceMap::get().is_empty())
 	}
 
@@ -126,14 +140,17 @@ impl sp_staking::StakingInterface for StakingMock {
 	}
 
 	fn stake(who: &Self::AccountId) -> Result<Stake<Balance>, DispatchError> {
-		match (
-			UnbondingBalanceMap::get().get(who).copied(),
-			BondedBalanceMap::get().get(who).copied(),
-		) {
+		match (UnbondingBalanceMap::get().get(who), BondedBalanceMap::get().get(who).copied()) {
 			(None, None) => Err(DispatchError::Other("balance not found")),
-			(Some(v), None) => Ok(Stake { total: v, active: 0 }),
+			(Some(v), None) => Ok(Stake {
+				total: v.into_iter().fold(0u128, |acc, &x| acc.saturating_add(x.1)),
+				active: 0,
+			}),
 			(None, Some(v)) => Ok(Stake { total: v, active: v }),
-			(Some(a), Some(b)) => Ok(Stake { total: a + b, active: b }),
+			(Some(a), Some(b)) => Ok(Stake {
+				total: a.into_iter().fold(0u128, |acc, &x| acc.saturating_add(x.1)) + b,
+				active: b,
+			}),
 		}
 	}