Introduce upgrade go-ahead and upgrade restriction signals (#3371)

* Introduce upgrade goahead and upgrade restriction signals

* Explicit encoding indicies for exposed enums

* typo: abscent -> absent

* Prune cooldowns as well

* Please hunspell

primitives/src/v1/mod.rs

@@ -182,6 +182,42 @@ pub mod well_known_keys {
 				.collect()
 		})
 	}
+
+	/// The signal that indicates whether the parachain should go-ahead with the proposed validation
+	/// code upgrade.
+	///
+	/// The storage entry stores a value of `UpgradeGoAhead` type.
+	pub fn upgrade_go_ahead_signal(para_id: Id) -> Vec<u8> {
+		let prefix = hex!["cd710b30bd2eab0352ddcc26417aa1949e94c040f5e73d9b7addd6cb603d15d3"];
+
+		para_id.using_encoded(|para_id: &[u8]| {
+			prefix
+				.as_ref()
+				.iter()
+				.chain(twox_64(para_id).iter())
+				.chain(para_id.iter())
+				.cloned()
+				.collect()
+		})
+	}
+
+	/// The signal that indicates whether the parachain is disallowed to signal an upgrade at this
+	/// relay-parent.
+	///
+	/// The storage entry stores a value of `UpgradeRestriction` type.
+	pub fn upgrade_restriction_signal(para_id: Id) -> Vec<u8> {
+		let prefix = hex!["cd710b30bd2eab0352ddcc26417aa194f27bbb460270642b5bcaf032ea04d56a"];
+
+		para_id.using_encoded(|para_id: &[u8]| {
+			prefix
+				.as_ref()
+				.iter()
+				.chain(twox_64(para_id).iter())
+				.chain(para_id.iter())
+				.cloned()
+				.collect()
+		})
+	}
 }
 
 /// Unique identifier for the Parachains Inherent
@@ -1035,6 +1071,38 @@ pub struct AbridgedHrmpChannel {
 	pub mqc_head: Option<Hash>,
 }
 
+/// A possible upgrade restriction that prevents a parachain from performing an upgrade.
+#[derive(Encode, Decode, PartialEq, RuntimeDebug)]
+pub enum UpgradeRestriction {
+	/// There is an upgrade restriction and there are no details about its specifics nor how long
+	/// it could last.
+	#[codec(index = 0)]
+	Present,
+}
+
+/// A struct that the relay-chain communicates to a parachain indicating what course of action the
+/// parachain should take in the coordinated parachain validation code upgrade process.
+///
+/// This data type appears in the last step of the upgrade process. After the parachain observes it
+/// and reacts to it the upgrade process concludes.
+#[derive(Encode, Decode, PartialEq, RuntimeDebug)]
+pub enum UpgradeGoAhead {
+	/// Abort the upgrade process. There is something wrong with the validation code previously
+	/// submitted by the parachain. This variant can also be used to prevent upgrades by the governance
+	/// should an emergency emerge.
+	///
+	/// The expected reaction on this variant is that the parachain will admit this message and
+	/// remove all the data about the pending upgrade. Depending on the nature of the problem (to
+	/// be examined offchain for now), it can try to send another validation code or just retry later.
+	#[codec(index = 0)]
+	Abort,
+	/// Apply the pending code change. The parablock that is built on a relay-parent that is descendant
+	/// of the relay-parent where the parachain observed this signal must use the upgraded validation
+	/// code.
+	#[codec(index = 1)]
+	GoAhead,
+}
+
 /// Consensus engine id for polkadot v1 consensus engine.
 pub const POLKADOT_ENGINE_ID: runtime_primitives::ConsensusEngineId = *b"POL1";
 

roadmap/implementers-guide/src/runtime/inclusion.md

@@ -45,7 +45,7 @@ PendingAvailabilityCommitments: map ParaId => CandidateCommitments;
 All failed checks should lead to an unrecoverable error making the block invalid.
 
 * `process_bitfields(expected_bits, Bitfields, core_lookup: Fn(CoreIndex) -> Option<ParaId>)`:
-  1. check that there is at most 1 bitfield per validator and that the number of bits in each bitfield is equal to expected_bits.
+  1. check that there is at most 1 bitfield per validator and that the number of bits in each bitfield is equal to `expected_bits`.
   1. check that there are no duplicates
   1. check all validator signatures.
   1. apply each bit of bitfield to the corresponding pending candidate. looking up parathread cores using the `core_lookup`. Disregard bitfields that have a `1` bit for any free cores.
@@ -70,7 +70,7 @@ All failed checks should lead to an unrecoverable error making the block invalid
   1. create a corresponding entry in the `PendingAvailabilityCommitments` with the commitments.
   1. Return a `Vec<CoreIndex>` of all scheduled cores of the list of passed assignments that a candidate was successfully backed for, sorted ascending by CoreIndex.
 * `enact_candidate(relay_parent_number: BlockNumber, CommittedCandidateReceipt)`:
-  1. If the receipt contains a code upgrade, Call `Paras::schedule_code_upgrade(para_id, code, relay_parent_number + config.validationl_upgrade_delay)`.
+  1. If the receipt contains a code upgrade, Call `Paras::schedule_code_upgrade(para_id, code, relay_parent_number, config)`.
     > TODO: Note that this is safe as long as we never enact candidates where the relay parent is across a session boundary. In that case, which we should be careful to avoid with contextual execution, the configuration might have changed and the para may de-sync from the host's understanding of it.
   1. Reward all backing validators of each candidate, contained within the `backers` field.
   1. call `Ump::receive_upward_messages` for each backed candidate, using the [`UpwardMessage`s](../types/messages.md#upward-message) from the [`CandidateCommitments`](../types/candidate.md#candidate-commitments).

roadmap/implementers-guide/src/runtime/paras.md

@@ -137,6 +137,35 @@ PastCodePruning: Vec<(ParaId, BlockNumber)>;
 FutureCodeUpgrades: map ParaId => Option<BlockNumber>;
 /// The actual future code of a para.
 FutureCodeHash: map ParaId => Option<ValidationCodeHash>;
+/// This is used by the relay-chain to communicate to a parachain a go-ahead with in the upgrade procedure.
+///
+/// This value is absent when there are no upgrades scheduled or during the time the relay chain
+/// performs the checks. It is set at the first relay-chain block when the corresponding parachain
+/// can switch its upgrade function. As soon as the parachain's block is included, the value
+/// gets reset to `None`.
+///
+/// NOTE that this field is used by parachains via merkle storage proofs, therefore changing
+/// the format will require migration of parachains.
+UpgradeGoAheadSignal: map hasher(twox_64_concat) ParaId => Option<UpgradeGoAhead>;
+/// This is used by the relay-chain to communicate that there are restrictions for performing
+/// an upgrade for this parachain.
+///
+/// This may be a because the parachain waits for the upgrade cooldown to expire. Another
+/// potential use case is when we want to perform some maintanance (such as storage migration)
+/// we could restrict upgrades to make the process simpler.
+///
+/// NOTE that this field is used by parachains via merkle storage proofs, therefore changing
+/// the format will require migration of parachains.
+UpgradeRestrictionSignal: map hasher(twox_64_concat) ParaId => Option<UpgradeRestriction>;
+/// The list of parachains that are awaiting for their upgrade restriction to cooldown.
+///
+/// Ordered ascending by block number.
+UpgradeCooldowns: Vec<(ParaId, T::BlockNumber)>;
+/// 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.
+UpcomingUpgrades: Vec<(ParaId, T::BlockNumber)>;
 /// The actions to perform during the start of a specific session index.
 ActionsQueue: map SessionIndex => Vec<ParaId>;
 /// Upcoming paras instantiation arguments.
@@ -171,6 +200,9 @@ CodeByHash: map ValidationCodeHash => Option<ValidationCode>
 
 1. Do pruning based on all entries in `PastCodePruning` with `BlockNumber <= now`. Update the
    corresponding `PastCodeMeta` and `PastCode` accordingly.
+1. Toggle the upgrade related signals
+  1. Collect all `(para_id, expected_at)` from `UpcomingUpgrades` where `expected_at <= now` and prune them. For each para pruned set `UpgradeGoAheadSignal` to `GoAhead`.
+  1. Collect all `(para_id, next_possible_upgrade_at)` from `UpgradeCooldowns` where `next_possible_upgrade_at <= now` and prune them. For each para pruned set `UpgradeRestrictionSignal` to `Present`.
 
 ## Routines
 
@@ -179,12 +211,12 @@ CodeByHash: map ValidationCodeHash => Option<ValidationCode>
 * `schedule_para_cleanup(ParaId)`: Schedule a para to be cleaned up after the next full session.
 * `schedule_parathread_upgrade(ParaId)`: Schedule a parathread to be upgraded to a parachain.
 * `schedule_parachain_downgrade(ParaId)`: Schedule a parachain to be downgraded to a parathread.
-* `schedule_code_upgrade(ParaId, CurrentCode, expected_at: BlockNumber)`: Schedule a future code
+* `schedule_code_upgrade(ParaId, CurrentCode, relay_parent: BlockNumber, HostConfiguration)`: Schedule a future code
   upgrade of the given parachain, to be applied after inclusion of a block of the same parachain
-  executed in the context of a relay-chain block with number >= `expected_at`.
+  executed in the context of a relay-chain block with number >= `relay_parent + config.validation_upgrade_delay`. If the upgrade is scheduled `UpgradeRestrictionSignal` is set and it will remain set until `relay_parent + config.validation_upgrade_frequency`.
 * `note_new_head(ParaId, HeadData, BlockNumber)`: note that a para has progressed to a new head,
   where the new head was executed in the context of a relay-chain block with given number. This will
-  apply pending code upgrades based on the block number provided.
+  apply pending code upgrades based on the block number provided. If an upgrade took place it will clear the `UpgradeGoAheadSignal`.
 * `validation_code_at(ParaId, at: BlockNumber, assume_intermediate: Option<BlockNumber>)`: Fetches
   the validation code to be used when validating a block in the context of the given relay-chain
   height. A second block number parameter may be used to tell the lookup to proceed as if an

runtime/parachains/src/inclusion.rs

@@ -701,7 +701,8 @@ impl<T: Config> Pallet<T> {
 			weight += <paras::Pallet<T>>::schedule_code_upgrade(
 				receipt.descriptor.para_id,
 				new_code,
-				relay_parent_number + config.validation_upgrade_delay,
+				relay_parent_number,
+				&config,
 			);
 		}
 
@@ -2120,9 +2121,19 @@ mod tests {
 					BackingKind::Threshold,
 				));
 
-				Paras::schedule_code_upgrade(chain_a, vec![1, 2, 3, 4].into(), 10);
+				{
+					let cfg = Configuration::config();
+					let expected_at = 10 + cfg.validation_upgrade_delay;
+					assert_eq!(expected_at, 10);
+					Paras::schedule_code_upgrade(
+						chain_a,
+						vec![1, 2, 3, 4].into(),
+						expected_at,
+						&cfg,
+					);
 
-				assert_eq!(Paras::last_code_upgrade(chain_a, true), Some(10));
+					assert_eq!(Paras::last_code_upgrade(chain_a, true), Some(expected_at));
+				}
 
 				assert_eq!(
 					ParaInclusion::process_candidates(

runtime/parachains/src/paras.rs

@@ -28,7 +28,8 @@ use frame_support::pallet_prelude::*;
 use frame_system::pallet_prelude::*;
 use parity_scale_codec::{Decode, Encode};
 use primitives::v1::{
-	ConsensusLog, HeadData, Id as ParaId, SessionIndex, ValidationCode, ValidationCodeHash,
+	ConsensusLog, HeadData, Id as ParaId, SessionIndex, UpgradeGoAhead, UpgradeRestriction,
+	ValidationCode, ValidationCodeHash,
 };
 use sp_core::RuntimeDebug;
 use sp_runtime::{traits::One, DispatchResult, SaturatedConversion};
@@ -374,6 +375,47 @@ pub mod pallet {
 	pub(super) type FutureCodeHash<T: Config> =
 		StorageMap<_, Twox64Concat, ParaId, ValidationCodeHash>;
 
+	/// This is used by the relay-chain to communicate to a parachain a go-ahead with in the upgrade procedure.
+	///
+	/// This value is absent when there are no upgrades scheduled or during the time the relay chain
+	/// performs the checks. It is set at the first relay-chain block when the corresponding parachain
+	/// can switch its upgrade function. As soon as the parachain's block is included, the value
+	/// gets reset to `None`.
+	///
+	/// NOTE that this field is used by parachains via merkle storage proofs, therefore changing
+	/// the format will require migration of parachains.
+	#[pallet::storage]
+	pub(super) type UpgradeGoAheadSignal<T: Config> =
+		StorageMap<_, Twox64Concat, ParaId, UpgradeGoAhead>;
+
+	/// This is used by the relay-chain to communicate that there are restrictions for performing
+	/// an upgrade for this parachain.
+	///
+	/// This may be a because the parachain waits for the upgrade cooldown to expire. Another
+	/// potential use case is when we want to perform some maintenance (such as storage migration)
+	/// we could restrict upgrades to make the process simpler.
+	///
+	/// NOTE that this field is used by parachains via merkle storage proofs, therefore changing
+	/// the format will require migration of parachains.
+	#[pallet::storage]
+	pub(super) type UpgradeRestrictionSignal<T: Config> =
+		StorageMap<_, Twox64Concat, ParaId, UpgradeRestriction>;
+
+	/// The list of parachains that are awaiting for their upgrade restriction to cooldown.
+	///
+	/// Ordered ascending by block number.
+	#[pallet::storage]
+	pub(super) type UpgradeCooldowns<T: Config> =
+		StorageValue<_, Vec<(ParaId, T::BlockNumber)>, 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.
+	///
+	/// Ordered ascending by block number.
+	#[pallet::storage]
+	pub(super) type UpcomingUpgrades<T: Config> =
+		StorageValue<_, Vec<(ParaId, T::BlockNumber)>, ValueQuery>;
+
 	/// The actions to perform during the start of a specific session index.
 	#[pallet::storage]
 	#[pallet::getter(fn actions_queue)]
@@ -478,16 +520,17 @@ pub mod pallet {
 			Ok(())
 		}
 
-		/// Schedule a code upgrade for block `expected_at`.
+		/// Schedule an upgrade as if it was scheduled in the given relay parent block.
 		#[pallet::weight(0)]
 		pub fn force_schedule_code_upgrade(
 			origin: OriginFor<T>,
 			para: ParaId,
 			new_code: ValidationCode,
-			expected_at: T::BlockNumber,
+			relay_parent_number: T::BlockNumber,
 		) -> DispatchResult {
 			ensure_root(origin)?;
-			Self::schedule_code_upgrade(para, new_code, expected_at);
+			let config = configuration::Pallet::<T>::config();
+			Self::schedule_code_upgrade(para, new_code, relay_parent_number, &config);
 			Self::deposit_event(Event::CodeUpgradeScheduled(para));
 			Ok(())
 		}
@@ -527,7 +570,8 @@ pub mod pallet {
 impl<T: Config> Pallet<T> {
 	/// Called by the initializer to initialize the configuration pallet.
 	pub(crate) fn initializer_initialize(now: T::BlockNumber) -> Weight {
-		Self::prune_old_code(now)
+		let weight = Self::prune_old_code(now);
+		weight + Self::process_scheduled_upgrade_changes(now)
 	}
 
 	/// Called by the initializer to finalize the configuration pallet.
@@ -618,6 +662,8 @@ impl<T: Config> Pallet<T> {
 
 					<Self as Store>::Heads::remove(&para);
 					<Self as Store>::FutureCodeUpgrades::remove(&para);
+					<Self as Store>::UpgradeGoAheadSignal::remove(&para);
+					<Self as Store>::UpgradeRestrictionSignal::remove(&para);
 					ParaLifecycles::<T>::remove(&para);
 					let removed_future_code_hash = <Self as Store>::FutureCodeHash::take(&para);
 					if let Some(removed_future_code_hash) = removed_future_code_hash {
@@ -634,6 +680,27 @@ impl<T: Config> Pallet<T> {
 			}
 		}
 
+		if !outgoing.is_empty() {
+			// Filter offboarded parachains from the upcoming upgrades and upgrade cooldowns list.
+			//
+			// We do it after the offboarding to get away with only a single read/write per list.
+			//
+			// 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.
+			<Self as Store>::UpcomingUpgrades::mutate(|upcoming_upgrades| {
+				*upcoming_upgrades = sp_std::mem::take(upcoming_upgrades)
+					.into_iter()
+					.filter(|&(ref para, _)| !outgoing.contains(para))
+					.collect();
+			});
+			<Self as Store>::UpgradeCooldowns::mutate(|upgrade_cooldowns| {
+				*upgrade_cooldowns = sp_std::mem::take(upgrade_cooldowns)
+					.into_iter()
+					.filter(|&(ref para, _)| !outgoing.contains(para))
+					.collect();
+			});
+		}
+
 		// Place the new parachains set in storage.
 		<Self as Store>::Parachains::set(parachains);
 
@@ -727,6 +794,33 @@ impl<T: Config> Pallet<T> {
 		T::DbWeight::get().reads_writes(1 + pruning_tasks_done, 2 * pruning_tasks_done)
 	}
 
+	/// Process the timers related to upgrades. Specifically, the upgrade go ahead signals toggle
+	/// and the upgrade cooldown restrictions.
+	///
+	/// Takes the current block number and returns the weight consumed.
+	fn process_scheduled_upgrade_changes(now: T::BlockNumber) -> Weight {
+		let upgrades_signaled = <Self as Store>::UpcomingUpgrades::mutate(
+			|upcoming_upgrades: &mut Vec<(ParaId, T::BlockNumber)>| {
+				let num = upcoming_upgrades.iter().take_while(|&(_, at)| at <= &now).count();
+				for (para, _) in upcoming_upgrades.drain(..num) {
+					<Self as Store>::UpgradeGoAheadSignal::insert(&para, UpgradeGoAhead::GoAhead);
+				}
+				num
+			},
+		);
+		let cooldowns_expired = <Self as Store>::UpgradeCooldowns::mutate(
+			|upgrade_cooldowns: &mut Vec<(ParaId, T::BlockNumber)>| {
+				let num = upgrade_cooldowns.iter().take_while(|&(_, at)| at <= &now).count();
+				for (para, _) in upgrade_cooldowns.drain(..num) {
+					<Self as Store>::UpgradeRestrictionSignal::remove(&para);
+				}
+				num
+			},
+		);
+
+		T::DbWeight::get().reads_writes(2, upgrades_signaled as u64 + cooldowns_expired as u64)
+	}
+
 	/// Verify that `schedule_para_initialize` can be called successfully.
 	///
 	/// Returns false if para is already registered in the system.
@@ -830,14 +924,37 @@ impl<T: Config> Pallet<T> {
 	pub(crate) fn schedule_code_upgrade(
 		id: ParaId,
 		new_code: ValidationCode,
-		expected_at: T::BlockNumber,
+		relay_parent_number: T::BlockNumber,
+		cfg: &configuration::HostConfiguration<T::BlockNumber>,
 	) -> Weight {
 		<Self as Store>::FutureCodeUpgrades::mutate(&id, |up| {
 			if up.is_some() {
 				T::DbWeight::get().reads_writes(1, 0)
 			} else {
+				let expected_at = relay_parent_number + cfg.validation_upgrade_delay;
+				let next_possible_upgrade_at =
+					relay_parent_number + cfg.validation_upgrade_frequency;
+
 				*up = Some(expected_at);
 
+				<Self as Store>::UpcomingUpgrades::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));
+				});
+
+				// From the moment of signalling of the upgrade until the cooldown expires, the
+				// parachain is disallowed to make further upgrades. Therefore set the upgrade
+				// permission signal to disallowed and activate the cooldown timer.
+				<Self as Store>::UpgradeRestrictionSignal::insert(&id, UpgradeRestriction::Present);
+				<Self as Store>::UpgradeCooldowns::mutate(|upgrade_cooldowns| {
+					let insert_idx = upgrade_cooldowns
+						.binary_search_by_key(&next_possible_upgrade_at, |&(_, b)| b)
+						.unwrap_or_else(|idx| idx);
+					upgrade_cooldowns.insert(insert_idx, (id, next_possible_upgrade_at));
+				});
+
 				let new_code_hash = new_code.hash();
 				let expected_at_u32 = expected_at.saturated_into();
 				let log = ConsensusLog::ParaScheduleUpgradeCode(id, new_code_hash, expected_at_u32);
@@ -845,14 +962,14 @@ impl<T: Config> Pallet<T> {
 
 				let (reads, writes) = Self::increase_code_ref(&new_code_hash, &new_code);
 				FutureCodeHash::<T>::insert(&id, new_code_hash);
-				T::DbWeight::get().reads_writes(1 + reads, 2 + writes)
+				T::DbWeight::get().reads_writes(2 + reads, 3 + writes)
 			}
 		})
 	}
 
 	/// Note that a para has progressed to a new head, where the new head was executed in the context
 	/// of a relay-chain block with given number. This will apply pending code upgrades based
-	/// on the block number provided.
+	/// on the relay-parent block number provided.
 	pub(crate) fn note_new_head(
 		id: ParaId,
 		new_head: HeadData,
@@ -863,6 +980,7 @@ impl<T: Config> Pallet<T> {
 		if let Some(expected_at) = <Self as Store>::FutureCodeUpgrades::get(&id) {
 			if expected_at <= execution_context {
 				<Self as Store>::FutureCodeUpgrades::remove(&id);
+				<Self as Store>::UpgradeGoAheadSignal::remove(&id);
 
 				// Both should always be `Some` in this case, since a code upgrade is scheduled.
 				let new_code_hash = FutureCodeHash::<T>::take(&id).unwrap_or_default();
@@ -1025,7 +1143,7 @@ mod tests {
 
 	use crate::{
 		configuration::HostConfiguration,
-		mock::{new_test_ext, MockGenesisConfig, Paras, ParasShared, System},
+		mock::{new_test_ext, Configuration, MockGenesisConfig, Paras, ParasShared, System},
 	};
 
 	fn run_to_block(to: BlockNumber, new_session: Option<Vec<BlockNumber>>) {
@@ -1332,6 +1450,7 @@ mod tests {
 	fn code_upgrade_applied_after_delay() {
 		let code_retention_period = 10;
 		let validation_upgrade_delay = 5;
+		let validation_upgrade_frequency = 10;
 
 		let original_code = ValidationCode(vec![1, 2, 3]);
 		let paras = vec![(
@@ -1349,6 +1468,7 @@ mod tests {
 				config: HostConfiguration {
 					code_retention_period,
 					validation_upgrade_delay,
+					validation_upgrade_frequency,
 					..Default::default()
 				},
 				..Default::default()
@@ -1368,12 +1488,23 @@ mod tests {
 			let expected_at = {
 				// this parablock is in the context of block 1.
 				let expected_at = 1 + validation_upgrade_delay;
-				Paras::schedule_code_upgrade(para_id, new_code.clone(), expected_at);
+				let next_possible_upgrade_at = 1 + validation_upgrade_frequency;
+				Paras::schedule_code_upgrade(
+					para_id,
+					new_code.clone(),
+					1,
+					&Configuration::config(),
+				);
 				Paras::note_new_head(para_id, Default::default(), 1);
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
 				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(<Paras as Store>::UpcomingUpgrades::get(), vec![(para_id, expected_at)]);
+				assert_eq!(
+					<Paras as Store>::UpgradeCooldowns::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);
@@ -1391,6 +1522,10 @@ mod tests {
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
 				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(
+					<Paras as Store>::UpgradeGoAheadSignal::get(&para_id),
+					Some(UpgradeGoAhead::GoAhead)
+				);
 				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 				check_code_is_stored(&original_code);
 				check_code_is_stored(&new_code);
@@ -1410,6 +1545,7 @@ mod tests {
 				);
 				assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
 				assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
+				assert!(<Paras as Store>::UpgradeGoAheadSignal::get(&para_id).is_none());
 				assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
 				check_code_is_stored(&original_code);
 				check_code_is_stored(&new_code);
@@ -1421,6 +1557,7 @@ mod tests {
 	fn code_upgrade_applied_after_delay_even_when_late() {
 		let code_retention_period = 10;
 		let validation_upgrade_delay = 5;
+		let validation_upgrade_frequency = 10;
 
 		let original_code = ValidationCode(vec![1, 2, 3]);
 		let paras = vec![(
@@ -1438,6 +1575,7 @@ mod tests {
 				config: HostConfiguration {
 					code_retention_period,
 					validation_upgrade_delay,
+					validation_upgrade_frequency,
 					..Default::default()
 				},
 				..Default::default()
@@ -1455,12 +1593,24 @@ mod tests {
 			let expected_at = {
 				// this parablock is in the context of block 1.
 				let expected_at = 1 + validation_upgrade_delay;
-				Paras::schedule_code_upgrade(para_id, new_code.clone(), expected_at);
+				let next_possible_upgrade_at = 1 + validation_upgrade_frequency;
+				Paras::schedule_code_upgrade(
+					para_id,
+					new_code.clone(),
+					1,
+					&Configuration::config(),
+				);
 				Paras::note_new_head(para_id, Default::default(), 1);
 
 				assert!(Paras::past_code_meta(&para_id).most_recent_change().is_none());
 				assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(expected_at));
 				assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
+				assert_eq!(<Paras as Store>::UpcomingUpgrades::get(), vec![(para_id, expected_at)]);
+				assert_eq!(
+					<Paras as Store>::UpgradeCooldowns::get(),
+					vec![(para_id, next_possible_upgrade_at)]
+				);
+				assert!(<Paras as Store>::UpgradeGoAheadSignal::get(&para_id).is_none());
 				assert_eq!(Paras::current_code(&para_id), Some(original_code.clone()));
 
 				expected_at
@@ -1471,6 +1621,12 @@ mod tests {
 			// the candidate is in the context of the first descendant of `expected_at`, and triggers
 			// the upgrade.
 			{
+				// The signal should be set to go-ahead until the new head is actually processed.
+				assert_eq!(
+					<Paras as Store>::UpgradeGoAheadSignal::get(&para_id),
+					Some(UpgradeGoAhead::GoAhead),
+				);
+
 				Paras::note_new_head(para_id, Default::default(), expected_at + 4);
 
 				assert_eq!(Paras::past_code_meta(&para_id).most_recent_change(), Some(expected_at),);
@@ -1502,6 +1658,7 @@ mod tests {
 				);
 				assert!(<Paras as Store>::FutureCodeUpgrades::get(&para_id).is_none());
 				assert!(<Paras as Store>::FutureCodeHash::get(&para_id).is_none());
+				assert!(<Paras as Store>::UpgradeGoAheadSignal::get(&para_id).is_none());
 				assert_eq!(Paras::current_code(&para_id), Some(new_code.clone()));
 			}
 		});
@@ -1510,6 +1667,7 @@ mod tests {
 	#[test]
 	fn submit_code_change_when_not_allowed_is_err() {
 		let code_retention_period = 10;
+		let validation_upgrade_delay = 7;
 
 		let paras = vec![(
 			0u32.into(),
@@ -1523,7 +1681,11 @@ mod tests {
 		let genesis_config = MockGenesisConfig {
 			paras: GenesisConfig { paras, ..Default::default() },
 			configuration: crate::configuration::GenesisConfig {
-				config: HostConfiguration { code_retention_period, ..Default::default() },
+				config: HostConfiguration {
+					code_retention_period,
+					validation_upgrade_delay,
+					..Default::default()
+				},
 				..Default::default()
 			},
 			..Default::default()
@@ -1535,14 +1697,22 @@ mod tests {
 			let newer_code = ValidationCode(vec![4, 5, 6, 7]);
 
 			run_to_block(1, None);
-
-			Paras::schedule_code_upgrade(para_id, new_code.clone(), 8);
-			assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(8));
+			Paras::schedule_code_upgrade(para_id, new_code.clone(), 1, &Configuration::config());
+			assert_eq!(
+				<Paras as Store>::FutureCodeUpgrades::get(&para_id),
+				Some(1 + validation_upgrade_delay)
+			);
 			assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
 			check_code_is_stored(&new_code);
 
-			Paras::schedule_code_upgrade(para_id, newer_code.clone(), 10);
-			assert_eq!(<Paras as Store>::FutureCodeUpgrades::get(&para_id), Some(8));
+			// We expect that if an upgrade is signalled while there is already one pending we just
+			// ignore it. Note that this is only true from perspective of this module.
+			run_to_block(2, None);
+			Paras::schedule_code_upgrade(para_id, newer_code.clone(), 2, &Configuration::config());
+			assert_eq!(
+				<Paras as Store>::FutureCodeUpgrades::get(&para_id),
+				Some(1 + validation_upgrade_delay), // did not change since the same assertion from the last time.
+			);
 			assert_eq!(<Paras as Store>::FutureCodeHash::get(&para_id), Some(new_code.hash()));
 			check_code_is_not_stored(&newer_code);
 		});
@@ -1551,6 +1721,7 @@ mod tests {
 	#[test]
 	fn full_parachain_cleanup_storage() {
 		let code_retention_period = 10;
+		let validation_upgrade_delay = 1 + 5;
 
 		let original_code = ValidationCode(vec![1, 2, 3]);
 		let paras = vec![(
@@ -1565,7 +1736,11 @@ mod tests {
 		let genesis_config = MockGenesisConfig {
 			paras: GenesisConfig { paras, ..Default::default() },
 			configuration: crate::configuration::GenesisConfig {
-				config: HostConfiguration { code_retention_period, ..Default::default() },
+				config: HostConfiguration {
+					code_retention_period,
+					validation_upgrade_delay,