Skip to content

Commit

Permalink
fix: lst unbond stores the unbond pool (#822)
Browse files Browse the repository at this point in the history 
* add new function to balances precompile

* fix consesus data provider

* cleanup

* cleanup

* clippy

* add tests for multiple pool unbonding
  • Loading branch information
@1xstj
1xstj authored Nov 19, 2024
1 parent 462aa84 commit b535f8c
Show file tree
Hide file tree
Showing 3 changed files with 216 additions and 14 deletions.
26 changes: 17 additions & 9 deletions pallets/tangle-lst/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -758,13 +758,11 @@ pub mod pallet {
UnbondingMembers::<T>::try_mutate(
member_account.clone(),
|member| -> DispatchResult {
let member = member.get_or_insert_with(|| PoolMember {
pool_id,
unbonding_eras: Default::default(),
});
let member = member
.get_or_insert_with(|| PoolMember { unbonding_eras: Default::default() });
member
.unbonding_eras
.try_insert(unbond_era, points_unbonded)
.try_insert(unbond_era, (pool_id, points_unbonded))
.map(|old| {
if old.is_some() {
defensive!("value checked to not exist in the map; qed");
Expand Down Expand Up @@ -847,10 +845,20 @@ pub mod pallet {
.ok_or(Error::<T>::PoolMemberNotFound)?;
let current_era = T::Staking::current_era();

let bonded_pool = BondedPool::<T>::get(member.pool_id)
// get the pool id from the unbonding map
let pool_id = member.get_by_pool_id(current_era, pool_id);

if pool_id.is_none() {
return Err(Error::<T>::PoolNotFound.into());
}

// checked above
let pool_id = pool_id.unwrap();

let bonded_pool = BondedPool::<T>::get(pool_id)
.defensive_ok_or::<Error<T>>(DefensiveError::PoolNotFound.into())?;
let mut sub_pools =
SubPoolsStorage::<T>::get(member.pool_id).ok_or(Error::<T>::SubPoolsNotFound)?;
SubPoolsStorage::<T>::get(pool_id).ok_or(Error::<T>::SubPoolsNotFound)?;

bonded_pool.ok_to_withdraw_unbonded_with(&caller, &member_account)?;

Expand Down Expand Up @@ -905,7 +913,7 @@ pub mod pallet {

Self::deposit_event(Event::<T>::Withdrawn {
member: member_account.clone(),
pool_id: member.pool_id,
pool_id,
points: sum_unlocked_points,
balance: balance_to_unbond,
});
Expand All @@ -922,7 +930,7 @@ pub mod pallet {
Pallet::<T>::dissolve_pool(bonded_pool);
None
} else {
SubPoolsStorage::<T>::insert(member.pool_id, sub_pools);
SubPoolsStorage::<T>::insert(pool_id, sub_pools);
Some(T::WeightInfo::withdraw_unbonded_update(num_slashing_spans))
}
} else {
Expand Down
190 changes: 190 additions & 0 deletions pallets/tangle-lst/src/tests/unbond.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -823,3 +823,193 @@ fn depositor_permissioned_partial_unbond_slashed() {
);
});
}

#[test]
fn multi_pool_unbonding_works() {
ExtBuilder::default()
.add_members(vec![(40, 40)])
.build_and_execute(|| {
// Create a second pool
assert_ok!(Lst::create(
RuntimeOrigin::signed(20),
20,
900,
901,
902,
Default::default(),
Default::default()
));

// Join the second pool
assert_ok!(Lst::join(RuntimeOrigin::signed(40), 40, 2));

// Unbond from both pools
assert_ok!(Lst::unbond(RuntimeOrigin::signed(40), 40, 1, 20));
assert_ok!(Lst::unbond(RuntimeOrigin::signed(40), 40, 2, 30));

// Check that unbonding entries are correctly tracked with pool IDs
let member = UnbondingMembers::<Runtime>::get(40).unwrap();
let unbonding_eras: Vec<_> = member.unbonding_eras.iter().collect();

assert_eq!(unbonding_eras.len(), 2);
assert_eq!(unbonding_eras[0].1.0, 1); // First entry should be from pool 1
assert_eq!(unbonding_eras[0].1.1, 20); // With 20 points
assert_eq!(unbonding_eras[1].1.0, 2); // Second entry should be from pool 2
assert_eq!(unbonding_eras[1].1.1, 30); // With 30 points

// Advance era and try to withdraw
CurrentEra::set(3);

// Withdraw from pool 1
assert_ok!(Lst::withdraw_unbonded(RuntimeOrigin::signed(40), 40, 1, 0));

// Check that only pool 1's unbonding was withdrawn
let member = UnbondingMembers::<Runtime>::get(40).unwrap();
let remaining_unbonding: Vec<_> = member.unbonding_eras.iter().collect();
assert_eq!(remaining_unbonding.len(), 1);
assert_eq!(remaining_unbonding[0].1.0, 2); // Only pool 2 entry should remain
assert_eq!(remaining_unbonding[0].1.1, 30);

// Withdraw from pool 2
assert_ok!(Lst::withdraw_unbonded(RuntimeOrigin::signed(40), 40, 2, 0));

// Check that all unbonding entries are cleared
let member = UnbondingMembers::<Runtime>::get(40);
assert!(member.is_none());

assert_eq!(
pool_events_since_last_call(),
vec![
Event::Created { depositor: 10, pool_id: 1 },
Event::Bonded { member: 10, pool_id: 1, bonded: 10, joined: true },
Event::Bonded { member: 40, pool_id: 1, bonded: 40, joined: true },
Event::Created { depositor: 20, pool_id: 2 },
Event::Bonded { member: 20, pool_id: 2, bonded: 20, joined: true },
Event::Bonded { member: 40, pool_id: 2, bonded: 40, joined: true },
Event::Unbonded { member: 40, pool_id: 1, points: 20, balance: 20, era: 3 },
Event::Unbonded { member: 40, pool_id: 2, points: 30, balance: 30, era: 3 },
Event::Withdrawn { member: 40, pool_id: 1, points: 20, balance: 20 },
Event::Withdrawn { member: 40, pool_id: 2, points: 30, balance: 30 },
Event::MemberRemoved { pool_id: 1, member: 40 },
Event::MemberRemoved { pool_id: 2, member: 40 }
]
);
});
}

#[test]
fn multi_pool_unbonding_with_slashing() {
ExtBuilder::default()
.add_members(vec![(40, 40)])
.build_and_execute(|| {
// Create a second pool
assert_ok!(Lst::create(
RuntimeOrigin::signed(20),
20,
900,
901,
902,
Default::default(),
Default::default()
));

// Join the second pool
assert_ok!(Lst::join(RuntimeOrigin::signed(40), 40, 2));

// Unbond from both pools
assert_ok!(Lst::unbond(RuntimeOrigin::signed(40), 40, 1, 20));
assert_ok!(Lst::unbond(RuntimeOrigin::signed(40), 40, 2, 30));

// Slash pool 1
StakingMock::slash_by(1, 10);

// Advance era
CurrentEra::set(3);

// Withdraw from both pools
assert_ok!(Lst::withdraw_unbonded(RuntimeOrigin::signed(40), 40, 1, 0));
assert_ok!(Lst::withdraw_unbonded(RuntimeOrigin::signed(40), 40, 2, 0));

assert_eq!(
pool_events_since_last_call(),
vec![
Event::Created { depositor: 10, pool_id: 1 },
Event::Bonded { member: 10, pool_id: 1, bonded: 10, joined: true },
Event::Bonded { member: 40, pool_id: 1, bonded: 40, joined: true },
Event::Created { depositor: 20, pool_id: 2 },
Event::Bonded { member: 20, pool_id: 2, bonded: 20, joined: true },
Event::Bonded { member: 40, pool_id: 2, bonded: 40, joined: true },
Event::Unbonded { member: 40, pool_id: 1, points: 20, balance: 20, era: 3 },
Event::Unbonded { member: 40, pool_id: 2, points: 30, balance: 30, era: 3 },
Event::PoolSlashed { pool_id: 1, balance: 10 },
Event::Withdrawn { member: 40, pool_id: 1, points: 20, balance: 15 }, // Slashed amount
Event::Withdrawn { member: 40, pool_id: 2, points: 30, balance: 30 }, // Unaffected by slash
Event::MemberRemoved { pool_id: 1, member: 40 },
Event::MemberRemoved { pool_id: 2, member: 40 }
]
);
});
}

#[test]
fn multi_pool_unbonding_with_destroying_pool() {
ExtBuilder::default()
.add_members(vec![(40, 40)])
.build_and_execute(|| {
// Create a second pool
assert_ok!(Lst::create(
RuntimeOrigin::signed(20),
20,
900,
901,
902,
Default::default(),
Default::default()
));

// Join the second pool
assert_ok!(Lst::join(RuntimeOrigin::signed(40), 40, 2));

// Set pool 1 to destroying
unsafe_set_state(1, PoolState::Destroying);

// Unbond from both pools
assert_ok!(Lst::unbond(RuntimeOrigin::signed(40), 40, 1, 40)); // Full unbond from destroying pool
assert_ok!(Lst::unbond(RuntimeOrigin::signed(40), 40, 2, 30)); // Partial unbond from active pool

// Check that unbonding entries are correctly tracked
let member = UnbondingMembers::<Runtime>::get(40).unwrap();
let unbonding_eras: Vec<_> = member.unbonding_eras.iter().collect();

assert_eq!(unbonding_eras.len(), 2);
assert_eq!(unbonding_eras[0].1.0, 1);
assert_eq!(unbonding_eras[0].1.1, 40); // Full amount from pool 1
assert_eq!(unbonding_eras[1].1.0, 2);
assert_eq!(unbonding_eras[1].1.1, 30); // Partial amount from pool 2

// Advance era and withdraw
CurrentEra::set(3);
assert_ok!(Lst::withdraw_unbonded(RuntimeOrigin::signed(40), 40, 1, 0));
assert_ok!(Lst::withdraw_unbonded(RuntimeOrigin::signed(40), 40, 2, 0));

// Check that member is removed from pool 1 but still exists in pool 2
assert!(PoolMembers::<Runtime>::contains_key(40));

assert_eq!(
pool_events_since_last_call(),
vec![
Event::Created { depositor: 10, pool_id: 1 },
Event::Bonded { member: 10, pool_id: 1, bonded: 10, joined: true },
Event::Bonded { member: 40, pool_id: 1, bonded: 40, joined: true },
Event::Created { depositor: 20, pool_id: 2 },
Event::Bonded { member: 20, pool_id: 2, bonded: 20, joined: true },
Event::Bonded { member: 40, pool_id: 2, bonded: 40, joined: true },
Event::Unbonded { member: 40, pool_id: 1, points: 40, balance: 40, era: 3 },
Event::Unbonded { member: 40, pool_id: 2, points: 30, balance: 30, era: 3 },
Event::Withdrawn { member: 40, pool_id: 1, points: 40, balance: 40 },
Event::MemberRemoved { pool_id: 1, member: 40 },
Event::Withdrawn { member: 40, pool_id: 2, points: 30, balance: 30 }
]
);
});
}
14 changes: 9 additions & 5 deletions pallets/tangle-lst/src/types/pools.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,9 @@ use super::*;
#[codec(mel_bound(T: Config))]
#[scale_info(skip_type_params(T))]
pub struct PoolMember<T: Config> {
/// The identifier of the pool to which `who` belongs.
pub pool_id: PoolId,
/// The eras in which this member is unbonding, mapped from era index to the number of
/// points scheduled to unbond in the given era.
pub unbonding_eras: BoundedBTreeMap<EraIndex, BalanceOf<T>, T::MaxUnbonding>,
pub unbonding_eras: BoundedBTreeMap<EraIndex, (PoolId, BalanceOf<T>), T::MaxUnbonding>,
}

impl<T: Config> PoolMember<T> {
Expand All @@ -26,7 +24,7 @@ impl<T: Config> PoolMember<T> {
self.unbonding_eras
.as_ref()
.iter()
.fold(BalanceOf::<T>::zero(), |acc, (_, v)| acc.saturating_add(*v))
.fold(BalanceOf::<T>::zero(), |acc, (_, v)| acc.saturating_add(v.1))
}

/// Withdraw any funds in [`Self::unbonding_eras`] who's deadline in reached and is fully
Expand All @@ -47,13 +45,19 @@ impl<T: Config> PoolMember<T> {
true
} else {
removed_points
.try_insert(*e, *p)
.try_insert(*e, p.1)
.expect("source map is bounded, this is a subset, will be bounded; qed");
false
}
});
removed_points
}

pub fn get_by_pool_id(&self, current_era: EraIndex, pool_id: PoolId) -> Option<PoolId> {
self.unbonding_eras
.get(&current_era)
.and_then(|(p, _b)| if *p == pool_id { Some(*p) } else { None })
}
}

/// A pool's possible states.
Expand Down

0 comments on commit b535f8c

Please sign in to comment.