From 8a67549e6b0830bd7e86f71ee2388cb8f76b0de2 Mon Sep 17 00:00:00 2001
From: Valentin Fernandez <tinchofernandez8@gmail.com>
Date: Mon, 9 Dec 2024 15:16:11 -0300
Subject: [PATCH] Fix issue 17

---
 src/lib.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 63013bf..a7a7ad6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -323,6 +323,11 @@ pub mod pallet {
 	pub type Candidates<T: Config> =
 		CountedStorageMap<_, Blake2_128Concat, T::AccountId, CandidateInfoOf<T>, OptionQuery>;
 
+	// Map of Candidates that have been removed in the current session.
+	#[pallet::storage]
+	pub type SessionRemovedCandidates<T: Config> =
+		StorageMap<_, Blake2_128Concat, T::AccountId, CandidateInfoOf<T>, OptionQuery>;
+
 	/// Last block authored by a collator.
 	#[pallet::storage]
 	pub type LastAuthoredBlock<T: Config> =
@@ -575,6 +580,8 @@ pub mod pallet {
 		NoStakeOnCandidate,
 		/// No rewards to claim as previous claim happened on the same session.
 		NoPendingClaim,
+		/// Candidate has not been removed in the current session.
+		NotRemovedCandidate,
 	}
 
 	#[pallet::hooks]
@@ -767,7 +774,11 @@ pub mod pallet {
 				Error::<T>::TooFewEligibleCollators
 			);
 			// Do remove their last authored block.
-			Self::try_remove_candidate(&who, true)?;
+			let candidate_info = Self::try_remove_candidate(&who, true)?;
+
+			// Store removed candidate in SessionRemovedCandidates to properly reward
+			// the candidate and its stakers at the end of the session.
+			SessionRemovedCandidates::<T>::insert(&who, candidate_info);
 
 			Ok(())
 		}
@@ -1149,6 +1160,14 @@ pub mod pallet {
 			Candidates::<T>::get(account).ok_or(Error::<T>::NotCandidate.into())
 		}
 
+		/// Checks whether a given `account` is a candidate and returns its position if successful.
+		pub fn take_removed_candidate(
+			account: &T::AccountId,
+		) -> Result<CandidateInfoOf<T>, DispatchError> {
+			SessionRemovedCandidates::<T>::take(account)
+				.ok_or(Error::<T>::NotRemovedCandidate.into())
+		}
+
 		/// Checks whether a given `account` is an invulnerable.
 		pub fn is_invulnerable(account: &T::AccountId) -> bool {
 			Invulnerables::<T>::get().binary_search(account).is_ok()
@@ -1321,6 +1340,11 @@ pub mod pallet {
 					}
 					let info = CandidateInfo { stake, stakers };
 					*maybe_candidate_info = Some(info.clone());
+
+					// If the candidate left in the current session and is now rejoining
+					// remove it from the SessionRemovedCandidates
+					SessionRemovedCandidates::<T>::remove(&who);
+
 					T::Currency::set_freeze(&FreezeReason::CandidacyBond.into(), who, bond)?;
 					Ok(info)
 				},
@@ -1700,7 +1724,14 @@ pub mod pallet {
 			if !rewardable_blocks.is_zero() && !total_rewards.is_zero() {
 				let collator_percentage = CollatorRewardPercentage::<T>::get();
 				for (collator, blocks) in ProducedBlocks::<T>::drain() {
-					if let Ok(collator_info) = Self::get_candidate(&collator) {
+					// Get the collator info of a candidate, in the case that the collator was removed from the
+					// candidate list during the session, the collator and its stakers must still be rewarded
+					// for the produced blocks in the session so the info can be obtained from SessionRemovedCandidates.
+					let info = Self::get_candidate(&collator)
+						.or_else(|_| Self::take_removed_candidate(&collator))
+						.ok();
+
+					if let Some(collator_info) = info {
 						if blocks > rewardable_blocks {
 							// The only case this could happen is if the candidate was an invulnerable during the session.
 							// Since blocks produced by invulnerables are not currently stored in ProducedBlocks this error
@@ -1891,7 +1922,12 @@ pub mod pallet {
 		) -> Result<T::AccountId, DispatchError> {
 			let (candidate, worst_bond) = Self::get_worst_candidate()?;
 			ensure!(bond > worst_bond, Error::<T>::InvalidCandidacyBond);
-			Self::try_remove_candidate(&candidate, false)?;
+			let candidate_info = Self::try_remove_candidate(&candidate, false)?;
+
+			// Store removed candidate in SessionRemovedCandidates to properly reward
+			// the candidate and its stakers at the end of the session.
+			SessionRemovedCandidates::<T>::insert(&candidate, candidate_info);
+
 			Ok(candidate)
 		}
 
@@ -2001,6 +2037,12 @@ pub mod pallet {
 			let removed = candidates_len_before.saturating_sub(active_candidates_count);
 			let result = Self::assemble_collators();
 
+			// Although the removed candidates are passively deleted from SessionRemovedCandidates
+			// during the distribution of session rewards, it is possible that a removed candidate
+			// is not removed if the candidate didn't produce and blocks during the session. For that
+			// reason the leftover keys in the SessionRemovedCandidates StorageMap must be cleared.
+			let _ = SessionRemovedCandidates::<T>::clear(T::MaxCandidates::get(), None);
+
 			frame_system::Pallet::<T>::register_extra_weight_unchecked(
 				T::WeightInfo::new_session(removed, candidates_len_before),
 				DispatchClass::Mandatory,