Fix parachain upgrade scheduling when done by the owner/root

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, EnactUpgradeDirectly, ParaGenesisArgs},
+	paras::{self, ParaGenesisArgs, UpgradeStrategy},
 	Origin, ParaLifecycle,
 };
 use sp_std::{prelude::*, result};
@@ -412,9 +412,8 @@ pub mod pallet {
 		/// 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 will not be aware that the code changes. Thus, this should only
-		/// be used if the parachain e.g. is broken and requires some manual fixing. Do not use this
-		/// for upgrading your parachain as this can brick the parachain.
+		/// 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.
@@ -429,7 +428,7 @@ pub mod pallet {
 			runtime_parachains::schedule_code_upgrade::<T>(
 				para,
 				new_code,
-				EnactUpgradeDirectly::Yes,
+				UpgradeStrategy::ApplyAtExpectedBlock,
 			)?;
 			Ok(())
 		}

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

@@ -22,8 +22,8 @@
 use crate::{
 	configuration::{self, HostConfiguration},
 	disputes, dmp, hrmp,
-	paras::{self, EnactUpgradeDirectly},
-	scheduler::{self, AvailabilityTimeoutStatus},
+	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,
-				EnactUpgradeDirectly::No,
+				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,
-					EnactUpgradeDirectly::No,
+					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, enact_upgrade_directly: EnactUpgradeDirectly::No }
+			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, EnactUpgradeDirectly};
+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: EnactUpgradeDirectly,
+	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,
-			EnactUpgradeDirectly::No,
+			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(),
-				EnactUpgradeDirectly::No,
+				UpgradeStrategy::SetGoAheadSignal,
 			);
 		} else {
 			let r = Pallet::<T>::schedule_para_initialize(

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

@@ -390,15 +390,28 @@ pub(crate) enum PvfCheckCause<BlockNumber> {
 		///
 		/// If set to `Yes` it means that no `GoAheadSignal` will be set and the parachain code
 		/// will also be overwritten directly.
-		enact_upgrade_directly: EnactUpgradeDirectly,
+		upgrade_strategy: UpgradeStrategy,
 	},
 }
 
-/// Should the upgrade of the validation code be enacted directly after the checks succeeded?
+/// 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 EnactUpgradeDirectly {
-	No,
-	Yes,
+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> {
@@ -761,14 +774,27 @@ 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]
 	#[pallet::getter(fn future_code_upgrade_at)]
 	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`].
@@ -812,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]
@@ -919,7 +947,7 @@ pub mod pallet {
 				new_code,
 				relay_parent_number,
 				&config,
-				EnactUpgradeDirectly::Yes,
+				UpgradeStrategy::ApplyAtExpectedBlock,
 			);
 			Self::deposit_event(Event::CodeUpgradeScheduled(para));
 			Ok(())
@@ -1218,7 +1246,7 @@ impl<T: Config> Pallet<T> {
 	pub(crate) fn schedule_code_upgrade_external(
 		id: ParaId,
 		new_code: ValidationCode,
-		enact_upgrade_directly: EnactUpgradeDirectly,
+		upgrade_strategy: UpgradeStrategy,
 	) -> DispatchResult {
 		// Check that we can schedule an upgrade at all.
 		ensure!(Self::can_upgrade_validation_code(id), Error::<T>::CannotUpgradeCode);
@@ -1230,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, enact_upgrade_directly);
+		Self::schedule_code_upgrade(id, new_code, upgrade_block, &config, upgrade_strategy);
 		Self::deposit_event(Event::CodeUpgradeScheduled(id));
 		Ok(())
 	}
@@ -1243,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.
@@ -1346,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));
 			});
 		}
 
@@ -1451,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,
@@ -1571,14 +1628,14 @@ impl<T: Config> Pallet<T> {
 				PvfCheckCause::Onboarding(id) => {
 					weight += Self::proceed_with_onboarding(*id, sessions_observed);
 				},
-				PvfCheckCause::Upgrade { id, included_at, enact_upgrade_directly } => {
+				PvfCheckCause::Upgrade { id, included_at, upgrade_strategy } => {
 					weight += Self::proceed_with_upgrade(
 						*id,
 						code_hash,
 						now,
 						*included_at,
 						cfg,
-						*enact_upgrade_directly,
+						*upgrade_strategy,
 					);
 				},
 			}
@@ -1612,32 +1669,37 @@ impl<T: Config> Pallet<T> {
 		now: BlockNumberFor<T>,
 		relay_parent_number: BlockNumberFor<T>,
 		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
-		enact_upgrade_directly: EnactUpgradeDirectly,
+		upgrade_strategy: UpgradeStrategy,
 	) -> Weight {
 		let mut weight = Weight::zero();
 
-		if enact_upgrade_directly == EnactUpgradeDirectly::Yes {
-			FutureCodeHash::<T>::remove(&id);
-			weight.saturating_add(Self::set_current_code(id, *code_hash, now));
-		} else {
-			// 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.
-			//
-			// 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,
-			);
+		// 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.
+		//
+		// 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,
+		);
+
+		if 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));
+			});
 
-			weight += T::DbWeight::get().reads_writes(1, 4);
+			weight += T::DbWeight::get().reads_writes(0, 2);
+		} else {
 			FutureCodeUpgrades::<T>::insert(&id, expected_at);
 
 			UpcomingUpgrades::<T>::mutate(|upcoming_upgrades| {
@@ -1647,11 +1709,13 @@ impl<T: Config> Pallet<T> {
 				upcoming_upgrades.insert(insert_idx, (id, expected_at));
 			});
 
-			let expected_at = expected_at.saturated_into();
-			let log = ConsensusLog::ParaScheduleUpgradeCode(id, *code_hash, expected_at);
-			<frame_system::Pallet<T>>::deposit_log(log.into());
+			weight += T::DbWeight::get().reads_writes(1, 3);
 		}
 
+		let expected_at = expected_at.saturated_into();
+		let log = ConsensusLog::ParaScheduleUpgradeCode(id, *code_hash, expected_at);
+		<frame_system::Pallet<T>>::deposit_log(log.into());
+
 		weight
 	}
 
@@ -1885,7 +1949,7 @@ impl<T: Config> Pallet<T> {
 		new_code: ValidationCode,
 		inclusion_block_number: BlockNumberFor<T>,
 		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
-		enact_upgrade_directly: EnactUpgradeDirectly,
+		upgrade_strategy: UpgradeStrategy,
 	) -> Weight {
 		let mut weight = T::DbWeight::get().reads(1);
 
@@ -1942,11 +2006,7 @@ impl<T: Config> Pallet<T> {
 		});
 
 		weight += Self::kick_off_pvf_check(
-			PvfCheckCause::Upgrade {
-				id,
-				included_at: inclusion_block_number,
-				enact_upgrade_directly,
-			},
+			PvfCheckCause::Upgrade { id, included_at: inclusion_block_number, upgrade_strategy },
 			code_hash,
 			new_code,
 			cfg,