From e3ee909e9e418ab4fb322f9229d083ce7aaa1b96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <git@kchr.de>
Date: Mon, 10 Jul 2023 17:08:12 +0200
Subject: [PATCH] paras: Prune upgrade cooldowns (#7470)

After an upgrade cooldown triggered to remove a signal we also need to
prune it from the list.
---
 polkadot/runtime/parachains/src/paras/mod.rs   | 14 ++++++++++----
 polkadot/runtime/parachains/src/paras/tests.rs | 18 +++++++++++++++---
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/polkadot/runtime/parachains/src/paras/mod.rs b/polkadot/runtime/parachains/src/paras/mod.rs
index 8c8eabf9d2e..554a393d605 100644
--- a/polkadot/runtime/parachains/src/paras/mod.rs
+++ b/polkadot/runtime/parachains/src/paras/mod.rs
@@ -1262,7 +1262,7 @@ impl<T: Config> Pallet<T> {
 		// Persist parachains into the storage explicitly.
 		drop(parachains);
 
-		return outgoing
+		outgoing
 	}
 
 	// note replacement of the code of para with given `id`, which occured in the
@@ -1389,9 +1389,15 @@ impl<T: Config> Pallet<T> {
 	/// See `process_scheduled_upgrade_changes` for more details.
 	fn process_scheduled_upgrade_cooldowns(now: T::BlockNumber) {
 		UpgradeCooldowns::<T>::mutate(|upgrade_cooldowns: &mut Vec<(ParaId, T::BlockNumber)>| {
-			for &(para, _) in upgrade_cooldowns.iter().take_while(|&(_, at)| at <= &now) {
-				UpgradeRestrictionSignal::<T>::remove(&para);
-			}
+			// Remove all expired signals and also prune the cooldowns.
+			upgrade_cooldowns.retain(|(para, at)| {
+				if at <= &now {
+					UpgradeRestrictionSignal::<T>::remove(&para);
+					false
+				} else {
+					true
+				}
+			});
 		});
 	}
 
diff --git a/polkadot/runtime/parachains/src/paras/tests.rs b/polkadot/runtime/parachains/src/paras/tests.rs
index 1acd8809412..7e3e80d1026 100644
--- a/polkadot/runtime/parachains/src/paras/tests.rs
+++ b/polkadot/runtime/parachains/src/paras/tests.rs
@@ -441,7 +441,7 @@ fn code_upgrade_applied_after_delay() {
 		run_to_block(2, Some(vec![1]));
 		assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 
-		let expected_at = {
+		let (expected_at, next_possible_upgrade_at) = {
 			// this parablock is in the context of block 1.
 			let expected_at = 1 + validation_upgrade_delay;
 			let next_possible_upgrade_at = 1 + validation_upgrade_cooldown;
@@ -460,7 +460,7 @@ fn code_upgrade_applied_after_delay() {
 			check_code_is_stored(&original_code);
 			check_code_is_stored(&new_code);
 
-			expected_at
+			(expected_at, next_possible_upgrade_at)
 		};
 
 		run_to_block(expected_at, None);
@@ -495,9 +495,21 @@ fn code_upgrade_applied_after_delay() {
 			assert!(FutureCodeHash::<Test>::get(&para_id).is_none());
 			assert!(UpgradeGoAheadSignal::<Test>::get(&para_id).is_none());
 			assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
+			assert_eq!(
+				UpgradeRestrictionSignal::<Test>::get(&para_id),
+				Some(UpgradeRestriction::Present),
+			);
+			assert_eq!(UpgradeCooldowns::<Test>::get(), vec![(para_id, next_possible_upgrade_at)]);
 			check_code_is_stored(&original_code);
 			check_code_is_stored(&new_code);
 		}
+
+		run_to_block(next_possible_upgrade_at + 1, None);
+
+		{
+			assert!(UpgradeRestrictionSignal::<Test>::get(&para_id).is_none());
+			assert!(UpgradeCooldowns::<Test>::get().is_empty());
+		}
 	});
 }
 
@@ -568,7 +580,7 @@ fn code_upgrade_applied_after_delay_even_when_late() {
 		// the upgrade.
 		{
 			// The signal should be set to go-ahead until the new head is actually processed.
-			assert_eq!(UpgradeGoAheadSignal::<Test>::get(&para_id), Some(UpgradeGoAhead::GoAhead),);
+			assert_eq!(UpgradeGoAheadSignal::<Test>::get(&para_id), Some(UpgradeGoAhead::GoAhead));
 
 			Paras::note_new_head(para_id, Default::default(), expected_at + 4);
 
-- 
GitLab