From 06e10a360371a15d93ca0c6d9402c8b9e76a01da Mon Sep 17 00:00:00 2001
From: Bryan Chen <xlchen1291@gmail.com>
Date: Wed, 21 Oct 2020 13:53:59 +1300
Subject: [PATCH 1/2] avoid reiterate all the pools

---
 rewards/src/lib.rs    | 4 +++-
 traits/src/rewards.rs | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

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<T: Trait> 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::<T>::mutate(pool, | pool_info | pool_info.total_rewards = pool_info.total_rewards.saturating_add(reward_to_accumulate));
 				}
 			});
-			T::WeightInfo::on_initialize(Pools::<T>::iter().count() as u32)
+			T::WeightInfo::on_initialize(count)
 		}
 	}
 }
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<AccountId, BlockNumber> {
 	/// 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`

From 58635ba3fc90ec1fd69ff5361cd84196f04d4273 Mon Sep 17 00:00:00 2001
From: Bryan Chen <xlchen1291@gmail.com>
Date: Wed, 21 Oct 2020 14:16:36 +1300
Subject: [PATCH 2/2] fix test

---
 rewards/src/mock.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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<AccountId, BlockNumber> 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;