Cleanup the counters storage item

Author: José Molina

src/benchmarking.rs

@@ -167,7 +167,9 @@ fn prepare_rewards<T: Config + pallet_session::Config>(
 			)
 			.unwrap_or_else(|e| panic!("Could not stake: {:?}", e));
 			Counters::<T>::mutate(candidate, |counter| {
-				counter.saturating_accrue(FixedU128::saturating_from_rational(amount, amount))
+				counter
+					.value
+					.saturating_accrue(FixedU128::saturating_from_rational(amount, amount));
 			})
 		}
 	}

src/lib.rs

@@ -78,6 +78,8 @@ pub mod weights;
 
 const LOG_TARGET: &str = "runtime::collator-staking";
 
+/// # Pallet Documentation
+///
 #[frame_support::pallet]
 pub mod pallet {
 	use frame_support::{
@@ -412,7 +414,7 @@ pub mod pallet {
 	/// Tracks the different types of locks that can be applied to an account's balance.
 	///
 	/// This struct keeps track of the different amounts that are locked for staking,
-	/// releasing (in process of being unlocked), and held as candidacy bond.
+	/// releasing (in the process of being unlocked), and held as candidacy bond.
 	#[derive(
 		PartialEq,
 		Eq,
@@ -431,6 +433,30 @@ pub mod pallet {
 		pub candidacy_bond: Balance,
 	}
 
+	/// Track reward distribution state for a candidate.
+	///
+	/// It maintains the state needed to calculate and distribute rewards for stakers
+	/// who have delegated to a specific collator.
+	#[derive(
+		PartialEq,
+		Eq,
+		Clone,
+		Encode,
+		Decode,
+		RuntimeDebug,
+		scale_info::TypeInfo,
+		MaxEncodedLen,
+		Default,
+	)]
+	pub struct Counter {
+		/// Accumulated reward ratio per token staked, represented as a fixed-point number.
+		/// Updated when new rewards are added for distribution.
+		pub value: FixedU128,
+		/// Number of stakers that have not yet claimed their rewards.
+		/// Decremented when stakers claim their rewards.
+		pub pending_claims: u32,
+	}
+
 	impl<Balance> LockedBalance<Balance>
 	where
 		Balance: Saturating + Copy,
@@ -564,7 +590,7 @@ pub mod pallet {
 	/// Represents accumulated rewards per token staked on a given collator over time.
 	#[pallet::storage]
 	pub type Counters<T: Config> =
-		StorageMap<_, Blake2_128Concat, T::AccountId, FixedU128, ValueQuery>;
+		StorageMap<_, Blake2_128Concat, T::AccountId, Counter, ValueQuery>;
 
 	/// The storage value `LastRewardedKey` is used to track the last key that was rewarded
 	/// automatically by the system.
@@ -1495,13 +1521,26 @@ pub mod pallet {
 			let current_session = CurrentSession::<T>::get();
 			UserStake::<T>::try_mutate(who, |user_stake_info| {
 				for candidate in &user_stake_info.candidates {
-					let counter = Counters::<T>::get(candidate);
-					CandidateStake::<T>::mutate(candidate, who, |info| {
-						let reward =
-							counter.saturating_sub(info.checkpoint).saturating_mul_int(info.stake);
-						candidate_rewards.push((candidate.clone(), reward));
-						total_rewards.saturating_accrue(reward);
-						info.checkpoint = counter;
+					Counters::<T>::mutate_exists(candidate, |maybe_counter| {
+						CandidateStake::<T>::mutate(candidate, who, |info| {
+							let value = maybe_counter.clone().unwrap_or_default().value;
+							let reward = value
+								.saturating_sub(info.checkpoint)
+								.saturating_mul_int(info.stake);
+							candidate_rewards.push((candidate.clone(), reward));
+							total_rewards.saturating_accrue(reward);
+							info.checkpoint = value;
+						});
+						// Remove the counter if everything has been claimed and the candidate left.
+						if let Some(counter) = maybe_counter {
+							if counter.pending_claims <= 1
+								&& Self::get_candidate(candidate).is_err()
+							{
+								*maybe_counter = None;
+							} else {
+								counter.pending_claims.saturating_dec();
+							}
+						}
 					});
 				}
 				user_stake_info.maybe_last_reward_session = Some(current_session);
@@ -1532,7 +1571,8 @@ pub mod pallet {
 			for candidate in &user_stake_info.candidates {
 				let counter = Counters::<T>::get(candidate);
 				let info = CandidateStake::<T>::get(candidate, who);
-				let reward = counter.saturating_sub(info.checkpoint).saturating_mul_int(info.stake);
+				let reward =
+					counter.value.saturating_sub(info.checkpoint).saturating_mul_int(info.stake);
 				total_rewards.saturating_accrue(reward);
 			}
 			total_rewards
@@ -1700,7 +1740,7 @@ pub mod pallet {
 							candidate_info.stakers.saturating_inc();
 						}
 						candidate_stake_info.stake = final_staker_stake;
-						candidate_stake_info.checkpoint = Counters::<T>::get(candidate);
+						candidate_stake_info.checkpoint = Counters::<T>::get(candidate).value;
 						candidate_info.stake.saturating_accrue(amount);
 						UserStake::<T>::try_mutate(staker, |user_stake_info| -> DispatchResult {
 							// In case the user recently unstaked we cannot allow those funds to be quickly
@@ -1897,7 +1937,13 @@ pub mod pallet {
 					}
 					Self::release_candidacy_bond(who, reason)?;
 
-					// Store removed candidate in SessionRemovedCandidates to properly reward
+					// Remove the counter if all rewards have been claimed.
+					let counter = Counters::<T>::get(who);
+					if counter.pending_claims.is_zero() {
+						Counters::<T>::remove(who);
+					}
+
+					// Store the removed candidate in SessionRemovedCandidates to properly reward
 					// the candidate and its stakers at the end of the session.
 					SessionRemovedCandidates::<T>::insert(who, candidate.clone());
 
@@ -2100,7 +2146,8 @@ pub mod pallet {
 							collator_info.stake,
 						);
 						Counters::<T>::mutate(&collator, |counter| {
-							counter.saturating_accrue(session_ratio)
+							counter.value.saturating_accrue(session_ratio);
+							counter.pending_claims = collator_info.stakers;
 						});
 					} else {
 						log::warn!("Collator {:?} is no longer a candidate", collator);

src/migrations/v2.rs

@@ -19,9 +19,9 @@
 
 use crate::migrations::PALLET_MIGRATIONS_ID;
 use crate::{
-	AutoCompoundSettings, BalanceOf, CandidacyBondReleaseOf, CandidacyBondReleases, CandidateStake,
-	CandidateStakeInfo, Candidates, ClaimableRewards, Config, FreezeReason, Layer, LockedBalances,
-	Pallet, ReleaseQueues, ReleaseRequestOf, WeightInfo,
+	AutoCompoundSettings, BalanceOf, CandidacyBondReleases, CandidateStake, CandidateStakeInfo,
+	Candidates, ClaimableRewards, Config, FreezeReason, Layer, LockedBalances, Pallet,
+	ReleaseQueues, WeightInfo,
 };
 use core::fmt::Debug;
 use frame_support::migrations::{MigrationId, SteppedMigration, SteppedMigrationError};
@@ -31,6 +31,8 @@ use frame_support::weights::WeightMeter;
 use sp_runtime::{FixedU128, Percent, Saturating};
 use sp_std::vec::Vec;
 
+#[cfg(feature = "try-runtime")]
+use crate::{CandidacyBondReleaseOf, ReleaseRequestOf};
 #[cfg(feature = "try-runtime")]
 use sp_runtime::TryRuntimeError;
 #[cfg(feature = "try-runtime")]

src/tests.rs

@@ -17,7 +17,7 @@ use crate as collator_staking;
 use crate::{
 	mock::*, AutoCompoundSettings, CandidacyBondRelease, CandidacyBondReleaseReason,
 	CandidacyBondReleases, CandidateInfo, CandidateStake, CandidateStakeInfo, Candidates,
-	ClaimableRewards, CollatorRewardPercentage, Config, Counters, CurrentSession,
+	ClaimableRewards, CollatorRewardPercentage, Config, Counter, Counters, CurrentSession,
 	DesiredCandidates, Error, Event, ExtraReward, IdentityCollator, Invulnerables,
 	LastAuthoredBlock, Layer, MinCandidacyBond, MinStake, NextSystemOperation, Operation,
 	ProducedBlocks, ReleaseQueues, ReleaseRequest, SessionRemovedCandidates, StakeTarget,
@@ -3453,8 +3453,14 @@ mod collator_rewards {
 				)
 			}));
 			assert_eq!(CollatorStaking::calculate_unclaimed_rewards(&4), 15);
+			assert_eq!(Counters::<Test>::get(4).pending_claims, 1);
 			assert_ok!(CollatorStaking::do_claim_rewards(&4));
-			assert_eq!(CandidateStake::<Test>::get(4, 4).checkpoint, Counters::<Test>::get(4));
+			// The candidate 4 left the candidacy, so the counter got erased.
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: FixedU128::zero(), pending_claims: 0 }
+			);
+			assert!(!Counters::<Test>::contains_key(4));
 			assert_eq!(ClaimableRewards::<Test>::get(), 0);
 			// Now we can see the reward.
 			System::assert_has_event(RuntimeEvent::CollatorStaking(Event::StakingRewardReceived {
@@ -3565,7 +3571,10 @@ mod collator_rewards {
 			// Reward for staker when claiming.
 			assert_eq!(CollatorStaking::calculate_unclaimed_rewards(&4), 28);
 			assert_ok!(CollatorStaking::do_claim_rewards(&4));
-			assert_eq!(CandidateStake::<Test>::get(4, 4).checkpoint, Counters::<Test>::get(4));
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: CandidateStake::<Test>::get(4, 4).checkpoint, pending_claims: 0 }
+			);
 			assert_eq!(ClaimableRewards::<Test>::get(), 0);
 			System::assert_has_event(RuntimeEvent::CollatorStaking(Event::StakingRewardReceived {
 				account: 4,
@@ -3677,7 +3686,10 @@ mod collator_rewards {
 			// Reward for staker.
 			assert_eq!(CollatorStaking::calculate_unclaimed_rewards(&4), 15);
 			assert_ok!(CollatorStaking::do_claim_rewards(&4));
-			assert_eq!(CandidateStake::<Test>::get(4, 4).checkpoint, Counters::<Test>::get(4));
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: CandidateStake::<Test>::get(4, 4).checkpoint, pending_claims: 0 }
+			);
 			assert_eq!(ClaimableRewards::<Test>::get(), 0);
 			System::assert_has_event(RuntimeEvent::CollatorStaking(Event::StakingRewardReceived {
 				account: 4,
@@ -3814,15 +3826,21 @@ mod collator_rewards {
 			}));
 			assert_eq!(CollatorStaking::calculate_unclaimed_rewards(&2), 12);
 			assert_ok!(CollatorStaking::do_claim_rewards(&2));
-			assert_eq!(CandidateStake::<Test>::get(4, 2).checkpoint, Counters::<Test>::get(4));
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: CandidateStake::<Test>::get(4, 2).checkpoint, pending_claims: 1 }
+			);
 			assert_eq!(ClaimableRewards::<Test>::get(), 16); // this remains to staker 3.
 			System::assert_has_event(RuntimeEvent::CollatorStaking(Event::StakingRewardReceived {
 				account: 2,
 				amount: 12,
 			}));
 			assert_eq!(CollatorStaking::calculate_unclaimed_rewards(&3), 15);
 			assert_ok!(CollatorStaking::do_claim_rewards(&3));
-			assert_eq!(CandidateStake::<Test>::get(4, 3).checkpoint, Counters::<Test>::get(4));
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: CandidateStake::<Test>::get(4, 3).checkpoint, pending_claims: 0 }
+			);
 			assert_eq!(ClaimableRewards::<Test>::get(), 1); // rounding issue
 			System::assert_has_event(RuntimeEvent::CollatorStaking(Event::StakingRewardReceived {
 				account: 3,
@@ -3916,7 +3934,10 @@ mod collator_rewards {
 
 			// Check the collator's counter and staker's checkpoint. Both should be zero, as no
 			// rewards were distributed.
-			assert_eq!(Counters::<Test>::get(3), FixedU128::zero());
+			assert_eq!(
+				Counters::<Test>::get(3),
+				Counter { value: FixedU128::zero(), pending_claims: 0 }
+			);
 			assert_eq!(CandidateStake::<Test>::get(3, 4).checkpoint, FixedU128::zero());
 
 			// Skip session 0, as there are no rewards for this session.
@@ -3934,7 +3955,10 @@ mod collator_rewards {
 			// Now we have 10 units of rewards being distributed. 20% goes to the collator, and 80%
 			// goes to stakers, so total 8 for stakers. The collator's counter should be the ratio
 			// between the rewards obtained and the total stake deposited in it.
-			assert_eq!(Counters::<Test>::get(4), FixedU128::from_rational(8, 40));
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: FixedU128::from_rational(8, 40), pending_claims: 1 }
+			);
 
 			// The current checkpoint does not vary, as the staker did not claim the rewards yet.
 			assert_eq!(CandidateStake::<Test>::get(4, 3).checkpoint, FixedU128::zero());
@@ -3962,7 +3986,10 @@ mod collator_rewards {
 			));
 
 			// The checkpoint should be equal to the candidate's current counter.
-			assert_eq!(Counters::<Test>::get(4), FixedU128::from_rational(8, 40));
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: FixedU128::from_rational(8, 40), pending_claims: 0 }
+			);
 			assert_eq!(
 				CandidateStake::<Test>::get(4, 5).checkpoint,
 				FixedU128::from_rational(8, 40)
@@ -4369,7 +4396,12 @@ mod claim_rewards_other {
 
 			// Check the collator's counter and staker's checkpoint. Both should be zero, as no
 			// rewards were distributed.
-			assert_eq!(Counters::<Test>::get(4), FixedU128::zero());
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: FixedU128::zero(), pending_claims: 0 }
+			);
+			// Which implies the counter should not exist yet.
+			assert!(!Counters::<Test>::contains_key(4));
 			assert_eq!(CandidateStake::<Test>::get(4, 3).checkpoint, FixedU128::zero());
 
 			// Skip session 0, as there are no rewards for this session
@@ -4392,7 +4424,10 @@ mod claim_rewards_other {
 			// At this point, rewards from session 1 should have been calculated
 			// Check counter is updated as expected with 20% to collator, 80% to stakers
 			let expected_checkpoint = FixedU128::from_rational(80, 40);
-			assert_eq!(Counters::<Test>::get(4), expected_checkpoint);
+			assert_eq!(
+				Counters::<Test>::get(4),
+				Counter { value: expected_checkpoint, pending_claims: 1 }
+			);
 
 			// The checkpoint should still be zero, as rewards aren't claimed
 			// or distributed yet
@@ -4402,8 +4437,7 @@ mod claim_rewards_other {
 			let weight = Weight::from_parts(u64::MAX, u64::MAX);
 			CollatorStaking::on_idle(21, weight);
 
-			// The checkpoint should now be updated to match the counter as
-			// rewards have been distributed
+			// The checkpoint should now be zero, since all stakers claimed
 			assert_eq!(CandidateStake::<Test>::get(4, 3).checkpoint, expected_checkpoint);
 
 			// The stake should have increased by the reward amount (40 * 80/40 = 80)