Fix parachain upgrade scheduling when done by the owner/root (#3341)

When using `schedule_code_upgrade` to change the code of a parachain in
the relay chain runtime, we had already fixed to not set the `GoAhead`
signal. This was done to not brick any parachain after the upgrade,
because they were seeing the signal without having any upgrade prepared.
The remaining problem is that the parachain code is only upgraded after
a parachain header was enacted, aka the parachain made some progress.
However, this is quite complicated if the parachain is bricked (which is
the most common scenario why to manually schedule a code upgrade). Thus,
this pull request replaces `SetGoAhead` with `UpgradeStrategy` to signal
to the logic kind of strategy want to use. The strategies are either
`SetGoAheadSignal` or `ApplyAtExpectedBlock`. `SetGoAheadSignal` sets
the go ahead signal as before and awaits a parachain block.
`ApplyAtExpectedBlock` schedules the upgrade and applies it directly at
the `expected_block` without waiting for the parachain to make any kind
of progress.

polkadot/runtime/common/src/paras_registrar/mod.rs

@@ -29,7 +29,7 @@ use frame_system::{self, ensure_root, ensure_signed};
 use primitives::{HeadData, Id as ParaId, ValidationCode, LOWEST_PUBLIC_ID, MIN_CODE_SIZE};
 use runtime_parachains::{
 	configuration, ensure_parachain,
-	paras::{self, ParaGenesisArgs, SetGoAhead},
+	paras::{self, ParaGenesisArgs, UpgradeStrategy},
 	Origin, ParaLifecycle,
 };
 use sp_std::{prelude::*, result};
@@ -408,6 +408,13 @@ pub mod pallet {
 
 		/// Schedule a parachain upgrade.
 		///
+		/// This will kick off a check of `new_code` by all validators. After the majority of the
+		/// validators have reported on the validity of the code, the code will either be enacted
+		/// or the upgrade will be rejected. If the code will be enacted, the current code of the
+		/// parachain will be overwritten directly. This means that any PoV will be checked by this
+		/// new code. The parachain itself will not be informed explictely that the validation code
+		/// has changed.
+		///
 		/// Can be called by Root, the parachain, or the parachain manager if the parachain is
 		/// unlocked.
 		#[pallet::call_index(7)]
@@ -418,7 +425,11 @@ pub mod pallet {
 			new_code: ValidationCode,
 		) -> DispatchResult {
 			Self::ensure_root_para_or_owner(origin, para)?;
-			runtime_parachains::schedule_code_upgrade::<T>(para, new_code, SetGoAhead::No)?;
+			runtime_parachains::schedule_code_upgrade::<T>(
+				para,
+				new_code,
+				UpgradeStrategy::ApplyAtExpectedBlock,
+			)?;
 			Ok(())
 		}
 

polkadot/runtime/parachains/src/inclusion/mod.rs

@@ -22,7 +22,7 @@
 use crate::{
 	configuration::{self, HostConfiguration},
 	disputes, dmp, hrmp,
-	paras::{self, SetGoAhead},
+	paras::{self, UpgradeStrategy},
 	scheduler,
 	shared::{self, AllowedRelayParentsTracker},
 	util::make_persisted_validation_data_with_parent,
@@ -839,7 +839,7 @@ impl<T: Config> Pallet<T> {
 				new_code,
 				now,
 				&config,
-				SetGoAhead::Yes,
+				UpgradeStrategy::SetGoAheadSignal,
 			));
 		}
 

polkadot/runtime/parachains/src/inclusion/tests.rs

@@ -1590,7 +1590,7 @@ fn candidate_checks() {
 					vec![9, 8, 7, 6, 5, 4, 3, 2, 1].into(),
 					expected_at,
 					&cfg,
-					SetGoAhead::Yes,
+					UpgradeStrategy::SetGoAheadSignal,
 				);
 			}
 
@@ -2857,7 +2857,7 @@ fn para_upgrade_delay_scheduled_from_inclusion() {
 		let cause = &active_vote_state.causes()[0];
 		// Upgrade block is the block of inclusion, not candidate's parent.
 		assert_matches!(cause,
-			paras::PvfCheckCause::Upgrade { id, included_at, set_go_ahead: SetGoAhead::Yes }
+			paras::PvfCheckCause::Upgrade { id, included_at, upgrade_strategy: UpgradeStrategy::SetGoAheadSignal }
 				if id == &chain_a && included_at == &7
 		);
 	});

polkadot/runtime/parachains/src/lib.rs

@@ -54,7 +54,7 @@ mod mock;
 mod ump_tests;
 
 pub use origin::{ensure_parachain, Origin};
-pub use paras::{ParaLifecycle, SetGoAhead};
+pub use paras::{ParaLifecycle, UpgradeStrategy};
 use primitives::{HeadData, Id as ParaId, ValidationCode};
 use sp_runtime::{DispatchResult, FixedU128};
 
@@ -104,7 +104,7 @@ pub fn schedule_parachain_downgrade<T: paras::Config>(id: ParaId) -> Result<(),
 pub fn schedule_code_upgrade<T: paras::Config>(
 	id: ParaId,
 	new_code: ValidationCode,
-	set_go_ahead: SetGoAhead,
+	set_go_ahead: UpgradeStrategy,
 ) -> DispatchResult {
 	paras::Pallet::<T>::schedule_code_upgrade_external(id, new_code, set_go_ahead)
 }

polkadot/runtime/parachains/src/paras/benchmarking.rs

@@ -129,7 +129,7 @@ benchmarks! {
 			ValidationCode(vec![0]),
 			expired,
 			&config,
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 	}: _(RawOrigin::Root, para_id, new_head)
 	verify {

polkadot/runtime/parachains/src/paras/benchmarking/pvf_check.rs

@@ -177,7 +177,7 @@ where
 				validation_code,
 				/* relay_parent_number */ 1u32.into(),
 				&configuration::Pallet::<T>::config(),
-				SetGoAhead::Yes,
+				UpgradeStrategy::SetGoAheadSignal,
 			);
 		} else {
 			let r = Pallet::<T>::schedule_para_initialize(

polkadot/runtime/parachains/src/paras/mod.rs

@@ -386,16 +386,32 @@ pub(crate) enum PvfCheckCause<BlockNumber> {
 		///
 		/// See https://github.com/paritytech/polkadot/issues/4601 for detailed explanation.
 		included_at: BlockNumber,
-		/// Whether or not the given para should be sent the `GoAhead` signal.
-		set_go_ahead: SetGoAhead,
+		/// Whether or not the upgrade should be enacted directly.
+		///
+		/// If set to `Yes` it means that no `GoAheadSignal` will be set and the parachain code
+		/// will also be overwritten directly.
+		upgrade_strategy: UpgradeStrategy,
 	},
 }
 
-/// Should the `GoAhead` signal be set after a successful check of the new wasm binary?
+/// The strategy on how to handle a validation code upgrade.
+///
+/// When scheduling a parachain code upgrade the upgrade first is checked by all validators. The
+/// validators ensure that the new validation code can be compiled and instantiated. After the
+/// majority of the validators have reported their checking result the upgrade is either scheduled
+/// or aborted. This strategy then comes into play around the relay chain block this upgrade was
+/// scheduled in.
 #[derive(Debug, Copy, Clone, PartialEq, TypeInfo, Decode, Encode)]
-pub enum SetGoAhead {
-	Yes,
-	No,
+pub enum UpgradeStrategy {
+	/// Set the `GoAhead` signal to inform the parachain that it is time to upgrade.
+	///
+	/// The upgrade will then be applied after the first parachain block was enacted that must have
+	/// observed the `GoAhead` signal.
+	SetGoAheadSignal,
+	/// Apply the upgrade directly at the expected relay chain block.
+	///
+	/// This doesn't wait for the parachain to make any kind of progress.
+	ApplyAtExpectedBlock,
 }
 
 impl<BlockNumber> PvfCheckCause<BlockNumber> {
@@ -758,7 +774,8 @@ pub mod pallet {
 	pub(super) type PastCodePruning<T: Config> =
 		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
 
-	/// The block number at which the planned code change is expected for a para.
+	/// The block number at which the planned code change is expected for a parachain.
+	///
 	/// The change will be applied after the first parablock for this ID included which executes
 	/// in the context of a relay chain block with a number >= `expected_at`.
 	#[pallet::storage]
@@ -766,6 +783,18 @@ pub mod pallet {
 	pub(super) type FutureCodeUpgrades<T: Config> =
 		StorageMap<_, Twox64Concat, ParaId, BlockNumberFor<T>>;
 
+	/// The list of upcoming future code upgrades.
+	///
+	/// Each item is a pair of the parachain and the expected block at which the upgrade should be
+	/// applied. The upgrade will be applied at the given relay chain block. In contrast to
+	/// [`FutureCodeUpgrades`] this code upgrade will be applied regardless the parachain making any
+	/// progress or not.
+	///
+	/// Ordered ascending by block number.
+	#[pallet::storage]
+	pub(super) type FutureCodeUpgradesAt<T: Config> =
+		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
+
 	/// The actual future code hash of a para.
 	///
 	/// Corresponding code can be retrieved with [`CodeByHash`].
@@ -809,8 +838,10 @@ pub mod pallet {
 	pub(super) type UpgradeCooldowns<T: Config> =
 		StorageValue<_, Vec<(ParaId, BlockNumberFor<T>)>, ValueQuery>;
 
-	/// The list of upcoming code upgrades. Each item is a pair of which para performs a code
-	/// upgrade and at which relay-chain block it is expected at.
+	/// The list of upcoming code upgrades.
+	///
+	/// Each item is a pair of which para performs a code upgrade and at which relay-chain block it
+	/// is expected at.
 	///
 	/// Ordered ascending by block number.
 	#[pallet::storage]
@@ -880,21 +911,9 @@ pub mod pallet {
 			new_code: ValidationCode,
 		) -> DispatchResult {
 			ensure_root(origin)?;
-			let maybe_prior_code_hash = CurrentCodeHash::<T>::get(&para);
 			let new_code_hash = new_code.hash();
 			Self::increase_code_ref(&new_code_hash, &new_code);
-			CurrentCodeHash::<T>::insert(&para, new_code_hash);
-
-			let now = frame_system::Pallet::<T>::block_number();
-			if let Some(prior_code_hash) = maybe_prior_code_hash {
-				Self::note_past_code(para, now, now, prior_code_hash);
-			} else {
-				log::error!(
-					target: LOG_TARGET,
-					"Pallet paras storage is inconsistent, prior code not found {:?}",
-					&para
-				);
-			}
+			Self::set_current_code(para, new_code_hash, frame_system::Pallet::<T>::block_number());
 			Self::deposit_event(Event::CurrentCodeUpdated(para));
 			Ok(())
 		}
@@ -928,7 +947,7 @@ pub mod pallet {
 				new_code,
 				relay_parent_number,
 				&config,
-				SetGoAhead::No,
+				UpgradeStrategy::ApplyAtExpectedBlock,
 			);
 			Self::deposit_event(Event::CodeUpgradeScheduled(para));
 			Ok(())
@@ -1227,7 +1246,7 @@ impl<T: Config> Pallet<T> {
 	pub(crate) fn schedule_code_upgrade_external(
 		id: ParaId,
 		new_code: ValidationCode,
-		set_go_ahead: SetGoAhead,
+		upgrade_strategy: UpgradeStrategy,
 	) -> DispatchResult {
 		// Check that we can schedule an upgrade at all.
 		ensure!(Self::can_upgrade_validation_code(id), Error::<T>::CannotUpgradeCode);
@@ -1239,7 +1258,7 @@ impl<T: Config> Pallet<T> {
 		let current_block = frame_system::Pallet::<T>::block_number();
 		// Schedule the upgrade with a delay just like if a parachain triggered the upgrade.
 		let upgrade_block = current_block.saturating_add(config.validation_upgrade_delay);
-		Self::schedule_code_upgrade(id, new_code, upgrade_block, &config, set_go_ahead);
+		Self::schedule_code_upgrade(id, new_code, upgrade_block, &config, upgrade_strategy);
 		Self::deposit_event(Event::CodeUpgradeScheduled(id));
 		Ok(())
 	}
@@ -1252,8 +1271,9 @@ impl<T: Config> Pallet<T> {
 
 	/// Called by the initializer to initialize the paras pallet.
 	pub(crate) fn initializer_initialize(now: BlockNumberFor<T>) -> Weight {
-		let weight = Self::prune_old_code(now);
-		weight + Self::process_scheduled_upgrade_changes(now)
+		Self::prune_old_code(now) +
+			Self::process_scheduled_upgrade_changes(now) +
+			Self::process_future_code_upgrades_at(now)
 	}
 
 	/// Called by the initializer to finalize the paras pallet.
@@ -1355,16 +1375,13 @@ impl<T: Config> Pallet<T> {
 			// NOTE both of those iterates over the list and the outgoing. We do not expect either
 			//      of these to be large. Thus should be fine.
 			UpcomingUpgrades::<T>::mutate(|upcoming_upgrades| {
-				*upcoming_upgrades = mem::take(upcoming_upgrades)
-					.into_iter()
-					.filter(|(para, _)| !outgoing.contains(para))
-					.collect();
+				upcoming_upgrades.retain(|(para, _)| !outgoing.contains(para));
 			});
 			UpgradeCooldowns::<T>::mutate(|upgrade_cooldowns| {
-				*upgrade_cooldowns = mem::take(upgrade_cooldowns)
-					.into_iter()
-					.filter(|(para, _)| !outgoing.contains(para))
-					.collect();
+				upgrade_cooldowns.retain(|(para, _)| !outgoing.contains(para));
+			});
+			FutureCodeUpgradesAt::<T>::mutate(|future_upgrades| {
+				future_upgrades.retain(|(para, _)| !outgoing.contains(para));
 			});
 		}
 
@@ -1460,6 +1477,37 @@ impl<T: Config> Pallet<T> {
 		T::DbWeight::get().reads_writes(1 + pruning_tasks_done, 2 * pruning_tasks_done)
 	}
 
+	/// Process the future code upgrades that should be applied directly.
+	///
+	/// Upgrades that should not be applied directly are being processed in
+	/// [`Self::process_scheduled_upgrade_changes`].
+	fn process_future_code_upgrades_at(now: BlockNumberFor<T>) -> Weight {
+		// account weight for `FutureCodeUpgradeAt::mutate`.
+		let mut weight = T::DbWeight::get().reads_writes(1, 1);
+		FutureCodeUpgradesAt::<T>::mutate(
+			|upcoming_upgrades: &mut Vec<(ParaId, BlockNumberFor<T>)>| {
+				let num = upcoming_upgrades.iter().take_while(|&(_, at)| at <= &now).count();
+				for (id, expected_at) in upcoming_upgrades.drain(..num) {
+					weight += T::DbWeight::get().reads_writes(1, 1);
+
+					// Both should always be `Some` in this case, since a code upgrade is scheduled.
+					let new_code_hash = if let Some(new_code_hash) = FutureCodeHash::<T>::take(&id)
+					{
+						new_code_hash
+					} else {
+						log::error!(target: LOG_TARGET, "Missing future code hash for {:?}", &id);
+						continue
+					};
+
+					weight += Self::set_current_code(id, new_code_hash, expected_at);
+				}
+				num
+			},
+		);
+
+		weight
+	}
+
 	/// Process the timers related to upgrades. Specifically, the upgrade go ahead signals toggle
 	/// and the upgrade cooldown restrictions. However, this function does not actually unset
 	/// the upgrade restriction, that will happen in the `initializer_finalize` function. However,
@@ -1580,14 +1628,14 @@ impl<T: Config> Pallet<T> {
 				PvfCheckCause::Onboarding(id) => {
 					weight += Self::proceed_with_onboarding(*id, sessions_observed);
 				},
-				PvfCheckCause::Upgrade { id, included_at, set_go_ahead } => {
+				PvfCheckCause::Upgrade { id, included_at, upgrade_strategy } => {
 					weight += Self::proceed_with_upgrade(
 						*id,
 						code_hash,
 						now,
 						*included_at,
 						cfg,
-						*set_go_ahead,
+						*upgrade_strategy,
 					);
 				},
 			}
@@ -1621,38 +1669,50 @@ impl<T: Config> Pallet<T> {
 		now: BlockNumberFor<T>,
 		relay_parent_number: BlockNumberFor<T>,
 		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
-		set_go_ahead: SetGoAhead,
+		upgrade_strategy: UpgradeStrategy,
 	) -> Weight {
 		let mut weight = Weight::zero();
 
-		// Compute the relay-chain block number starting at which the code upgrade is ready to be
-		// applied.
+		// Compute the relay-chain block number starting at which the code upgrade is ready to
+		// be applied.
 		//
-		// The first parablock that has a relay-parent higher or at the same height of `expected_at`
-		// will trigger the code upgrade. The parablock that comes after that will be validated
-		// against the new validation code.
+		// The first parablock that has a relay-parent higher or at the same height of
+		// `expected_at` will trigger the code upgrade. The parablock that comes after that will
+		// be validated against the new validation code.
 		//
-		// Here we are trying to choose the block number that will have `validation_upgrade_delay`
-		// blocks from the relay-parent of inclusion of the the block that scheduled code upgrade
-		// but no less than `minimum_validation_upgrade_delay`. We want this delay out of caution
-		// so that when the last vote for pre-checking comes the parachain will have some time until
-		// the upgrade finally takes place.
+		// Here we are trying to choose the block number that will have
+		// `validation_upgrade_delay` blocks from the relay-parent of inclusion of the the block
+		// that scheduled code upgrade but no less than `minimum_validation_upgrade_delay`. We
+		// want this delay out of caution so that when the last vote for pre-checking comes the
+		// parachain will have some time until the upgrade finally takes place.
 		let expected_at = cmp::max(
 			relay_parent_number + cfg.validation_upgrade_delay,
 			now + cfg.minimum_validation_upgrade_delay,
 		);
 
-		weight += T::DbWeight::get().reads_writes(1, 4);
-		FutureCodeUpgrades::<T>::insert(&id, expected_at);
+		match upgrade_strategy {
+			UpgradeStrategy::ApplyAtExpectedBlock => {
+				FutureCodeUpgradesAt::<T>::mutate(|future_upgrades| {
+					let insert_idx = future_upgrades
+						.binary_search_by_key(&expected_at, |&(_, b)| b)
+						.unwrap_or_else(|idx| idx);
+					future_upgrades.insert(insert_idx, (id, expected_at));
+				});
 
-		// Only set an upcoming upgrade if `GoAhead` signal should be set for the respective para.
-		if set_go_ahead == SetGoAhead::Yes {
-			UpcomingUpgrades::<T>::mutate(|upcoming_upgrades| {
-				let insert_idx = upcoming_upgrades
-					.binary_search_by_key(&expected_at, |&(_, b)| b)
-					.unwrap_or_else(|idx| idx);
-				upcoming_upgrades.insert(insert_idx, (id, expected_at));
-			});
+				weight += T::DbWeight::get().reads_writes(0, 2);
+			},
+			UpgradeStrategy::SetGoAheadSignal => {
+				FutureCodeUpgrades::<T>::insert(&id, expected_at);
+
+				UpcomingUpgrades::<T>::mutate(|upcoming_upgrades| {
+					let insert_idx = upcoming_upgrades
+						.binary_search_by_key(&expected_at, |&(_, b)| b)
+						.unwrap_or_else(|idx| idx);
+					upcoming_upgrades.insert(insert_idx, (id, expected_at));
+				});
+
+				weight += T::DbWeight::get().reads_writes(1, 3);
+			},
 		}
 
 		let expected_at = expected_at.saturated_into();
@@ -1892,7 +1952,7 @@ impl<T: Config> Pallet<T> {
 		new_code: ValidationCode,
 		inclusion_block_number: BlockNumberFor<T>,
 		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
-		set_go_ahead: SetGoAhead,
+		upgrade_strategy: UpgradeStrategy,
 	) -> Weight {
 		let mut weight = T::DbWeight::get().reads(1);
 
@@ -1949,7 +2009,7 @@ impl<T: Config> Pallet<T> {
 		});
 
 		weight += Self::kick_off_pvf_check(
-			PvfCheckCause::Upgrade { id, included_at: inclusion_block_number, set_go_ahead },
+			PvfCheckCause::Upgrade { id, included_at: inclusion_block_number, upgrade_strategy },
 			code_hash,
 			new_code,
 			cfg,
@@ -2061,24 +2121,10 @@ impl<T: Config> Pallet<T> {
 					log::error!(target: LOG_TARGET, "Missing future code hash for {:?}", &id);
 					return T::DbWeight::get().reads_writes(3, 1 + 3)
 				};
-				let maybe_prior_code_hash = CurrentCodeHash::<T>::get(&id);
-				CurrentCodeHash::<T>::insert(&id, &new_code_hash);
 
-				let log = ConsensusLog::ParaUpgradeCode(id, new_code_hash);
-				<frame_system::Pallet<T>>::deposit_log(log.into());
+				let weight = Self::set_current_code(id, new_code_hash, expected_at);
 
-				// `now` is only used for registering pruning as part of `fn note_past_code`
-				let now = <frame_system::Pallet<T>>::block_number();
-
-				let weight = if let Some(prior_code_hash) = maybe_prior_code_hash {
-					Self::note_past_code(id, expected_at, now, prior_code_hash)
-				} else {
-					log::error!(target: LOG_TARGET, "Missing prior code hash for para {:?}", &id);
-					Weight::zero()
-				};
-
-				// add 1 to writes due to heads update.
-				weight + T::DbWeight::get().reads_writes(3, 1 + 3)
+				weight + T::DbWeight::get().reads_writes(3, 3)
 			} else {
 				T::DbWeight::get().reads_writes(1, 1 + 0)
 			}
@@ -2094,6 +2140,34 @@ impl<T: Config> Pallet<T> {
 		weight.saturating_add(T::OnNewHead::on_new_head(id, &new_head))
 	}
 
+	/// Set the current code for the given parachain.
+	// `at` for para-triggered replacement is the block number of the relay-chain
+	// block in whose context the parablock was executed
+	// (i.e. number of `relay_parent` in the receipt)
+	pub(crate) fn set_current_code(
+		id: ParaId,
+		new_code_hash: ValidationCodeHash,
+		at: BlockNumberFor<T>,
+	) -> Weight {
+		let maybe_prior_code_hash = CurrentCodeHash::<T>::get(&id);
+		CurrentCodeHash::<T>::insert(&id, &new_code_hash);
+
+		let log = ConsensusLog::ParaUpgradeCode(id, new_code_hash);
+		<frame_system::Pallet<T>>::deposit_log(log.into());
+
+		// `now` is only used for registering pruning as part of `fn note_past_code`
+		let now = <frame_system::Pallet<T>>::block_number();
+
+		let weight = if let Some(prior_code_hash) = maybe_prior_code_hash {
+			Self::note_past_code(id, at, now, prior_code_hash)
+		} else {
+			log::error!(target: LOG_TARGET, "Missing prior code hash for para {:?}", &id);
+			Weight::zero()
+		};
+
+		weight + T::DbWeight::get().writes(1)
+	}
+
 	/// Returns the list of PVFs (aka validation code) that require casting a vote by a validator in
 	/// the active validator set.
 	pub(crate) fn pvfs_require_precheck() -> Vec<ValidationCodeHash> {

polkadot/runtime/parachains/src/paras/tests.rs

@@ -451,7 +451,7 @@ fn code_upgrade_applied_after_delay() {
 				new_code.clone(),
 				1,
 				&Configuration::config(),
-				SetGoAhead::Yes,
+				UpgradeStrategy::SetGoAheadSignal,
 			);
 			// Include votes for super-majority.
 			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
@@ -521,7 +521,7 @@ fn code_upgrade_applied_after_delay() {
 }
 
 #[test]
-fn code_upgrade_applied_without_setting_go_ahead_signal() {
+fn upgrade_strategy_apply_at_expected_block_works() {
 	let code_retention_period = 10;
 	let validation_upgrade_delay = 5;
 	let validation_upgrade_cooldown = 10;
@@ -560,77 +560,42 @@ fn code_upgrade_applied_without_setting_go_ahead_signal() {
 		run_to_block(2, Some(vec![1]));
 		assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 
-		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;
-			// `set_go_ahead` parameter set to `false` which prevents signaling the parachain
-			// with the `GoAhead` signal.
-			Paras::schedule_code_upgrade(
-				para_id,
-				new_code.clone(),
-				1,
-				&Configuration::config(),
-				SetGoAhead::No,
-			);
-			// Include votes for super-majority.
-			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
-
-			Paras::note_new_head(para_id, Default::default(), 1);
-
-			assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
-			assert_eq!(FutureCodeUpgrades::<Test>::get(&para_id), Some(expected_at));
-			assert_eq!(FutureCodeHash::<Test>::get(&para_id), Some(new_code.hash()));
-			assert_eq!(UpcomingUpgrades::<Test>::get(), vec![]);
-			assert_eq!(UpgradeCooldowns::<Test>::get(), vec![(para_id, next_possible_upgrade_at)]);
-			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
-			check_code_is_stored(&original_code);
-			check_code_is_stored(&new_code);
-
-			(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;
+		// `set_go_ahead` parameter set to `false` which prevents signaling the parachain
+		// with the `GoAhead` signal.
+		Paras::schedule_code_upgrade(
+			para_id,
+			new_code.clone(),
+			1,
+			&Configuration::config(),
+			UpgradeStrategy::ApplyAtExpectedBlock,
+		);
+		// Include votes for super-majority.
+		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
+		assert!(FutureCodeUpgradesAt::<Test>::get().iter().any(|(id, _)| *id == para_id));
 
+		// Going to the expected block triggers the upgrade directly.
 		run_to_block(expected_at, None);
 
-		// the candidate is in the context of the parent of `expected_at`,
-		// thus does not trigger the code upgrade. However, now the `UpgradeGoAheadSignal`
-		// should not be set.
-		{
-			Paras::note_new_head(para_id, Default::default(), expected_at - 1);
+		// Reporting a head doesn't change anything.
+		Paras::note_new_head(para_id, Default::default(), expected_at - 1);
 
-			assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
-			assert_eq!(FutureCodeUpgrades::<Test>::get(&para_id), Some(expected_at));
-			assert_eq!(FutureCodeHash::<Test>::get(&para_id), Some(new_code.hash()));
-			assert!(UpgradeGoAheadSignal::<Test>::get(&para_id).is_none());
-			assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
-			check_code_is_stored(&original_code);
-			check_code_is_stored(&new_code);
-		}
-
-		run_to_block(expected_at + 1, None);
-
-		// the candidate is in the context of `expected_at`, and triggers
-		// the upgrade.
-		{
-			Paras::note_new_head(para_id, Default::default(), expected_at);
-
-			assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(expected_at));
-			assert_eq!(
-				PastCodeHash::<Test>::get(&(para_id, expected_at)),
-				Some(original_code.hash()),
-			);
-			assert!(FutureCodeUpgrades::<Test>::get(&para_id).is_none());
-			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);
-		}
+		assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(expected_at));
+		assert_eq!(PastCodeHash::<Test>::get(&(para_id, expected_at)), Some(original_code.hash()));
+		assert!(FutureCodeUpgrades::<Test>::get(&para_id).is_none());
+		assert!(FutureCodeUpgradesAt::<Test>::get().iter().all(|(id, _)| *id != para_id));
+		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);
 
@@ -688,7 +653,7 @@ fn code_upgrade_applied_after_delay_even_when_late() {
 				new_code.clone(),
 				1,
 				&Configuration::config(),
-				SetGoAhead::Yes,
+				UpgradeStrategy::SetGoAheadSignal,
 			);
 			// Include votes for super-majority.
 			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
@@ -772,7 +737,7 @@ fn submit_code_change_when_not_allowed_is_err() {
 			new_code.clone(),
 			1,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		// Include votes for super-majority.
 		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
@@ -790,7 +755,7 @@ fn submit_code_change_when_not_allowed_is_err() {
 			newer_code.clone(),
 			2,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		assert_eq!(
 			FutureCodeUpgrades::<Test>::get(&para_id),
@@ -854,7 +819,7 @@ fn upgrade_restriction_elapsed_doesnt_mean_can_upgrade() {
 			new_code.clone(),
 			0,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		// Include votes for super-majority.
 		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
@@ -879,7 +844,7 @@ fn upgrade_restriction_elapsed_doesnt_mean_can_upgrade() {
 			newer_code.clone(),
 			30,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		assert_eq!(FutureCodeUpgrades::<Test>::get(&para_id), Some(0 + validation_upgrade_delay));
 	});
@@ -940,7 +905,7 @@ fn full_parachain_cleanup_storage() {
 				new_code.clone(),
 				1,
 				&Configuration::config(),
-				SetGoAhead::Yes,
+				UpgradeStrategy::SetGoAheadSignal,
 			);
 			// Include votes for super-majority.
 			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
@@ -1036,7 +1001,7 @@ fn cannot_offboard_ongoing_pvf_check() {
 			new_code.clone(),
 			RELAY_PARENT,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		assert!(!Paras::pvfs_require_precheck().is_empty());
 
@@ -1194,7 +1159,7 @@ fn code_hash_at_returns_up_to_end_of_code_retention_period() {
 			new_code.clone(),
 			0,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		// Include votes for super-majority.
 		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
@@ -1303,7 +1268,7 @@ fn pvf_check_coalescing_onboarding_and_upgrade() {
 			validation_code.clone(),
 			RELAY_PARENT,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		assert!(!Paras::pvfs_require_precheck().is_empty());
 
@@ -1413,7 +1378,7 @@ fn pvf_check_upgrade_reject() {
 			new_code.clone(),
 			RELAY_PARENT,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		check_code_is_stored(&new_code);
 
@@ -1599,7 +1564,7 @@ fn include_pvf_check_statement_refunds_weight() {
 			new_code.clone(),
 			RELAY_PARENT,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 
 		let mut stmts = IntoIterator::into_iter([0, 1, 2, 3])
@@ -1700,7 +1665,7 @@ fn poke_unused_validation_code_doesnt_remove_code_with_users() {
 			validation_code.clone(),
 			1,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		Paras::note_new_head(para_id, HeadData::default(), 1);
 
@@ -1771,7 +1736,7 @@ fn add_trusted_validation_code_insta_approval() {
 			validation_code.clone(),
 			1,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		Paras::note_new_head(para_id, HeadData::default(), 1);
 
@@ -1813,7 +1778,7 @@ fn add_trusted_validation_code_enacts_existing_pvf_vote() {
 			validation_code.clone(),
 			1,
 			&Configuration::config(),
-			SetGoAhead::Yes,
+			UpgradeStrategy::SetGoAheadSignal,
 		);
 		Paras::note_new_head(para_id, HeadData::default(), 1);
 

prdoc/pr_3341.prdoc

@@ -0,0 +1,18 @@
+title: "Fix `schedule_code_upgrade` when called by the owner/root"
+
+doc:
+  - audience: Runtime User
+    description: |
+      Fixes `schedule_code_upgrade` when being used by the owner/root. The call is used for
+      manually upgrading the validation code of a parachain on the relay chain. It was failing
+      before because the relay chain waited for the parachain to make progress. However, this
+      call is mostly used for when a parachain are bricked which means that they are not able
+      anymore to build any blocks. The fix is to schedule the validation code upgrade and then
+      to enact it at the scheduled block. The enacting happens now without requiring the parachain
+      to make any progress.
+
+crates:
+  - name: polkadot-runtime-common
+    bump: minor
+  - name: polkadot-runtime-parachains
+    bump: major