Fix notify inactive collators failures at the end of a round (#3128)

* fix: 🐛 call ParachainStaking::on_finalize on mock blocks rolling

* fix: 🐛 pass correct block number to ParachainStaking::on_finalize

* test: ✅ fix happy path test for notify inactive collator

* revert: 🔥 remove debug code

* feat: ✨ add round length-aware block rolling functions

* test: ✅ increase MaxOfflineRounds to 2 to match RewardPaymentDelay

* test: ✅ call ParachainStaking::on_finalize before Balances and System on_finalize

* test: ✅ add test reproducing a bug in notify_inactive_collator

* refactor: 🔥 remove unused use statement

* test: ✅ minimize changes to mocking framework for ParachainStaking

* docs: 📝 improve comment

* fix: 🐛 add historical staking information

* refactor: ♻️ make roll_to_round functions round length-aware

* test: ✨ improve mock fidelity

* refactor: 🎨 use POINTS_PER_ROUND const in set_author invocations

* docs: 📝 restore comments in tests

* perf: ⚡ use a WasInactive storage to keep track of inactive collators

* test: ✅ fix benchmark test

* fix: 🐛 update benchmarks

* fix: 🐛 correctly update weight calculation on mark_collators_as_inactive invocation

* refactor: ⚡ use a more realistic upper bound to the number of collators in benchmarks

* chore: 🔧 compute new weights

* test: ✅ update test with new weights for operations on initialize

* perf: ⚡ improve estimation of mark_collators_as_inactive weights

* test: ✅ update test test_on_initialize_weights

pallets/parachain-staking/src/benchmarks.rs

@@ -2275,7 +2275,7 @@ benchmarks! {
 	}
 
 	notify_inactive_collator {
-		use crate::{AtStake, CollatorSnapshot, AwardedPts};
+		use crate::{WasInactive};
 
 		// Blocks per-round must be greater than TotalSelected
 		Pallet::<T>::set_blocks_per_round(RawOrigin::Root.into(), 101u32)?;
@@ -2325,11 +2325,8 @@ benchmarks! {
 
 		// Manually change these values for inactive_collator,
 		// so that it can be marked as inactive.
-		<AtStake<T>>::insert(1, &inactive_collator, CollatorSnapshot::default());
-		<AwardedPts<T>>::insert(1, &inactive_collator, 0);
-
-		<AtStake<T>>::insert(2, &inactive_collator, CollatorSnapshot::default());
-		<AwardedPts<T>>::insert(2, &inactive_collator, 0);
+		<WasInactive<T>>::insert(1, &inactive_collator, ());
+		<WasInactive<T>>::insert(2, &inactive_collator, ());
 
 		// Enable killswitch
 		<EnableMarkingOffline<T>>::set(true);
@@ -2338,6 +2335,38 @@ benchmarks! {
 	verify {
 		assert!(!Pallet::<T>::candidate_info(&inactive_collator).expect("must exist").is_active());
 	}
+
+	mark_collators_as_inactive {
+		let x in 0..50; // num collators
+
+		// must come after 'let foo in 0..` statements for macro
+		use crate::{AtStake, CollatorSnapshot, AwardedPts};
+
+		let round = 2;
+		let prev = round - 1;
+
+
+
+		for i in 0..x {
+			let collator = create_funded_collator::<T>(
+				"collator",
+				USER_SEED + i,
+				min_candidate_stk::<T>() * 1_000_000u32.into(),
+				true,
+				999999,
+			)?;
+
+			// All collators were inactinve in previous round
+			<AtStake<T>>::insert(prev, &collator, CollatorSnapshot::default());
+			<AwardedPts<T>>::insert(prev, &collator, 0);
+		}
+
+	}: {
+		let cur = 2;
+		let inactive_info = Pallet::<T>::mark_collators_as_inactive(cur);
+	}
+	verify {
+	}
 }
 
 #[cfg(test)]

pallets/parachain-staking/src/lib.rs

@@ -489,20 +489,23 @@ pub mod pallet {
 					selected_collators_number: collator_count,
 					total_balance: total_staked,
 				});
+				// record inactive collators
+				weight = weight.saturating_add(Self::mark_collators_as_inactive(round.current));
 				// account for Round write
 				weight = weight.saturating_add(T::DbWeight::get().reads_writes(0, 1));
 			} else {
 				weight = weight.saturating_add(Self::handle_delayed_payouts(round.current));
 			}
 
 			// add on_finalize weight
-			//   read:  Author, Points, AwardedPts
-			//   write: Points, AwardedPts
-			weight = weight.saturating_add(T::DbWeight::get().reads_writes(3, 2));
+			//   read:  Author, Points, AwardedPts, WasInactive
+			//   write: Points, AwardedPts, WasInactive
+			weight = weight.saturating_add(T::DbWeight::get().reads_writes(4, 3));
 			weight
 		}
 		fn on_finalize(_n: BlockNumberFor<T>) {
 			Self::award_points_to_block_author();
+			Self::cleanup_inactive_collator_info();
 		}
 	}
 
@@ -643,6 +646,13 @@ pub mod pallet {
 		OptionQuery,
 	>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn was_inactive)]
+	/// Records collators' inactivity.
+	/// Data persists for MaxOfflineRounds + 1 rounds before being pruned.
+	pub type WasInactive<T: Config> =
+		StorageDoubleMap<_, Twox64Concat, RoundIndex, Twox64Concat, T::AccountId, (), OptionQuery>;
+
 	#[pallet::storage]
 	#[pallet::getter(fn delayed_payouts)]
 	/// Delayed payouts
@@ -1417,17 +1427,9 @@ pub mod pallet {
 			// the collator should be notified as inactive
 			let mut inactive_counter: RoundIndex = 0u32;
 
-			// Iter rounds to check
-			//
-			// - The collator has AtStake associated and their AwardedPts are zero
-			//
-			// If the previous condition is met in all rounds of rounds_to_check,
-			// the collator is notified as inactive
+			// Iter rounds and check whether the collator has been inactive
 			for r in rounds_to_check {
-				let stake = <AtStake<T>>::get(r, &collator);
-				let pts = <AwardedPts<T>>::get(r, &collator);
-
-				if stake.is_some() && pts.is_zero() {
+				if <WasInactive<T>>::get(r, &collator).is_some() {
 					inactive_counter = inactive_counter.saturating_add(1);
 				}
 			}
@@ -1704,8 +1706,8 @@ pub mod pallet {
 			let return_stake = |bond: Bond<T::AccountId, BalanceOf<T>>| {
 				// remove delegation from delegator state
 				let mut delegator = DelegatorState::<T>::get(&bond.owner).expect(
-					"Collator state and delegator state are consistent. 
-						Collator state has a record of this delegation. Therefore, 
+					"Collator state and delegator state are consistent.
+						Collator state has a record of this delegation. Therefore,
 						Delegator state also has a record. qed.",
 				);
 
@@ -2333,18 +2335,48 @@ pub mod pallet {
 				});
 			};
 		}
-	}
 
-	/// Add reward points to block authors:
-	/// * 20 points to the block producer for producing a block in the chain
-	impl<T: Config> Pallet<T> {
+		/// Add reward points to block authors:
+		/// * 20 points to the block producer for producing a block in the chain
 		fn award_points_to_block_author() {
 			let author = T::BlockAuthor::get();
 			let now = <Round<T>>::get().current;
 			let score_plus_20 = <AwardedPts<T>>::get(now, &author).saturating_add(20);
 			<AwardedPts<T>>::insert(now, author, score_plus_20);
 			<Points<T>>::mutate(now, |x| *x = x.saturating_add(20));
 		}
+
+		/// Marks collators as inactive for the previous round if they received zero awarded points.
+		pub fn mark_collators_as_inactive(cur: RoundIndex) -> Weight {
+			// This function is called after round index increment,
+			// We don't need to saturate here because the genesis round is 1.
+			let prev = cur - 1;
+
+			let mut collators_at_stake_count = 0u32;
+			for (account, _) in <AtStake<T>>::iter_prefix(prev) {
+				collators_at_stake_count = collators_at_stake_count.saturating_add(1u32);
+				if <AwardedPts<T>>::get(prev, &account).is_zero() {
+					<WasInactive<T>>::insert(prev, account, ());
+				}
+			}
+
+			<T as Config>::WeightInfo::mark_collators_as_inactive(collators_at_stake_count)
+		}
+
+		/// Cleans up historical staking information that is older than MaxOfflineRounds
+		/// by removing entries from the WasIactive storage map.
+		fn cleanup_inactive_collator_info() {
+			let now = <Round<T>>::get().current;
+			let minimum_rounds_required = T::MaxOfflineRounds::get() + 1;
+
+			if now < minimum_rounds_required {
+				return;
+			}
+
+			let _ = <WasInactive<T>>::iter_prefix(now - minimum_rounds_required)
+				.drain()
+				.next();
+		}
 	}
 
 	impl<T: Config> nimbus_primitives::CanAuthor<T::AccountId> for Pallet<T> {

pallets/parachain-staking/src/mock.rs

@@ -116,9 +116,13 @@ const GENESIS_BLOCKS_PER_ROUND: BlockNumber = 5;
 const GENESIS_COLLATOR_COMMISSION: Perbill = Perbill::from_percent(20);
 const GENESIS_PARACHAIN_BOND_RESERVE_PERCENT: Percent = Percent::from_percent(30);
 const GENESIS_NUM_SELECTED_CANDIDATES: u32 = 5;
+
+pub const POINTS_PER_BLOCK: u32 = 20;
+pub const POINTS_PER_ROUND: u32 = GENESIS_BLOCKS_PER_ROUND * POINTS_PER_BLOCK;
+
 parameter_types! {
 	pub const MinBlocksPerRound: u32 = 3;
-	pub const MaxOfflineRounds: u32 = 1;
+	pub const MaxOfflineRounds: u32 = 2;
 	pub const LeaveCandidatesDelay: u32 = 2;
 	pub const CandidateBondLessDelay: u32 = 2;
 	pub const LeaveDelegatorsDelay: u32 = 2;
@@ -278,6 +282,7 @@ impl ExtBuilder {
 
 /// Rolls forward one block. Returns the new block number.
 fn roll_one_block() -> BlockNumber {
+	ParachainStaking::on_finalize(System::block_number());
 	Balances::on_finalize(System::block_number());
 	System::on_finalize(System::block_number());
 	System::set_block_number(System::block_number() + 1);
@@ -312,15 +317,31 @@ pub(crate) fn roll_blocks(num_blocks: u32) -> BlockNumber {
 /// This will complete the block in which the round change occurs.
 /// Returns the number of blocks played.
 pub(crate) fn roll_to_round_begin(round: BlockNumber) -> BlockNumber {
-	let block = (round - 1) * GENESIS_BLOCKS_PER_ROUND;
-	roll_to(block)
+	let r = ParachainStaking::round();
+
+	// Return 0 if target round has already passed
+	if round < r.current + 1 {
+		return 0;
+	}
+
+	// Calculate target block by adding round length for each round difference
+	let target = r.first + (round - r.current) * r.length;
+	roll_to(target)
 }
 
 /// Rolls block-by-block to the end of the specified round.
 /// The block following will be the one in which the specified round change occurs.
 pub(crate) fn roll_to_round_end(round: BlockNumber) -> BlockNumber {
-	let block = round * GENESIS_BLOCKS_PER_ROUND - 1;
-	roll_to(block)
+	let r = ParachainStaking::round();
+
+	// Return 0 if target round has already passed
+	if round < r.current {
+		return 0;
+	}
+
+	// Calculate target block by adding round length for each round difference
+	let target = r.first + (round - r.current + 1) * r.length - 1;
+	roll_to(target)
 }
 
 pub(crate) fn inflation_configs(