[Pools] Ensure members can always exit the pool gracefully

prdoc/pr_4998.prdoc

@@ -0,0 +1,20 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Ensure members can always exit the pool gracefully
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      Ensures when a member wants to withdraw all funds but the pool is not able to provide all their funds, the member
+      can receive as much as possible and exit pool. Also handles cases where some extra funds held in member's account
+      is released when they are removed.
+
+crates:
+  - name: pallet-delegated-staking
+    bump: patch
+  - name: pallet-nomination-pools
+    bump: minor
+  - name: sp-staking
+    bump: minor
+

substrate/frame/delegated-staking/src/impls.rs

@@ -32,28 +32,34 @@ impl<T: Config> DelegationInterface for Pallet<T> {
 			.ok()
 	}
 
+	fn agent_transferable_balance(agent: Agent<Self::AccountId>) -> Option<Self::Balance> {
+		AgentLedgerOuter::<T>::get(&agent.get())
+			.map(|a| a.ledger.unclaimed_withdrawals)
+			.ok()
+	}
+
 	fn delegator_balance(delegator: Delegator<Self::AccountId>) -> Option<Self::Balance> {
 		Delegation::<T>::get(&delegator.get()).map(|d| d.amount)
 	}
 
 	/// Delegate funds to an `Agent`.
-	fn delegate(
-		who: Delegator<Self::AccountId>,
+	fn register_agent(
 		agent: Agent<Self::AccountId>,
 		reward_account: &Self::AccountId,
-		amount: Self::Balance,
 	) -> DispatchResult {
 		Pallet::<T>::register_agent(
 			RawOrigin::Signed(agent.clone().get()).into(),
 			reward_account.clone(),
-		)?;
+		)
+	}
 
-		// Delegate the funds from who to the `Agent` account.
-		Pallet::<T>::delegate_to_agent(RawOrigin::Signed(who.get()).into(), agent.get(), amount)
+	/// Remove `Agent` registration.
+	fn remove_agent(agent: Agent<Self::AccountId>) -> DispatchResult {
+		Pallet::<T>::remove_agent(RawOrigin::Signed(agent.clone().get()).into())
 	}
 
 	/// Add more delegation to the `Agent` account.
-	fn delegate_extra(
+	fn delegate(
 		who: Delegator<Self::AccountId>,
 		agent: Agent<Self::AccountId>,
 		amount: Self::Balance,
@@ -118,7 +124,7 @@ impl<T: Config> DelegationMigrator for Pallet<T> {
 
 	/// Only used for testing.
 	#[cfg(feature = "runtime-benchmarks")]
-	fn drop_agent(agent: Agent<Self::AccountId>) {
+	fn migrate_to_direct_staker(agent: Agent<Self::AccountId>) {
 		<Agents<T>>::remove(agent.clone().get());
 		<Delegators<T>>::iter()
 			.filter(|(_, delegation)| delegation.agent == agent.clone().get())

substrate/frame/delegated-staking/src/lib.rs

@@ -304,6 +304,27 @@ pub mod pallet {
 			Ok(())
 		}
 
+		/// Remove an account from being an `Agent`.
+		///
+		/// This can only be called if the agent has no delegated funds, no pending slashes and no
+		/// unclaimed withdrawals.
+		pub fn remove_agent(origin: OriginFor<T>) -> DispatchResult {
+			let who = ensure_signed(origin)?;
+			let ledger = AgentLedger::<T>::get(&who).ok_or(Error::<T>::NotAgent)?;
+
+			ensure!(
+				ledger.total_delegated == Zero::zero() &&
+					ledger.pending_slash == Zero::zero() &&
+					ledger.unclaimed_withdrawals == Zero::zero(),
+				Error::<T>::NotAllowed
+			);
+
+			// remove provider reference
+			let _ = frame_system::Pallet::<T>::dec_providers(&who).defensive();
+			<Agents<T>>::remove(who);
+			Ok(())
+		}
+
 		/// Migrate from a `Nominator` account to `Agent` account.
 		///
 		/// The origin needs to
@@ -599,44 +620,19 @@ impl<T: Config> Pallet<T> {
 		ensure!(delegation.agent == agent, Error::<T>::NotAgent);
 		ensure!(delegation.amount >= amount, Error::<T>::NotEnoughFunds);
 
-		// if we do not already have enough funds to be claimed, try withdraw some more.
-		// keep track if we killed the staker in the process.
-		let stash_killed = if agent_ledger.ledger.unclaimed_withdrawals < amount {
+		// if we do not already have enough funds to be claimed, try to withdraw some more.
+		if agent_ledger.ledger.unclaimed_withdrawals < amount {
 			// withdraw account.
-			let killed = T::CoreStaking::withdraw_unbonded(agent.clone(), num_slashing_spans)
+			let _ = T::CoreStaking::withdraw_unbonded(agent.clone(), num_slashing_spans)
 				.map_err(|_| Error::<T>::WithdrawFailed)?;
 			// reload agent from storage since withdrawal might have changed the state.
 			agent_ledger = agent_ledger.reload()?;
-			Some(killed)
-		} else {
-			None
-		};
+		}
 
 		// if we still do not have enough funds to release, abort.
 		ensure!(agent_ledger.ledger.unclaimed_withdrawals >= amount, Error::<T>::NotEnoughFunds);
+		agent_ledger.remove_unclaimed_withdraw(amount)?.update();
 
-		// Claim withdraw from agent. Kill agent if no delegation left.
-		// TODO: Ideally if there is a register, there should be an unregister that should
-		// clean up the agent. Can be improved in future.
-		if agent_ledger.remove_unclaimed_withdraw(amount)?.update_or_kill()? {
-			match stash_killed {
-				Some(killed) => {
-					// this implies we did a `CoreStaking::withdraw` before release. Ensure
-					// we killed the staker as well.
-					ensure!(killed, Error::<T>::BadState);
-				},
-				None => {
-					// We did not do a `CoreStaking::withdraw` before release. Ensure staker is
-					// already killed in `CoreStaking`.
-					ensure!(T::CoreStaking::status(&agent).is_err(), Error::<T>::BadState);
-				},
-			}
-
-			// Remove provider reference for `who`.
-			let _ = frame_system::Pallet::<T>::dec_providers(&agent).defensive();
-		}
-
-		// book keep delegation
 		delegation.amount = delegation
 			.amount
 			.checked_sub(&amount)

substrate/frame/delegated-staking/src/tests.rs

@@ -943,7 +943,7 @@ mod pool_integration {
 				vec![
 					PoolsEvent::Withdrawn { member: 302, pool_id, balance: 100, points: 100 },
 					PoolsEvent::Withdrawn { member: 303, pool_id, balance: 200, points: 200 },
-					PoolsEvent::MemberRemoved { pool_id: 1, member: 303 },
+					PoolsEvent::MemberRemoved { pool_id: 1, member: 303, balance: 0 },
 				]
 			);
 
@@ -1055,7 +1055,7 @@ mod pool_integration {
 						balance: creator_stake,
 						points: creator_stake,
 					},
-					PoolsEvent::MemberRemoved { pool_id, member: creator },
+					PoolsEvent::MemberRemoved { pool_id, member: creator, balance: 0 },
 					PoolsEvent::Destroyed { pool_id },
 				]
 			);

substrate/frame/delegated-staking/src/types.rs

@@ -251,25 +251,10 @@ impl<T: Config> AgentLedgerOuter<T> {
 		self.ledger.update(&key)
 	}
 
-	/// Save self and remove if no delegation left.
-	///
-	/// Returns:
-	/// - true if agent killed.
-	/// - error if the delegate is in an unexpected state.
-	pub(crate) fn update_or_kill(self) -> Result<bool, DispatchError> {
+	/// Update agent ledger.
+	pub(crate) fn update(self) {
 		let key = self.key;
-		// see if delegate can be killed
-		if self.ledger.total_delegated == Zero::zero() {
-			ensure!(
-				self.ledger.unclaimed_withdrawals == Zero::zero() &&
-					self.ledger.pending_slash == Zero::zero(),
-				Error::<T>::BadState
-			);
-			<Agents<T>>::remove(key);
-			return Ok(true)
-		}
 		self.ledger.update(&key);
-		Ok(false)
 	}
 
 	/// Reloads self from storage.

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

@@ -106,9 +106,11 @@ pub trait StakeStrategy {
 
 	/// Balance that can be transferred from pool account to member.
 	///
-	/// This is part of the pool balance that is not actively staked. That is, tokens that are
-	/// in unbonding period or unbonded.
-	fn transferable_balance(pool_account: Pool<Self::AccountId>) -> Self::Balance;
+	/// This is part of the pool balance that can be withdrawn.
+	fn transferable_balance(
+		pool_account: Pool<Self::AccountId>,
+		member_account: Member<Self::AccountId>,
+	) -> Self::Balance;
 
 	/// Total balance of the pool including amount that is actively staked.
 	fn total_balance(pool_account: Pool<Self::AccountId>) -> Option<Self::Balance>;
@@ -181,6 +183,9 @@ pub trait StakeStrategy {
 		num_slashing_spans: u32,
 	) -> DispatchResult;
 
+	/// Dissolve the pool account.
+	fn dissolve(pool_account: Pool<Self::AccountId>) -> DispatchResult;
+
 	/// Check if there is any pending slash for the pool.
 	fn pending_slash(pool_account: Pool<Self::AccountId>) -> Self::Balance;
 
@@ -253,12 +258,15 @@ impl<T: Config, Staking: StakingInterface<Balance = BalanceOf<T>, AccountId = T:
 		StakeStrategyType::Transfer
 	}
 
-	fn transferable_balance(pool_account: Pool<Self::AccountId>) -> BalanceOf<T> {
+	fn transferable_balance(
+		pool_account: Pool<Self::AccountId>,
+		_: Member<Self::AccountId>,
+	) -> BalanceOf<T> {
 		T::Currency::balance(&pool_account.0).saturating_sub(Self::active_stake(pool_account))
 	}
 
 	fn total_balance(pool_account: Pool<Self::AccountId>) -> Option<BalanceOf<T>> {
-		Some(T::Currency::total_balance(&pool_account.0))
+		Some(T::Currency::total_balance(&pool_account.get()))
 	}
 
 	fn member_delegation_balance(
@@ -300,6 +308,17 @@ impl<T: Config, Staking: StakingInterface<Balance = BalanceOf<T>, AccountId = T:
 		Ok(())
 	}
 
+	fn dissolve(pool_account: Pool<Self::AccountId>) -> DispatchResult {
+		defensive_assert!(
+			T::Currency::total_balance(&pool_account.clone().get()).is_zero(),
+			"dissolving pool should not have any balance"
+		);
+
+		// Defensively force set balance to zero.
+		T::Currency::set_balance(&pool_account.get(), Zero::zero());
+		Ok(())
+	}
+
 	fn pending_slash(_: Pool<Self::AccountId>) -> Self::Balance {
 		// for transfer stake strategy, slashing is greedy and never deferred.
 		Zero::zero()
@@ -360,11 +379,14 @@ impl<
 		StakeStrategyType::Delegate
 	}
 
-	fn transferable_balance(pool_account: Pool<Self::AccountId>) -> BalanceOf<T> {
-		Delegation::agent_balance(pool_account.clone().into())
+	fn transferable_balance(
+		pool_account: Pool<Self::AccountId>,
+		member_account: Member<Self::AccountId>,
+	) -> BalanceOf<T> {
+		Delegation::agent_transferable_balance(pool_account.clone().into())
 			// pool should always be an agent.
 			.defensive_unwrap_or_default()
-			.saturating_sub(Self::active_stake(pool_account))
+			.min(Delegation::delegator_balance(member_account.into()).unwrap_or_default())
 	}
 
 	fn total_balance(pool_account: Pool<Self::AccountId>) -> Option<BalanceOf<T>> {
@@ -384,12 +406,13 @@ impl<
 	) -> DispatchResult {
 		match bond_type {
 			BondType::Create => {
-				// first delegation
-				Delegation::delegate(who.into(), pool_account.into(), reward_account, amount)
+				// first delegation. Register agent first.
+				Delegation::register_agent(pool_account.clone().into(), reward_account)?;
+				Delegation::delegate(who.into(), pool_account.into(), amount)
 			},
 			BondType::Extra => {
 				// additional delegation
-				Delegation::delegate_extra(who.into(), pool_account.into(), amount)
+				Delegation::delegate(who.into(), pool_account.into(), amount)
 			},
 		}
 	}
@@ -403,6 +426,10 @@ impl<
 		Delegation::withdraw_delegation(who.into(), pool_account.into(), amount, num_slashing_spans)
 	}
 
+	fn dissolve(pool_account: Pool<Self::AccountId>) -> DispatchResult {
+		Delegation::remove_agent(pool_account.into())
+	}
+
 	fn pending_slash(pool_account: Pool<Self::AccountId>) -> Self::Balance {
 		Delegation::pending_slash(pool_account.into()).defensive_unwrap_or_default()
 	}
@@ -433,6 +460,6 @@ impl<
 
 	#[cfg(feature = "runtime-benchmarks")]
 	fn remove_as_agent(pool: Pool<Self::AccountId>) {
-		Delegation::drop_agent(pool.into())
+		Delegation::migrate_to_direct_staker(pool.into())
 	}
 }

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

@@ -1831,7 +1831,9 @@ pub mod pallet {
 		/// A member has been removed from a pool.
 		///
 		/// The removal can be voluntary (withdrawn all unbonded funds) or involuntary (kicked).
-		MemberRemoved { pool_id: PoolId, member: T::AccountId },
+		/// Any funds that was still delegated (dangling delegation) is cleared and is represented
+		/// by `balance`.
+		MemberRemoved { pool_id: PoolId, member: T::AccountId, balance: BalanceOf<T> },
 		/// The roles of a pool have been updated to the given new roles. Note that the depositor
 		/// can never change.
 		RolesUpdated {
@@ -2342,9 +2344,10 @@ pub mod pallet {
 				// don't exist. This check is also defensive in cases where the unbond pool does not
 				// update its balance (e.g. a bug in the slashing hook.) We gracefully proceed in
 				// order to ensure members can leave the pool and it can be destroyed.
-				.min(T::StakeAdapter::transferable_balance(Pool::from(
-					bonded_pool.bonded_account(),
-				)));
+				.min(T::StakeAdapter::transferable_balance(
+					Pool::from(bonded_pool.bonded_account()),
+					Member::from(member_account.clone()),
+				));
 
 			// this can fail if the pool uses `DelegateStake` strategy and the member delegation
 			// is not claimed yet. See `Call::migrate_delegation()`.
@@ -2368,9 +2371,27 @@ pub mod pallet {
 
 				// member being reaped.
 				PoolMembers::<T>::remove(&member_account);
+
+				// Ensure any dangling delegation is withdrawn.
+				let dangling_withdrawal = match T::StakeAdapter::member_delegation_balance(
+					Member::from(member_account.clone()),
+				) {
+					Some(dangling_delegation) => {
+						T::StakeAdapter::member_withdraw(
+							Member::from(member_account.clone()),
+							Pool::from(bonded_pool.bonded_account()),
+							dangling_delegation,
+							num_slashing_spans,
+						)?;
+						dangling_delegation
+					},
+					None => Zero::zero(),
+				};
+
 				Self::deposit_event(Event::<T>::MemberRemoved {
 					pool_id: member.pool_id,
 					member: member_account.clone(),
+					balance: dangling_withdrawal,
 				});
 
 				if member_account == bonded_pool.roles.depositor {
@@ -3078,16 +3099,11 @@ impl<T: Config> Pallet<T> {
 			T::Currency::total_balance(&reward_account) == Zero::zero(),
 			"could not transfer all amount to depositor while dissolving pool"
 		);
-		defensive_assert!(
-			T::StakeAdapter::total_balance(Pool::from(bonded_pool.bonded_account()))
-				.unwrap_or_default() ==
-				Zero::zero(),
-			"dissolving pool should not have any balance"
-		);
 		// NOTE: Defensively force set balance to zero.
 		T::Currency::set_balance(&reward_account, Zero::zero());
-		// NOTE: With `DelegateStake` strategy, this won't do anything.
-		T::Currency::set_balance(&bonded_pool.bonded_account(), Zero::zero());
+
+		// dissolve pool account.
+		let _ = T::StakeAdapter::dissolve(Pool::from(bonded_account)).defensive();
 
 		Self::deposit_event(Event::<T>::Destroyed { pool_id: bonded_pool.id });
 		// Remove bonded pool metadata.
@@ -3525,13 +3541,8 @@ impl<T: Config> Pallet<T> {
 		// this is their balance in the pool
 		let expected_balance = pool_member.total_balance();
 
-		defensive_assert!(
-			actual_balance >= expected_balance,
-			"actual balance should always be greater or equal to the expected"
-		);
-
 		// return the amount to be slashed.
-		Ok(actual_balance.defensive_saturating_sub(expected_balance))
+		Ok(actual_balance.saturating_sub(expected_balance))
 	}
 
 	/// Apply freeze on reward account to restrict it from going below ED.