diff --git a/rewards/src/lib.rs b/rewards/src/lib.rs index 44dc30538..48053b3bf 100644 --- a/rewards/src/lib.rs +++ b/rewards/src/lib.rs @@ -85,12 +85,14 @@ decl_module! { pub struct Module for enum Call where origin: T::Origin { fn on_initialize(now: T::BlockNumber) -> Weight { + let mut count = 0; T::Handler::accumulate_reward(now, | pool, reward_to_accumulate | { if !reward_to_accumulate.is_zero() { + count += 1; Pools::::mutate(pool, | pool_info | pool_info.total_rewards = pool_info.total_rewards.saturating_add(reward_to_accumulate)); } }); - T::WeightInfo::on_initialize(Pools::::iter().count() as u32) + T::WeightInfo::on_initialize(count) } } } diff --git a/rewards/src/mock.rs b/rewards/src/mock.rs index 6be5c5692..104732134 100644 --- a/rewards/src/mock.rs +++ b/rewards/src/mock.rs @@ -79,7 +79,7 @@ impl RewardHandler for Handler { fn accumulate_reward( now: BlockNumber, - callback: impl Fn(Self::PoolId, Self::Balance), + mut callback: impl FnMut(Self::PoolId, Self::Balance), ) -> Vec<(Self::CurrencyId, Self::Balance)> { if now % 2 == 0 { let mut total_accumulated_rewards = 0; diff --git a/traits/src/rewards.rs b/traits/src/rewards.rs index 7e545f2a5..e53b905a6 100644 --- a/traits/src/rewards.rs +++ b/traits/src/rewards.rs @@ -19,7 +19,7 @@ pub trait RewardHandler { /// Accumulate rewards fn accumulate_reward( now: BlockNumber, - callback: impl Fn(Self::PoolId, Self::Balance), + callback: impl FnMut(Self::PoolId, Self::Balance), ) -> Vec<(Self::CurrencyId, Self::Balance)>; /// Payout the reward to `who`