Fix notify inactive collators failures at the end of a round

pallets/parachain-staking/src/lib.rs

@@ -503,6 +503,7 @@ pub mod pallet {
 		}
 		fn on_finalize(_n: BlockNumberFor<T>) {
 			Self::award_points_to_block_author();
+			Self::cleanup_stake_info();
 		}
 	}
 
@@ -643,6 +644,13 @@ pub mod pallet {
 		OptionQuery,
 	>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn had_stake)]
+	/// Records collator's delegation stake presence at round start.
+	/// Data persists for MaxOfflineRounds + 1 rounds before being pruned.
+	pub type HadStake<T: Config> =
+		StorageDoubleMap<_, Twox64Concat, RoundIndex, Twox64Concat, T::AccountId, (), OptionQuery>;
+
 	#[pallet::storage]
 	#[pallet::getter(fn delayed_payouts)]
 	/// Delayed payouts
@@ -1424,7 +1432,7 @@ pub mod pallet {
 			// If the previous condition is met in all rounds of rounds_to_check,
 			// the collator is notified as inactive
 			for r in rounds_to_check {
-				let stake = <AtStake<T>>::get(r, &collator);
+				let stake = <HadStake<T>>::get(r, &collator);
 				let pts = <AwardedPts<T>>::get(r, &collator);
 
 				if stake.is_some() && pts.is_zero() {
@@ -1704,8 +1712,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.",
 				);
 
@@ -2171,6 +2179,7 @@ pub mod pallet {
 					total: total_counted,
 				};
 				<AtStake<T>>::insert(now, account, snapshot);
+				<HadStake<T>>::insert(now, account, ());
 				Self::deposit_event(Event::CollatorChosen {
 					round: now,
 					collator_account: account.clone(),
@@ -2345,6 +2354,21 @@ pub mod pallet {
 			<AwardedPts<T>>::insert(now, author, score_plus_20);
 			<Points<T>>::mutate(now, |x| *x = x.saturating_add(20));
 		}
+
+		/// Cleans up historical staking information that is older than MaxOfflineRounds
+		/// by removing entries from the HadStake storage map.
+		fn cleanup_stake_info() {
+			let now = <Round<T>>::get().current;
+			let minimum_rounds_required = T::MaxOfflineRounds::get() + 1;
+
+			if now < minimum_rounds_required {
+				return;
+			}
+
+			let _ = <HadStake<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

@@ -118,7 +118,7 @@ const GENESIS_PARACHAIN_BOND_RESERVE_PERCENT: Percent = Percent::from_percent(30
 const GENESIS_NUM_SELECTED_CANDIDATES: u32 = 5;
 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;
@@ -277,7 +277,10 @@ impl ExtBuilder {
 }
 
 /// Rolls forward one block. Returns the new block number.
-fn roll_one_block() -> BlockNumber {
+fn roll_one_block(finalize_parachain_staking: bool) -> BlockNumber {
+	if finalize_parachain_staking {
+		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);
@@ -289,40 +292,69 @@ fn roll_one_block() -> BlockNumber {
 }
 
 /// Rolls to the desired block. Returns the number of blocks played.
-pub(crate) fn roll_to(n: BlockNumber) -> BlockNumber {
+/// Allows to decide whther to finalize blocks with ParachainStaking::on_finalize or not.
+fn roll_to_block(n: BlockNumber, finalize_parachain_staking: bool) -> BlockNumber {
 	let mut num_blocks = 0;
 	let mut block = System::block_number();
 	while block < n {
-		block = roll_one_block();
+		block = roll_one_block(finalize_parachain_staking);
 		num_blocks += 1;
 	}
 	num_blocks
 }
 
+/// Rolls to the desired block. Returns the number of blocks played.
+pub(crate) fn roll_to(n: BlockNumber) -> BlockNumber {
+	roll_to_block(n, false)
+}
+
 /// Rolls desired number of blocks. Returns the final block.
 pub(crate) fn roll_blocks(num_blocks: u32) -> BlockNumber {
 	let mut block = System::block_number();
 	for _ in 0..num_blocks {
-		block = roll_one_block();
+		block = roll_one_block(false);
 	}
 	block
 }
 
 /// Rolls block-by-block to the beginning of the specified round.
 /// This will complete the block in which the round change occurs.
 /// Returns the number of blocks played.
+///
+/// WARNING: this function is NOT aware of changes to [`ParachainStaking::round()`].
 pub(crate) fn roll_to_round_begin(round: BlockNumber) -> BlockNumber {
 	let block = (round - 1) * GENESIS_BLOCKS_PER_ROUND;
 	roll_to(block)
 }
 
 /// 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.
+///
+/// WARNING: this function is NOT aware of changes to [`ParachainStaking::round()`].
 pub(crate) fn roll_to_round_end(round: BlockNumber) -> BlockNumber {
 	let block = round * GENESIS_BLOCKS_PER_ROUND - 1;
 	roll_to(block)
 }
 
+/// Rolls block-by-block to the beginning of the next round.
+/// This will complete the block in which the round change occurs.
+/// Returns the number of blocks played.
+///
+/// NOTE: this function is aware of changes to [`ParachainStaking::round()`].
+pub(crate) fn roll_to_next_round_begin() -> BlockNumber {
+	let block = ParachainStaking::round().first + ParachainStaking::round().length;
+	roll_to_block(block, true)
+}
+
+/// Rolls block-by-block to the end of the current round.
+/// The block following will be the one in which the specified round change occurs.
+///
+/// NOTE: this function is aware of changes to [`ParachainStaking::round()`].
+pub(crate) fn roll_to_current_round_end() -> BlockNumber {
+	let block = ParachainStaking::round().first + ParachainStaking::round().length - 1;
+	roll_to_block(block, true)
+}
+
 pub(crate) fn inflation_configs(
 	pbr: AccountId,
 	pbr_percent: u8,

pallets/parachain-staking/src/tests.rs

@@ -25,14 +25,15 @@
 use crate::auto_compound::{AutoCompoundConfig, AutoCompoundDelegations};
 use crate::delegation_requests::{CancelledScheduledRequest, DelegationAction, ScheduledRequest};
 use crate::mock::{
-	inflation_configs, roll_blocks, roll_to, roll_to_round_begin, roll_to_round_end, set_author,
-	set_block_author, AccountId, Balances, BlockNumber, ExtBuilder, ParachainStaking, RuntimeEvent,
-	RuntimeOrigin, Test,
+	inflation_configs, roll_blocks, roll_to, roll_to_current_round_end, roll_to_next_round_begin,
+	roll_to_round_begin, roll_to_round_end, set_author, set_block_author, AccountId, Balances,
+	BlockNumber, ExtBuilder, ParachainStaking, RuntimeEvent, RuntimeOrigin, Test,
 };
 use crate::{
 	assert_events_emitted, assert_events_emitted_match, assert_events_eq, assert_no_events,
 	AtStake, Bond, CollatorStatus, DelegationScheduledRequests, DelegatorAdded,
-	EnableMarkingOffline, Error, Event, InflationDistributionInfo, Range, DELEGATOR_LOCK_ID,
+	EnableMarkingOffline, Error, Event, HadStake, InflationDistributionInfo, Range,
+	DELEGATOR_LOCK_ID,
 };
 use frame_support::traits::{Currency, ExistenceRequirement, WithdrawReasons};
 use frame_support::{assert_err, assert_noop, assert_ok, pallet_prelude::*, BoundedVec};
@@ -1173,8 +1174,62 @@ fn enable_marking_offline_fails_bad_origin() {
 		});
 }
 
+#[test]
+fn hadstake_is_cleaned_up_after_max_offline_rounds() {
+	const ACTIVE_COLLATOR: AccountId = 1;
+
+	ExtBuilder::default()
+		.with_balances(vec![(1, 20), (2, 20), (3, 20), (4, 20), (5, 20)])
+		.with_candidates(vec![(1, 20), (2, 20), (3, 20), (4, 20), (5, 20)])
+		.build()
+		.execute_with(|| {
+			assert_eq!(<Test as crate::Config>::MaxOfflineRounds::get(), 2);
+			assert_eq!(<Test as crate::Config>::RewardPaymentDelay::get(), 2);
+
+			// ACTIVE_COLLATOR authors all the blocks
+			set_block_author(ACTIVE_COLLATOR);
+
+			// Round 2
+			roll_to_next_round_begin();
+
+			// ACTIVE_COLLATOR has a stake in round 1
+			assert!(<AtStake<Test>>::contains_key(1, ACTIVE_COLLATOR));
+			assert!(<HadStake<Test>>::contains_key(1, ACTIVE_COLLATOR));
+
+			// Round 3
+			roll_to_next_round_begin();
+
+			// ACTIVE_COLLATOR has a stake in round 2
+			assert!(<AtStake<Test>>::contains_key(2, ACTIVE_COLLATOR));
+			assert!(<HadStake<Test>>::contains_key(2, ACTIVE_COLLATOR));
+
+			// End of round 3
+			roll_to_current_round_end();
+
+			// ACTIVE_COLLATOR has a no stake in round 1 anymore due to the distribution of rewards
+			// HadStake is still present
+			assert!(!<AtStake<Test>>::contains_key(1, ACTIVE_COLLATOR));
+			assert!(<HadStake<Test>>::contains_key(1, ACTIVE_COLLATOR));
+
+			// Round 4
+			roll_to_next_round_begin();
+			roll_to_current_round_end();
+
+			assert!(
+				!<HadStake<Test>>::contains_key(1, ACTIVE_COLLATOR),
+				"Round 1 HadStake should be cleaned up after MaxOfflineRounds"
+			);
+			assert!(<HadStake<Test>>::contains_key(2, ACTIVE_COLLATOR));
+			assert!(<HadStake<Test>>::contains_key(3, ACTIVE_COLLATOR));
+			assert!(<HadStake<Test>>::contains_key(4, ACTIVE_COLLATOR));
+		});
+}
+
 #[test]
 fn notify_inactive_collator_works() {
+	const INACTIVE_COLLATOR: AccountId = 1;
+	const ACTIVE_COLLATOR: AccountId = 2;
+
 	ExtBuilder::default()
 		.with_balances(vec![(1, 20), (2, 20), (3, 20), (4, 20), (5, 20)])
 		.with_candidates(vec![(1, 20), (2, 20), (3, 20), (4, 20), (5, 20)])
@@ -1183,38 +1238,92 @@ fn notify_inactive_collator_works() {
 			// Enable killswitch
 			<EnableMarkingOffline<Test>>::set(true);
 
-			// Round 2
-			roll_to_round_begin(2);
+			assert_eq!(<Test as crate::Config>::MaxOfflineRounds::get(), 2);
+			assert_eq!(<Test as crate::Config>::RewardPaymentDelay::get(), 2);
 
-			// Change block author
-			set_block_author(1);
+			// Round 2 - INACTIVE_COLLATOR authors blocks
+			set_block_author(INACTIVE_COLLATOR);
+			roll_to_next_round_begin();
 
-			// Finalize the first block of round 2
-			ParachainStaking::on_finalize(5);
+			// Change block author
+			set_block_author(ACTIVE_COLLATOR);
 
-			// We don't produce blocks on round 3
-			roll_to_round_begin(3);
+			// INACTIVE_COLLATOR does not produce blocks on round 2 and 3
+			roll_to_next_round_begin();
+			roll_to_next_round_begin();
 			roll_blocks(1);
 
-			// We don't produce blocks on round 4
-			roll_to_round_begin(4);
+			// On round 4 notify inactive collator
+			assert_ok!(ParachainStaking::notify_inactive_collator(
+				RuntimeOrigin::signed(1),
+				INACTIVE_COLLATOR
+			));
+
+			// Check the collator was marked as offline as it hasn't produced blocks
+			assert_events_eq!(Event::CandidateWentOffline {
+				candidate: INACTIVE_COLLATOR
+			},);
+		});
+}
+
+#[test]
+fn notify_inactive_collator_succeeds_even_after_rewards_are_distributed() {
+	const INACTIVE_COLLATOR: AccountId = 1;
+	const ACTIVE_COLLATOR: AccountId = 2;
+
+	ExtBuilder::default()
+		.with_balances(vec![(1, 20), (2, 20), (3, 20), (4, 20), (5, 20)])
+		.with_candidates(vec![(1, 20), (2, 20), (3, 20), (4, 20), (5, 20)])
+		.build()
+		.execute_with(|| {
+			// Enable killswitch
+			<EnableMarkingOffline<Test>>::set(true);
+
+			// We need (strictly) more blocks per round than collators so rewards
+			// can be distributed before the end of a round
+			assert_ok!(ParachainStaking::set_blocks_per_round(
+				RuntimeOrigin::root(),
+				6u32
+			));
+
+			// ACTIVE_COLLATOR authors all the blocks while INACTIVE_COLLATOR stays inactive
+			set_block_author(ACTIVE_COLLATOR);
+
+			// Round 2
+			roll_to_next_round_begin();
 			roll_blocks(1);
 
-			// Round 6 - notify the collator as inactive
-			roll_to_round_begin(6);
+			// INACTIVE_COLLATOR has a stake in round 1
+			assert!(<AtStake<Test>>::contains_key(1, INACTIVE_COLLATOR));
+
+			// Round 3
+			roll_to_next_round_begin();
 			roll_blocks(1);
 
-			assert_eq!(<Test as crate::Config>::MaxOfflineRounds::get(), 1);
-			assert_eq!(<Test as crate::Config>::RewardPaymentDelay::get(), 2);
+			// INACTIVE_COLLATOR has a stake in round 2
+			assert!(<AtStake<Test>>::contains_key(2, INACTIVE_COLLATOR));
 
-			// Call 'notify_inactive_collator' extrinsic
+			// End of round 3
+			roll_to_current_round_end();
+
+			// INACTIVE_COLLATOR has a no stake in round 1 anymore due to the distribution of rewards
+			assert!(!<AtStake<Test>>::contains_key(1, INACTIVE_COLLATOR));
+
+			// Call 'notify_inactive_collator' extrinsic on INACTIVE_COLLATOR
 			assert_ok!(ParachainStaking::notify_inactive_collator(
 				RuntimeOrigin::signed(1),
-				1
+				INACTIVE_COLLATOR
 			));
 
-			// Check the collator was marked as offline as it hasn't produced blocks
-			assert_events_eq!(Event::CandidateWentOffline { candidate: 1 },);
+			assert_events_eq!(
+				Event::Rewarded {
+					account: 2,
+					rewards: 0,
+				},
+				Event::CandidateWentOffline {
+					candidate: INACTIVE_COLLATOR
+				},
+			);
 		});
 }