Skip to content

Commit

Permalink
Part.3
Browse files Browse the repository at this point in the history
  • Loading branch information
@aurexav
aurexav committed Oct 17, 2024
1 parent 48443c5 commit 7e56ca8
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 95 deletions.
55 changes: 46 additions & 9 deletions pallet/staking/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,7 @@ pub mod pallet {
#[pallet::hooks]
impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
fn on_idle(_: BlockNumberFor<T>, mut remaining_weight: Weight) -> Weight {
Self::idle_allocate_ring_staking_reward(&mut remaining_weight);
Self::idle_unstake(&mut remaining_weight);

remaining_weight
Expand All @@ -199,9 +200,9 @@ pub mod pallet {
pub fn unstake_all_for(origin: OriginFor<T>, who: T::AccountId) -> DispatchResult {
ensure_signed(origin)?;

let l = <Ledgers<T>>::take(&who).ok_or(<Error<T>>::NoRecord)?;
let leger = <Ledgers<T>>::take(&who).ok_or(<Error<T>>::NoRecord)?;

Self::unstake_all_for_inner(who, l)?;
Self::unstake_all_for_inner(who, leger)?;

Ok(())
}
Expand All @@ -216,7 +217,9 @@ pub mod pallet {
) -> DispatchResult {
ensure_signed(origin)?;

Self::allocate_ring_staking_reward_of_inner(who)?;
let amount = <PendingRewards<T>>::take(&who).ok_or(<Error<T>>::NoReward)?;

Self::allocate_ring_staking_reward_of_inner(who, amount)?;

Ok(())
}
Expand Down Expand Up @@ -305,9 +308,10 @@ pub mod pallet {
T::KtonStaking::allocate(None, reward_to_kton_staking);
}

pub(crate) fn allocate_ring_staking_reward_of_inner(who: T::AccountId) -> DispatchResult {
let amount = <PendingRewards<T>>::take(&who).ok_or(<Error<T>>::NoReward)?;

pub(crate) fn allocate_ring_staking_reward_of_inner(
who: T::AccountId,
amount: Balance,
) -> DispatchResult {
T::RingStaking::allocate(Some(who.clone()), amount);

Self::deposit_event(Event::Payout { who, amount });
Expand Down Expand Up @@ -342,6 +346,37 @@ pub mod pallet {
}
}

fn idle_allocate_ring_staking_reward(remaining_weight: &mut Weight) {
// At least 1 read weight is required.
if let Some(rw) = remaining_weight.checked_sub(&T::DbWeight::get().reads(1)) {
*remaining_weight = rw;
} else {
return;
}

#[cfg(test)]
let weight = Weight::zero().add_ref_time(1);
#[cfg(not(test))]
let weight = T::WeightInfo::allocate_ring_staking_reward_of();
let mut reward_to_allocate = Vec::new();

for (who, amount) in <PendingRewards<T>>::iter() {
if let Some(rw) = remaining_weight.checked_sub(&weight) {
*remaining_weight = rw;

reward_to_allocate.push((who, amount));
} else {
break;
}
}

for (who, amount) in reward_to_allocate {
let _ = Self::allocate_ring_staking_reward_of_inner(who.clone(), amount);

<PendingRewards<T>>::remove(&who);
}
}

fn idle_unstake(remaining_weight: &mut Weight) {
// At least 1 read weight is required.
if let Some(rw) = remaining_weight.checked_sub(&T::DbWeight::get().reads(1)) {
Expand All @@ -350,9 +385,9 @@ pub mod pallet {
return;
}

#[cfg(feature = "test")]
#[cfg(test)]
let weight = Weight::zero().add_ref_time(1);
#[cfg(not(feature = "test"))]
#[cfg(not(test))]
let weight = T::WeightInfo::unstake_all_for();
let mut ledgers_to_migrate = Vec::new();

Expand All @@ -367,7 +402,9 @@ pub mod pallet {
}

for (who, l) in ledgers_to_migrate {
let _ = Self::unstake_all_for_inner(who, l);
let _ = Self::unstake_all_for_inner(who.clone(), l);

<Ledgers<T>>::remove(&who);
}
}

Expand Down
17 changes: 7 additions & 10 deletions pallet/staking/src/mock.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ use dc_types::UNIT;
// polkadot-sdk
use frame_support::{
assert_ok, derive_impl,
traits::{OnFinalize, OnInitialize},
traits::{OnFinalize, OnIdle, OnInitialize},
};
use sp_core::H160;
use sp_io::TestExternalities;
Expand Down Expand Up @@ -244,7 +244,7 @@ impl crate::Election<AccountId> for RingStaking {
}
}
impl crate::Reward<AccountId> for RingStaking {
fn distribute(who: Option<AccountId>, amount: Balance) {
fn allocate(who: Option<AccountId>, amount: Balance) {
let Some(who) = who else { return };
let _ = Balances::transfer_keep_alive(
RuntimeOrigin::signed(Treasury::account_id()),
Expand All @@ -255,7 +255,7 @@ impl crate::Reward<AccountId> for RingStaking {
}
pub enum KtonStaking {}
impl crate::Reward<AccountId> for KtonStaking {
fn distribute(_: Option<AccountId>, amount: Balance) {
fn allocate(_: Option<AccountId>, amount: Balance) {
let _ = Balances::transfer_keep_alive(
RuntimeOrigin::signed(TreasuryAcct::get()),
<KtonStakingContract<Runtime>>::get().unwrap(),
Expand Down Expand Up @@ -329,13 +329,9 @@ impl ExtBuilder {
}
.assimilate_storage(&mut storage)
.unwrap();
crate::GenesisConfig::<Runtime> {
collator_count: 3,
collators: (1..=3).map(|i| AccountId(i)).collect(),
..Default::default()
}
.assimilate_storage(&mut storage)
.unwrap();
crate::GenesisConfig::<Runtime> { collator_count: 3, ..Default::default() }
.assimilate_storage(&mut storage)
.unwrap();

let mut ext = TestExternalities::from(storage);

Expand Down Expand Up @@ -368,6 +364,7 @@ pub fn initialize_block(number: BlockNumber) {
}

pub fn finalize_block(number: BlockNumber) {
AllPalletsWithSystem::on_idle(number, Weight::MAX);
AllPalletsWithSystem::on_finalize(number);
}

Expand Down
145 changes: 69 additions & 76 deletions pallet/staking/src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
// darwinia
use crate::{mock::*, *};
// polkadot-sdk
use frame_support::{assert_noop, assert_ok};
use frame_support::{assert_noop, assert_ok, traits::OnIdle};
use sp_runtime::DispatchError;

#[test]
Expand Down Expand Up @@ -77,99 +77,53 @@ fn get_top_collators_should_work() {
}

#[test]
fn collator_caches_should_work() {
fn elect_should_work() {
ExtBuilder.build().execute_with(|| {
assert!(call_on_cache!(<Previous<Runtime>>::get()).unwrap().is_empty());
assert!(call_on_cache!(<Current<Runtime>>::get()).unwrap().is_empty());
assert_eq!(
call_on_cache!(<Next<Runtime>>::get()).unwrap(),
vec![AccountId(1), AccountId(2), AccountId(3)]
);
assert_eq!(
<CacheStates<Runtime>>::get(),
(CacheState::Previous, CacheState::Current, CacheState::Next)
);

Staking::shift_cache_states();

assert!(call_on_cache!(<Previous<Runtime>>::get()).unwrap().is_empty());
assert_eq!(
call_on_cache!(<Current<Runtime>>::get()).unwrap(),
vec![AccountId(1), AccountId(2), AccountId(3)]
);
assert!(call_on_cache!(<Next<Runtime>>::get()).unwrap().is_empty());
assert_eq!(
<CacheStates<Runtime>>::get(),
(CacheState::Next, CacheState::Previous, CacheState::Current)
);

Staking::shift_cache_states();

assert_eq!(
call_on_cache!(<Previous<Runtime>>::get()).unwrap(),
vec![AccountId(1), AccountId(2), AccountId(3)]
);
assert!(call_on_cache!(<Current<Runtime>>::get()).unwrap().is_empty());
assert!(call_on_cache!(<Next<Runtime>>::get()).unwrap().is_empty());
assert_eq!(
<CacheStates<Runtime>>::get(),
(CacheState::Current, CacheState::Next, CacheState::Previous)
);

Staking::shift_cache_states();
NEXT_COLLATOR_ID.with(|v| *v.borrow_mut() = 4);

assert!(call_on_cache!(<Previous<Runtime>>::get()).unwrap().is_empty());
assert!(call_on_cache!(<Current<Runtime>>::get()).unwrap().is_empty());
assert_eq!(
call_on_cache!(<Next<Runtime>>::get()).unwrap(),
vec![AccountId(1), AccountId(2), AccountId(3)]
);
assert_eq!(
<CacheStates<Runtime>>::get(),
(CacheState::Previous, CacheState::Current, CacheState::Next)
<Runtime as Config>::RingStaking::elect(<CollatorCount<Runtime>>::get()).unwrap(),
vec![AccountId(4), AccountId(5), AccountId(6)]
);
});
}

#[test]
fn elect_should_work() {
fn on_idle_allocate_ring_staking_reward_should_work() {
ExtBuilder.build().execute_with(|| {
assert_eq!(
call_on_cache!(<Next<Runtime>>::get()).unwrap(),
vec![AccountId(1), AccountId(2), AccountId(3)]
);
(1..=512).for_each(|i| <PendingRewards<Runtime>>::insert(AccountId(i), 1));

NEXT_COLLATOR_ID.with(|v| *v.borrow_mut() = 4);
new_session();
System::reset_events();
AllPalletsWithSystem::on_idle(0, Weight::zero().add_ref_time(128));
assert_eq!(events().into_iter().filter(|e| matches!(e, Event::Payout { .. })).count(), 128);

assert_eq!(
call_on_cache!(<Next<Runtime>>::get()).unwrap(),
vec![AccountId(4), AccountId(5), AccountId(6)]
);
System::reset_events();
AllPalletsWithSystem::on_idle(0, Weight::MAX);
assert_eq!(events().into_iter().filter(|e| matches!(e, Event::Payout { .. })).count(), 384);
});
}

#[test]
fn auto_payout_should_work() {
fn on_idle_unstake_should_work() {
ExtBuilder.build().execute_with(|| {
Efflux::block(1);

(1..=3).for_each(|i| <PendingRewards<Runtime>>::insert(AccountId(i), i as Balance));
(1..=512).for_each(|i| {
<Ledgers<Runtime>>::insert(
AccountId(i),
Ledger { ring: i as _, deposits: BoundedVec::new() },
)
});

System::reset_events();
Efflux::block(1);
dbg!(<PendingRewards<Runtime>>::iter().collect::<Vec<_>>());
assert_eq!(events(), vec![Event::Payout { who: AccountId(2), amount: 2 }]);
AllPalletsWithSystem::on_idle(0, Weight::zero().add_ref_time(128));
assert_eq!(<Ledgers<Runtime>>::iter().count(), 384);

System::reset_events();
Efflux::block(1);
dbg!(<PendingRewards<Runtime>>::iter().collect::<Vec<_>>());
assert_eq!(events(), vec![Event::Payout { who: AccountId(3), amount: 3 }]);
AllPalletsWithSystem::on_idle(0, Weight::MAX);
assert_eq!(<Ledgers<Runtime>>::iter().count(), 0);

System::reset_events();
Efflux::block(1);
dbg!(<PendingRewards<Runtime>>::iter().collect::<Vec<_>>());
assert_eq!(events(), vec![Event::Payout { who: AccountId(1), amount: 1 }]);
(1..512).for_each(|who| {
assert_eq!(Balances::free_balance(AccountId(who)), 100 + who as Balance);
});
});
}

Expand All @@ -186,14 +140,27 @@ fn unstake_all_for_should_work() {
<Error<Runtime>>::NoRecord
);

<Ledgers<Runtime>>::insert(AccountId(1), Ledger { ring: 1, deposits: Default::default() });
<Ledgers<Runtime>>::insert(AccountId(1), Ledger { ring: 1, deposits: BoundedVec::new() });
assert_ok!(Staking::unstake_all_for(RuntimeOrigin::signed(AccountId(1)), AccountId(1)));
});
}

#[test]
fn payout_for_should_work() {
ExtBuilder.build().execute_with(|| {});
fn allocate_ring_staking_reward_of_should_work() {
ExtBuilder.build().execute_with(|| {
let who = AccountId(1);

assert_noop!(
Staking::allocate_ring_staking_reward_of(RuntimeOrigin::signed(who), who),
<Error<Runtime>>::NoReward
);

<PendingRewards<Runtime>>::insert(who, 1);
System::reset_events();

assert_ok!(Staking::allocate_ring_staking_reward_of(RuntimeOrigin::signed(who), who));
assert_eq!(events(), vec![Event::Payout { who, amount: 1 }]);
});
}

#[test]
Expand All @@ -212,3 +179,29 @@ fn set_collator_count_should_work() {
assert_eq!(Staking::collator_count(), 1);
});
}

#[test]
fn set_ring_staking_contract_should_work() {
ExtBuilder.build().execute_with(|| {
assert_noop!(
Staking::set_ring_staking_contract(RuntimeOrigin::signed(AccountId(1)), AccountId(1)),
DispatchError::BadOrigin
);

assert_ok!(Staking::set_ring_staking_contract(RuntimeOrigin::root(), AccountId(1)));
assert_eq!(Staking::ring_staking_contract(), Some(AccountId(1)));
});
}

#[test]
fn set_kton_staking_contract_should_work() {
ExtBuilder.build().execute_with(|| {
assert_noop!(
Staking::set_kton_staking_contract(RuntimeOrigin::signed(AccountId(1)), AccountId(1)),
DispatchError::BadOrigin
);

assert_ok!(Staking::set_kton_staking_contract(RuntimeOrigin::root(), AccountId(1)));
assert_eq!(Staking::kton_staking_contract(), Some(AccountId(1)));
});
}

0 comments on commit 7e56ca8

Please sign in to comment.