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>

Cargo.lock

@@ -9735,6 +9735,7 @@ dependencies = [
  "pallet-bags-list",
  "pallet-balances",
  "pallet-election-provider-multi-phase",
+ "pallet-nomination-pools",
  "pallet-session",
  "pallet-staking",
  "pallet-timestamp",

polkadot/runtime/westend/src/lib.rs

@@ -1647,7 +1647,7 @@ pub mod migrations {
 		pallet_referenda::migration::v1::MigrateV0ToV1<Runtime, ()>,
 		pallet_grandpa::migrations::MigrateV4ToV5<Runtime>,
 		parachains_configuration::migration::v10::MigrateToV10<Runtime>,
-		pallet_nomination_pools::migration::versioned::V7ToV8<Runtime>,
+		pallet_nomination_pools::migration::unversioned::TotalValueLockedSync<Runtime>,
 		UpgradeSessionKeys,
 		frame_support::migrations::RemovePallet<
 			ImOnlinePalletName,

prdoc/pr_3052.prdoc

@@ -0,0 +1,15 @@
+title: "Fixes a scenario where a nomination pool's `TotalValueLocked` is out of sync due to staking's implicit withdraw"
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      The nomination pools pallet `TotalValueLocked` 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 on the `OnStakingUpdate` traits, `on_withdraw`, which allows the
+      nomination pools pallet to adjust the `TotalValueLocked` every time there is an implicit or
+      explicit withdrawal from a bonded pool's stash.
+
+crates: 
+  - name: "pallet-nomination-pools"
+  - name: "pallet-staking"
+  - name: "sp-staking"

substrate/frame/election-provider-multi-phase/test-staking-e2e/Cargo.toml

@@ -35,7 +35,23 @@ frame-election-provider-support = { path = "../../election-provider-support" }
 
 pallet-election-provider-multi-phase = { path = ".." }
 pallet-staking = { path = "../../staking" }
+pallet-nomination-pools = { path = "../../nomination-pools" }
 pallet-bags-list = { path = "../../bags-list" }
 pallet-balances = { path = "../../balances" }
 pallet-timestamp = { path = "../../timestamp" }
 pallet-session = { path = "../../session" }
+
+[features]
+try-runtime = [
+	"frame-election-provider-support/try-runtime",
+	"frame-support/try-runtime",
+	"frame-system/try-runtime",
+	"pallet-bags-list/try-runtime",
+	"pallet-balances/try-runtime",
+	"pallet-election-provider-multi-phase/try-runtime",
+	"pallet-nomination-pools/try-runtime",
+	"pallet-session/try-runtime",
+	"pallet-staking/try-runtime",
+	"pallet-timestamp/try-runtime",
+	"sp-runtime/try-runtime",
+]

substrate/frame/election-provider-multi-phase/test-staking-e2e/src/lib.rs

@@ -53,9 +53,9 @@ fn log_current_time() {
 
 #[test]
 fn block_progression_works() {
-	let (mut ext, pool_state, _) = ExtBuilder::default().build_offchainify();
+	let (ext, pool_state, _) = ExtBuilder::default().build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		assert_eq!(active_era(), 0);
 		assert_eq!(Session::current_index(), 0);
 		assert!(ElectionProviderMultiPhase::current_phase().is_off());
@@ -70,9 +70,9 @@ fn block_progression_works() {
 		assert!(ElectionProviderMultiPhase::current_phase().is_signed());
 	});
 
-	let (mut ext, pool_state, _) = ExtBuilder::default().build_offchainify();
+	let (ext, pool_state, _) = ExtBuilder::default().build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		assert_eq!(active_era(), 0);
 		assert_eq!(Session::current_index(), 0);
 		assert!(ElectionProviderMultiPhase::current_phase().is_off());
@@ -93,12 +93,12 @@ fn offchainify_works() {
 
 	let staking_builder = StakingExtBuilder::default();
 	let epm_builder = EpmExtBuilder::default();
-	let (mut ext, pool_state, _) = ExtBuilder::default()
+	let (ext, pool_state, _) = ExtBuilder::default()
 		.epm(epm_builder)
 		.staking(staking_builder)
 		.build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		// test ocw progression and solution queue if submission when unsigned phase submission is
 		// not delayed.
 		for _ in 0..100 {
@@ -142,9 +142,9 @@ fn offchainify_works() {
 /// restarts. Note that in this test case, the emergency throttling is disabled.
 fn enters_emergency_phase_after_forcing_before_elect() {
 	let epm_builder = EpmExtBuilder::default().disable_emergency_throttling();
-	let (mut ext, pool_state, _) = ExtBuilder::default().epm(epm_builder).build_offchainify();
+	let (ext, pool_state, _) = ExtBuilder::default().epm(epm_builder).build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		log!(
 			trace,
 			"current validators (staking): {:?}",
@@ -213,12 +213,12 @@ fn continous_slashes_below_offending_threshold() {
 	let staking_builder = StakingExtBuilder::default().validator_count(10);
 	let epm_builder = EpmExtBuilder::default().disable_emergency_throttling();
 
-	let (mut ext, pool_state, _) = ExtBuilder::default()
+	let (ext, pool_state, _) = ExtBuilder::default()
 		.epm(epm_builder)
 		.staking(staking_builder)
 		.build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		assert_eq!(Session::validators().len(), 10);
 		let mut active_validator_set = Session::validators();
 
@@ -271,12 +271,12 @@ fn set_validation_intention_after_chilled() {
 	use frame_election_provider_support::SortedListProvider;
 	use pallet_staking::{Event, Forcing, Nominators};
 
-	let (mut ext, pool_state, _) = ExtBuilder::default()
+	let (ext, pool_state, _) = ExtBuilder::default()
 		.epm(EpmExtBuilder::default())
 		.staking(StakingExtBuilder::default())
 		.build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		assert_eq!(active_era(), 0);
 		// validator is part of the validator set.
 		assert!(Session::validators().contains(&41));
@@ -334,10 +334,10 @@ fn set_validation_intention_after_chilled() {
 fn ledger_consistency_active_balance_below_ed() {
 	use pallet_staking::{Error, Event};
 
-	let (mut ext, pool_state, _) =
+	let (ext, pool_state, _) =
 		ExtBuilder::default().staking(StakingExtBuilder::default()).build_offchainify();
 
-	ext.execute_with(|| {
+	execute_with(ext, || {
 		assert_eq!(Staking::ledger(11.into()).unwrap().active, 1000);
 
 		// unbonding total of active stake fails because the active ledger balance would fall
@@ -387,3 +387,141 @@ fn ledger_consistency_active_balance_below_ed() {
 		assert!(Staking::ledger(11.into()).is_err());
 	});
 }
+
+#[test]
+/// Automatic withdrawal of unlocking funds in staking propagates to the nomination pools and its
+/// state correctly.
+///
+/// The staking pallet may withdraw unlocking funds from a pool's bonded account without a pool
+/// member or operator calling explicitly `Call::withdraw*`. This test verifies that the member's
+/// are eventually paid and the `TotalValueLocked` is kept in sync in those cases.
+fn automatic_unbonding_pools() {
+	use pallet_nomination_pools::TotalValueLocked;
+
+	// closure to fetch the staking unlocking chunks of an account.
+	let unlocking_chunks_of = |account: AccountId| -> usize {
+		Staking::ledger(sp_staking::StakingAccount::Controller(account))
+			.unwrap()
+			.unlocking
+			.len()
+	};
+
+	let (ext, pool_state, _) = ExtBuilder::default()
+		.pools(PoolsExtBuilder::default().max_unbonding(1))
+		.staking(StakingExtBuilder::default().max_unlocking(1).bonding_duration(2))
+		.build_offchainify();
+
+	execute_with(ext, || {
+		assert_eq!(<Runtime as pallet_staking::Config>::MaxUnlockingChunks::get(), 1);
+		assert_eq!(<Runtime as pallet_staking::Config>::BondingDuration::get(), 2);
+		assert_eq!(<Runtime as pallet_nomination_pools::Config>::MaxUnbonding::get(), 1);
+
+		// init state of pool members.
+		let init_free_balance_2 = Balances::free_balance(2);
+		let init_free_balance_3 = Balances::free_balance(3);
+
+		let pool_bonded_account = Pools::create_bonded_account(1);
+
+		// creates a pool with 5 bonded, owned by 1.
+		assert_ok!(Pools::create(RuntimeOrigin::signed(1), 5, 1, 1, 1));
+		assert_eq!(locked_amount_for(pool_bonded_account), 5);
+
+		let init_tvl = TotalValueLocked::<Runtime>::get();
+
+		// 2 joins the pool.
+		assert_ok!(Pools::join(RuntimeOrigin::signed(2), 10, 1));
+		assert_eq!(locked_amount_for(pool_bonded_account), 15);
+
+		// 3 joins the pool.
+		assert_ok!(Pools::join(RuntimeOrigin::signed(3), 10, 1));
+		assert_eq!(locked_amount_for(pool_bonded_account), 25);
+
+		assert_eq!(TotalValueLocked::<Runtime>::get(), 25);
+
+		// currently unlocking 0 chunks in the bonded pools ledger.
+		assert_eq!(unlocking_chunks_of(pool_bonded_account), 0);
+
+		// unbond 2 from pool.
+		assert_ok!(Pools::unbond(RuntimeOrigin::signed(2), 2, 10));
+
+		// amount is still locked in the pool, needs to wait for unbonding period.
+		assert_eq!(locked_amount_for(pool_bonded_account), 25);
+
+		// max chunks in the ledger are now filled up (`MaxUnlockingChunks == 1`).
+		assert_eq!(unlocking_chunks_of(pool_bonded_account), 1);
+
+		// tries to unbond 3 from pool. it will fail since there are no unlocking chunks left
+		// available and the current in the queue haven't been there for more than bonding
+		// duration.
+		assert_err!(
+			Pools::unbond(RuntimeOrigin::signed(3), 3, 10),
+			pallet_staking::Error::<Runtime>::NoMoreChunks
+		);
+
+		assert_eq!(current_era(), 0);
+
+		// progress over bonding duration.
+		for _ in 0..=<Runtime as pallet_staking::Config>::BondingDuration::get() {
+			start_next_active_era(pool_state.clone()).unwrap();
+		}
+		assert_eq!(current_era(), 3);
+		System::reset_events();
+
+		let locked_before_withdraw_pool = locked_amount_for(pool_bonded_account);
+		assert_eq!(Balances::free_balance(pool_bonded_account), 26);
+
+		// now unbonding 3 will work, although the pool's ledger still has the unlocking chunks
+		// filled up.
+		assert_ok!(Pools::unbond(RuntimeOrigin::signed(3), 3, 10));
+		assert_eq!(unlocking_chunks_of(pool_bonded_account), 1);
+
+		assert_eq!(
+			staking_events(),
+			[
+				// auto-withdraw happened as expected to release 2's unbonding funds, but the funds
+				// were not transfered to 2 and stay in the pool's tranferrable balance instead.
+				pallet_staking::Event::Withdrawn { stash: 7939698191839293293, amount: 10 },
+				pallet_staking::Event::Unbonded { stash: 7939698191839293293, amount: 10 }
+			]
+		);
+
+		// balance of the pool remains the same, it hasn't withdraw explicitly from the pool yet.
+		assert_eq!(Balances::free_balance(pool_bonded_account), 26);
+		// but the locked amount in the pool's account decreases due to the auto-withdraw:
+		assert_eq!(locked_before_withdraw_pool - 10, locked_amount_for(pool_bonded_account));
+
+		// TVL correctly updated.
+		assert_eq!(TotalValueLocked::<Runtime>::get(), 25 - 10);
+
+		// however, note that the withdrawing from the pool still works for 2, the funds are taken
+		// from the pool's free balance.
+		assert_eq!(Balances::free_balance(pool_bonded_account), 26);
+		assert_ok!(Pools::withdraw_unbonded(RuntimeOrigin::signed(2), 2, 10));
+		assert_eq!(Balances::free_balance(pool_bonded_account), 16);
+
+		assert_eq!(Balances::free_balance(2), 20);
+		assert_eq!(TotalValueLocked::<Runtime>::get(), 15);
+
+		// 3 cannot withdraw yet.
+		assert_err!(
+			Pools::withdraw_unbonded(RuntimeOrigin::signed(3), 3, 10),
+			pallet_nomination_pools::Error::<Runtime>::CannotWithdrawAny
+		);
+
+		// progress over bonding duration.
+		for _ in 0..=<Runtime as pallet_staking::Config>::BondingDuration::get() {
+			start_next_active_era(pool_state.clone()).unwrap();
+		}
+		assert_eq!(current_era(), 6);
+		System::reset_events();
+
+		assert_ok!(Pools::withdraw_unbonded(RuntimeOrigin::signed(3), 3, 10));
+
+		// final conditions are the expected.
+		assert_eq!(Balances::free_balance(pool_bonded_account), 6); // 5 init bonded + ED
+		assert_eq!(Balances::free_balance(2), init_free_balance_2);
+		assert_eq!(Balances::free_balance(3), init_free_balance_3);
+
+		assert_eq!(TotalValueLocked::<Runtime>::get(), init_tvl);
+	});
+}

substrate/frame/election-provider-multi-phase/test-staking-e2e/src/mock.rs

@@ -69,6 +69,7 @@ frame_support::construct_runtime!(
 		System: frame_system,
 		ElectionProviderMultiPhase: pallet_election_provider_multi_phase,
 		Staking: pallet_staking,
+		Pools: pallet_nomination_pools,
 		Balances: pallet_balances,
 		BagsList: pallet_bags_list,
 		Session: pallet_session,
@@ -114,7 +115,7 @@ impl pallet_balances::Config for Runtime {
 	type MaxFreezes = traits::ConstU32<1>;
 	type RuntimeHoldReason = RuntimeHoldReason;
 	type RuntimeFreezeReason = RuntimeFreezeReason;
-	type FreezeIdentifier = ();
+	type FreezeIdentifier = RuntimeFreezeReason;
 	type WeightInfo = ();
 }
 
@@ -233,7 +234,7 @@ const THRESHOLDS: [VoteWeight; 9] = [10, 20, 30, 40, 50, 60, 1_000, 2_000, 10_00
 parameter_types! {
 	pub static BagThresholds: &'static [sp_npos_elections::VoteWeight] = &THRESHOLDS;
 	pub const SessionsPerEra: sp_staking::SessionIndex = 2;
-	pub const BondingDuration: sp_staking::EraIndex = 28;
+	pub static BondingDuration: sp_staking::EraIndex = 28;
 	pub const SlashDeferDuration: sp_staking::EraIndex = 7; // 1/4 the bonding duration.
 	pub const OffendingValidatorsThreshold: Perbill = Perbill::from_percent(40);
 	pub HistoryDepth: u32 = 84;
@@ -247,6 +248,45 @@ impl pallet_bags_list::Config for Runtime {
 	type Score = VoteWeight;
 }
 
+pub struct BalanceToU256;
+impl sp_runtime::traits::Convert<Balance, sp_core::U256> for BalanceToU256 {
+	fn convert(n: Balance) -> sp_core::U256 {
+		n.into()
+	}
+}
+
+pub struct U256ToBalance;
+impl sp_runtime::traits::Convert<sp_core::U256, Balance> for U256ToBalance {
+	fn convert(n: sp_core::U256) -> Balance {
+		n.try_into().unwrap()
+	}
+}
+
+parameter_types! {
+	pub const PoolsPalletId: frame_support::PalletId = frame_support::PalletId(*b"py/nopls");
+	pub static MaxUnbonding: u32 = 8;
+}
+
+impl pallet_nomination_pools::Config for Runtime {
+	type RuntimeEvent = RuntimeEvent;
+	type WeightInfo = ();
+	type Currency = Balances;
+	type RuntimeFreezeReason = RuntimeFreezeReason;
+	type RewardCounter = sp_runtime::FixedU128;
+	type BalanceToU256 = BalanceToU256;
+	type U256ToBalance = U256ToBalance;
+	type Staking = Staking;
+	type PostUnbondingPoolsWindow = ConstU32<2>;
+	type PalletId = PoolsPalletId;
+	type MaxMetadataLen = ConstU32<256>;
+	type MaxUnbonding = MaxUnbonding;
+	type MaxPointsToBalance = frame_support::traits::ConstU8<10>;
+}
+
+parameter_types! {
+	pub static MaxUnlockingChunks: u32 = 32;
+}
+
 /// Upper limit on the number of NPOS nominations.
 const MAX_QUOTA_NOMINATIONS: u32 = 16;
 
@@ -273,10 +313,10 @@ impl pallet_staking::Config for Runtime {
 	type VoterList = BagsList;
 	type NominationsQuota = pallet_staking::FixedNominationsQuota<MAX_QUOTA_NOMINATIONS>;
 	type TargetList = pallet_staking::UseValidatorsMap<Self>;
-	type MaxUnlockingChunks = ConstU32<32>;
+	type MaxUnlockingChunks = MaxUnlockingChunks;
 	type MaxControllersInDeprecationBatch = ConstU32<100>;
 	type HistoryDepth = HistoryDepth;
-	type EventListeners = ();
+	type EventListeners = Pools;
 	type WeightInfo = pallet_staking::weights::SubstrateWeight<Runtime>;
 	type BenchmarkingConfig = pallet_staking::TestBenchmarkingConfig;
 }
@@ -394,6 +434,14 @@ impl StakingExtBuilder {
 		self.validator_count = n;
 		self
 	}
+	pub fn max_unlocking(self, max: u32) -> Self {
+		<MaxUnlockingChunks>::set(max);
+		self
+	}
+	pub fn bonding_duration(self, eras: EraIndex) -> Self {
+		<BondingDuration>::set(eras);
+		self
+	}
 }
 
 pub struct EpmExtBuilder {}
@@ -417,6 +465,21 @@ impl EpmExtBuilder {
 	}
 }
 
+pub struct PoolsExtBuilder {}
+
+impl Default for PoolsExtBuilder {
+	fn default() -> Self {
+		PoolsExtBuilder {}
+	}
+}
+
+impl PoolsExtBuilder {
+	pub fn max_unbonding(self, max: u32) -> Self {
+		<MaxUnbonding>::set(max);
+		self
+	}
+}
+
 pub struct BalancesExtBuilder {
 	balances: Vec<(AccountId, Balance)>,
 }
@@ -464,6 +527,7 @@ pub struct ExtBuilder {
 	staking_builder: StakingExtBuilder,
 	epm_builder: EpmExtBuilder,
 	balances_builder: BalancesExtBuilder,
+	pools_builder: PoolsExtBuilder,
 }
 
 impl Default for ExtBuilder {
@@ -472,6 +536,7 @@ impl Default for ExtBuilder {
 			staking_builder: StakingExtBuilder::default(),
 			epm_builder: EpmExtBuilder::default(),
 			balances_builder: BalancesExtBuilder::default(),
+			pools_builder: PoolsExtBuilder::default(),
 		}
 	}
 }
@@ -548,6 +613,11 @@ impl ExtBuilder {
 		self
 	}
 
+	pub fn pools(mut self, builder: PoolsExtBuilder) -> Self {
+		self.pools_builder = builder;
+		self
+	}
+
 	pub fn balances(mut self, builder: BalancesExtBuilder) -> Self {
 		self.balances_builder = builder;
 		self
@@ -567,20 +637,20 @@ impl ExtBuilder {
 
 		(ext, pool_state, offchain_state)
 	}
+}
 
-	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
-		let mut ext = self.build();
-		ext.execute_with(test);
+pub(crate) fn execute_with(mut ext: sp_io::TestExternalities, test: impl FnOnce() -> ()) {
+	ext.execute_with(test);
 
-		#[cfg(feature = "try-runtime")]
-		ext.execute_with(|| {
-			let bn = System::block_number();
+	#[cfg(feature = "try-runtime")]
+	ext.execute_with(|| {
+		let bn = System::block_number();
 
-			assert_ok!(<MultiPhase as Hooks<u64>>::try_state(bn));
-			assert_ok!(<Staking as Hooks<u64>>::try_state(bn));
-			assert_ok!(<Session as Hooks<u64>>::try_state(bn));
-		});
-	}
+		assert_ok!(<ElectionProviderMultiPhase as Hooks<BlockNumber>>::try_state(bn));
+		assert_ok!(<Staking as Hooks<BlockNumber>>::try_state(bn));
+		assert_ok!(<Pools as Hooks<BlockNumber>>::try_state(bn));
+		assert_ok!(<Session as Hooks<BlockNumber>>::try_state(bn));
+	});
 }
 
 // Progress to given block, triggering session and era changes as we progress and ensuring that
@@ -606,6 +676,7 @@ pub fn roll_to(n: BlockNumber, delay_solution: bool) {
 		if b != n {
 			Staking::on_finalize(System::block_number());
 		}
+		Pools::on_initialize(b);
 
 		log_current_time();
 	}
@@ -860,6 +931,11 @@ pub(crate) fn set_minimum_election_score(
 	.map_err(|_| ())
 }
 
+pub(crate) fn locked_amount_for(account_id: AccountId) -> Balance {
+	let lock = pallet_balances::Locks::<Runtime>::get(account_id);
+	lock[0].amount
+}
+
 pub(crate) fn staking_events() -> Vec<pallet_staking::Event<Runtime>> {
 	System::events()
 		.into_iter()

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

@@ -1293,27 +1293,6 @@ impl<T: Config> BondedPool<T> {
 			});
 		};
 	}
-
-	/// Withdraw all the funds that are already unlocked from staking for the
-	/// [`BondedPool::bonded_account`].
-	///
-	/// Also reduces the [`TotalValueLocked`] by the difference of the
-	/// [`T::Staking::total_stake`] of the [`BondedPool::bonded_account`] that might occur by
-	/// [`T::Staking::withdraw_unbonded`].
-	///
-	/// Returns the result of [`T::Staking::withdraw_unbonded`]
-	fn withdraw_from_staking(&self, num_slashing_spans: u32) -> Result<bool, DispatchError> {
-		let bonded_account = self.bonded_account();
-
-		let prev_total = T::Staking::total_stake(&bonded_account.clone()).unwrap_or_default();
-		let outcome = T::Staking::withdraw_unbonded(bonded_account.clone(), num_slashing_spans);
-		let diff = prev_total
-			.defensive_saturating_sub(T::Staking::total_stake(&bonded_account).unwrap_or_default());
-		TotalValueLocked::<T>::mutate(|tvl| {
-			tvl.saturating_reduce(diff);
-		});
-		outcome
-	}
 }
 
 /// A reward pool.
@@ -1733,7 +1712,7 @@ pub mod pallet {
 		CountedStorageMap<_, Twox64Concat, PoolId, BondedPoolInner<T>>;
 
 	/// Reward pools. This is where there rewards for each pool accumulate. When a members payout is
-	/// claimed, the balance comes out fo the reward pool. Keyed by the bonded pools account.
+	/// claimed, the balance comes out of the reward pool. Keyed by the bonded pools account.
 	#[pallet::storage]
 	pub type RewardPools<T: Config> = CountedStorageMap<_, Twox64Concat, PoolId, RewardPool<T>>;
 
@@ -1753,8 +1732,8 @@ pub mod pallet {
 
 	/// A reverse lookup from the pool's account id to its id.
 	///
-	/// This is only used for slashing. In all other instances, the pool id is used, and the
-	/// accounts are deterministically derived from it.
+	/// This is only used for slashing and on automatic withdraw update. In all other instances, the
+	/// pool id is used, and the accounts are deterministically derived from it.
 	#[pallet::storage]
 	pub type ReversePoolIdLookup<T: Config> =
 		CountedStorageMap<_, Twox64Concat, T::AccountId, PoolId, OptionQuery>;
@@ -2223,7 +2202,7 @@ pub mod pallet {
 			// For now we only allow a pool to withdraw unbonded if its not destroying. If the pool
 			// is destroying then `withdraw_unbonded` can be used.
 			ensure!(pool.state != PoolState::Destroying, Error::<T>::NotDestroying);
-			pool.withdraw_from_staking(num_slashing_spans)?;
+			T::Staking::withdraw_unbonded(pool.bonded_account(), num_slashing_spans)?;
 
 			Ok(())
 		}
@@ -2275,7 +2254,8 @@ pub mod pallet {
 
 			// Before calculating the `balance_to_unbond`, we call withdraw unbonded to ensure the
 			// `transferrable_balance` is correct.
-			let stash_killed = bonded_pool.withdraw_from_staking(num_slashing_spans)?;
+			let stash_killed =
+				T::Staking::withdraw_unbonded(bonded_pool.bonded_account(), num_slashing_spans)?;
 
 			// defensive-only: the depositor puts enough funds into the stash so that it will only
 			// be destroyed when they are leaving.
@@ -3628,4 +3608,14 @@ impl<T: Config> sp_staking::OnStakingUpdate<T::AccountId, BalanceOf<T>> for Pall
 		}
 		Self::deposit_event(Event::<T>::PoolSlashed { pool_id, balance: slashed_bonded });
 	}
+
+	/// Reduces the overall `TotalValueLocked` if a withdrawal happened for a pool involved in the
+	/// staking withdraw.
+	fn on_withdraw(pool_account: &T::AccountId, amount: BalanceOf<T>) {
+		if ReversePoolIdLookup::<T>::get(pool_account).is_some() {
+			TotalValueLocked::<T>::mutate(|tvl| {
+				tvl.saturating_reduce(amount);
+			});
+		}
+	}
 }

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()
+	}
+}

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

@@ -134,11 +134,24 @@ impl sp_staking::StakingInterface for StakingMock {
 
 	fn withdraw_unbonded(who: Self::AccountId, _: u32) -> Result<bool, DispatchError> {
 		let mut unbonding_map = UnbondingBalanceMap::get();
+
+		// closure to calculate the current unlocking funds across all eras/accounts.
+		let unlocking = |pair: &Vec<(EraIndex, Balance)>| -> Balance {
+			pair.iter()
+				.try_fold(Zero::zero(), |acc: Balance, (_at, amount)| acc.checked_add(*amount))
+				.unwrap()
+		};
+
 		let staker_map = unbonding_map.get_mut(&who).ok_or("Nothing to unbond")?;
+		let unlocking_before = unlocking(&staker_map);
 
 		let current_era = Self::current_era();
+
 		staker_map.retain(|(unlocking_at, _amount)| *unlocking_at > current_era);
 
+		// if there was a withdrawal, notify the pallet.
+		Pools::on_withdraw(&who, unlocking_before.saturating_sub(unlocking(&staker_map)));
+
 		UnbondingBalanceMap::set(&unbonding_map);
 		Ok(UnbondingBalanceMap::get().is_empty() && BondedBalanceMap::get().is_empty())
 	}

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

@@ -41,7 +41,7 @@ use sp_runtime::{
 use sp_staking::{
 	currency_to_vote::CurrencyToVote,
 	offence::{DisableStrategy, OffenceDetails, OnOffenceHandler},
-	EraIndex, Page, SessionIndex, Stake,
+	EraIndex, OnStakingUpdate, Page, SessionIndex, Stake,
 	StakingAccount::{self, Controller, Stash},
 	StakingInterface,
 };
@@ -150,6 +150,9 @@ impl<T: Config> Pallet<T> {
 			// Already checked that this won't overflow by entry condition.
 			let value = old_total.defensive_saturating_sub(new_total);
 			Self::deposit_event(Event::<T>::Withdrawn { stash, amount: value });
+
+			// notify listeners.
+			T::EventListeners::on_withdraw(controller, value);
 		}
 
 		Ok(used_weight)

substrate/frame/staking/src/pallet/mod.rs

@@ -275,7 +275,7 @@ pub mod pallet {
 		/// Something that listens to staking updates and performs actions based on the data it
 		/// receives.
 		///
-		/// WARNING: this only reports slashing events for the time being.
+		/// WARNING: this only reports slashing and withdraw events for the time being.
 		type EventListeners: sp_staking::OnStakingUpdate<Self::AccountId, BalanceOf<Self>>;
 
 		/// Some parameters of the benchmarking.

substrate/primitives/staking/src/lib.rs

@@ -150,6 +150,9 @@ pub trait OnStakingUpdate<AccountId, Balance> {
 		_slashed_total: Balance,
 	) {
 	}
+
+	/// Fired when a portion of a staker's balance has been withdrawn.
+	fn on_withdraw(_stash: &AccountId, _amount: Balance) {}
 }
 
 /// A generic representation of a staking implementation.