fix: GoAhead signal only set when runtime upgrade is enacted from parachain side (#1176)

The runtime code of a parachain can be replaced on the relay-chain via:

[cumulus]:
[enact_authorized_upgrade](https://github.com/paritytech/polkadot-sdk/blob/1a38d6d6be42b30e8be3ffccec75a4ec995fef9d/cumulus/pallets/parachain-system/src/lib.rs#L661);
this is used for a runtime upgrade when a parachain is not bricked.

[polkadot] (these are used when a parachain is bricked):
-
[force_set_current_code](https://github.com/paritytech/polkadot-sdk/blob/1a38d6d6be42b30e8be3ffccec75a4ec995fef9d/polkadot/runtime/parachains/src/paras/mod.rs#L823):
immediately changes the runtime code of a given para without a pvf check
(root).
-
[force_schedule_code_upgrade](https://github.com/paritytech/polkadot-sdk/blob/1a38d6d6be42b30e8be3ffccec75a4ec995fef9d/polkadot/runtime/parachains/src/paras/mod.rs#L864):
schedules a change to the runtime code of a given para including a pvf
check of the new code (root).
-
[schedule_code_upgrade](https://github.com/paritytech/polkadot-sdk/blob/1a38d6d6be42b30e8be3ffccec75a4ec995fef9d/polkadot/runtime/common/src/paras_registrar.rs#L395):
schedules a change to the runtime code of a given para including a pvf
check of the new code. Besides root, the parachain or parachain manager
can call this extrinsic given that the parachain is unlocked.

Polkadot signals a parachain to be ready for a runtime upgrade through
the
[GoAhead](https://github.com/paritytech/polkadot-sdk/blob/e49493442a9377be9344c06a4990e17423783d41/polkadot/primitives/src/v5/mod.rs#L1229)
signal.

When in cumulus `enact_authorized_upgrade` is executed, the same
underlying helper function of `force_schedule_code_upgrade` &
`schedule_code_upgrade`:
[schedule_code_upgrade](https://github.com/paritytech/polkadot/blob/09b61286da11921a3dda0a8e4015ceb9ef9cffca/runtime/parachains/src/paras/mod.rs#L1778),
is called on the relay-chain, which sets the `GoAhead` signal (if the
pvf is accepted).

If Cumulus receives the `GoAhead` signal from polkadot without having
the `PendingValidationCode` ready, it will panic
([ref](https://github.com/paritytech/polkadot/pull/7412)). For
`enact_authorized_upgrade` we know for sure the `PendingValidationCode`
is set. On the contrary, for `force_schedule_code_upgrade` &
`schedule_code_upgrade` this is not the case.

This PR includes a flag such that the `GoAhead` signal will only be set
when a runtime upgrade is enacted by the parachain
(`enact_authorized_upgrade`).

additional info: https://github.com/paritytech/polkadot/pull/7412

Closes #641

---------

Co-authored-by: Bastian Köcher <git@kchr.de>
Co-authored-by: Bastian Köcher <info@kchr.de>

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};
 use runtime_parachains::{
 	configuration, ensure_parachain,
-	paras::{self, ParaGenesisArgs},
+	paras::{self, ParaGenesisArgs, SetGoAhead},
 	Origin, ParaLifecycle,
 };
 use sp_std::{prelude::*, result};
@@ -412,7 +412,7 @@ pub mod pallet {
 			new_code: ValidationCode,
 		) -> DispatchResult {
 			Self::ensure_root_para_or_owner(origin, para)?;
-			runtime_parachains::schedule_code_upgrade::<T>(para, new_code)?;
+			runtime_parachains::schedule_code_upgrade::<T>(para, new_code, SetGoAhead::No)?;
 			Ok(())
 		}
 

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

@@ -21,7 +21,8 @@
 
 use crate::{
 	configuration::{self, HostConfiguration},
-	disputes, dmp, hrmp, paras,
+	disputes, dmp, hrmp,
+	paras::{self, SetGoAhead},
 	scheduler::{self, AvailabilityTimeoutStatus},
 	shared::{self, AllowedRelayParentsTracker},
 };
@@ -448,8 +449,9 @@ impl fmt::Debug for UmpAcceptanceCheckErr {
 				"the ump queue would have grown past the max size permitted by config ({} > {})",
 				total_size, limit,
 			),
-			UmpAcceptanceCheckErr::IsOffboarding =>
-				write!(fmt, "upward message rejected because the para is off-boarding",),
+			UmpAcceptanceCheckErr::IsOffboarding => {
+				write!(fmt, "upward message rejected because the para is off-boarding")
+			},
 		}
 	}
 }
@@ -885,6 +887,7 @@ impl<T: Config> Pallet<T> {
 				new_code,
 				now,
 				&config,
+				SetGoAhead::Yes,
 			));
 		}
 

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

@@ -1255,7 +1255,13 @@ fn candidate_checks() {
 				let cfg = Configuration::config();
 				let expected_at = 10 + cfg.validation_upgrade_delay;
 				assert_eq!(expected_at, 12);
-				Paras::schedule_code_upgrade(chain_a, vec![1, 2, 3, 4].into(), expected_at, &cfg);
+				Paras::schedule_code_upgrade(
+					chain_a,
+					vec![1, 2, 3, 4].into(),
+					expected_at,
+					&cfg,
+					SetGoAhead::Yes,
+				);
 			}
 
 			assert_noop!(
@@ -2235,7 +2241,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 }
+			paras::PvfCheckCause::Upgrade { id, included_at, set_go_ahead: SetGoAhead::Yes }
 				if id == &chain_a && included_at == &7
 		);
 	});

polkadot/runtime/parachains/src/lib.rs

@@ -53,7 +53,7 @@ mod mock;
 mod ump_tests;
 
 pub use origin::{ensure_parachain, Origin};
-pub use paras::ParaLifecycle;
+pub use paras::{ParaLifecycle, SetGoAhead};
 use primitives::{HeadData, Id as ParaId, ValidationCode};
 use sp_runtime::{DispatchResult, FixedU128};
 
@@ -89,8 +89,9 @@ 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,
 ) -> DispatchResult {
-	paras::Pallet::<T>::schedule_code_upgrade_external(id, new_code)
+	paras::Pallet::<T>::schedule_code_upgrade_external(id, new_code, set_go_ahead)
 }
 
 /// Sets the current parachain head with the given id.

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

@@ -124,7 +124,13 @@ benchmarks! {
 		let expired = frame_system::Pallet::<T>::block_number().saturating_sub(One::one());
 		let config = HostConfiguration::<BlockNumberFor<T>>::default();
 		generate_disordered_pruning::<T>();
-		Pallet::<T>::schedule_code_upgrade(para_id, ValidationCode(vec![0]), expired, &config);
+		Pallet::<T>::schedule_code_upgrade(
+			para_id,
+			ValidationCode(vec![0]),
+			expired,
+			&config,
+			SetGoAhead::Yes,
+		);
 	}: _(RawOrigin::Root, para_id, new_head)
 	verify {
 		assert_last_event::<T>(Event::NewHeadNoted(para_id).into());

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

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

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

@@ -386,9 +386,18 @@ 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,
 	},
 }
 
+/// Should the `GoAhead` signal be set after a successful check of the new wasm binary?
+#[derive(Debug, Copy, Clone, PartialEq, TypeInfo, Decode, Encode)]
+pub enum SetGoAhead {
+	Yes,
+	No,
+}
+
 impl<BlockNumber> PvfCheckCause<BlockNumber> {
 	/// Returns the ID of the para that initiated or subscribed to the pre-checking vote.
 	fn para_id(&self) -> ParaId {
@@ -888,7 +897,13 @@ pub mod pallet {
 		) -> DispatchResult {
 			ensure_root(origin)?;
 			let config = configuration::Pallet::<T>::config();
-			Self::schedule_code_upgrade(para, new_code, relay_parent_number, &config);
+			Self::schedule_code_upgrade(
+				para,
+				new_code,
+				relay_parent_number,
+				&config,
+				SetGoAhead::No,
+			);
 			Self::deposit_event(Event::CodeUpgradeScheduled(para));
 			Ok(())
 		}
@@ -1186,6 +1201,7 @@ impl<T: Config> Pallet<T> {
 	pub(crate) fn schedule_code_upgrade_external(
 		id: ParaId,
 		new_code: ValidationCode,
+		set_go_ahead: SetGoAhead,
 	) -> DispatchResult {
 		// Check that we can schedule an upgrade at all.
 		ensure!(Self::can_upgrade_validation_code(id), Error::<T>::CannotUpgradeCode);
@@ -1193,7 +1209,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);
+		Self::schedule_code_upgrade(id, new_code, upgrade_block, &config, set_go_ahead);
 		Self::deposit_event(Event::CodeUpgradeScheduled(id));
 		Ok(())
 	}
@@ -1534,8 +1550,15 @@ impl<T: Config> Pallet<T> {
 				PvfCheckCause::Onboarding(id) => {
 					weight += Self::proceed_with_onboarding(*id, sessions_observed);
 				},
-				PvfCheckCause::Upgrade { id, included_at } => {
-					weight += Self::proceed_with_upgrade(*id, code_hash, now, *included_at, cfg);
+				PvfCheckCause::Upgrade { id, included_at, set_go_ahead } => {
+					weight += Self::proceed_with_upgrade(
+						*id,
+						code_hash,
+						now,
+						*included_at,
+						cfg,
+						*set_go_ahead,
+					);
 				},
 			}
 		}
@@ -1568,6 +1591,7 @@ impl<T: Config> Pallet<T> {
 		now: BlockNumberFor<T>,
 		relay_parent_number: BlockNumberFor<T>,
 		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
+		set_go_ahead: SetGoAhead,
 	) -> Weight {
 		let mut weight = Weight::zero();
 
@@ -1591,12 +1615,15 @@ impl<T: Config> Pallet<T> {
 		weight += T::DbWeight::get().reads_writes(1, 4);
 		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));
-		});
+		// 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));
+			});
+		}
 
 		let expected_at = expected_at.saturated_into();
 		let log = ConsensusLog::ParaScheduleUpgradeCode(id, *code_hash, expected_at);
@@ -1835,6 +1862,7 @@ impl<T: Config> Pallet<T> {
 		new_code: ValidationCode,
 		inclusion_block_number: BlockNumberFor<T>,
 		cfg: &configuration::HostConfiguration<BlockNumberFor<T>>,
+		set_go_ahead: SetGoAhead,
 	) -> Weight {
 		let mut weight = T::DbWeight::get().reads(1);
 
@@ -1884,7 +1912,7 @@ impl<T: Config> Pallet<T> {
 		});
 
 		weight += Self::kick_off_pvf_check(
-			PvfCheckCause::Upgrade { id, included_at: inclusion_block_number },
+			PvfCheckCause::Upgrade { id, included_at: inclusion_block_number, set_go_ahead },
 			code_hash,
 			new_code,
 			cfg,

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

@@ -436,7 +436,13 @@ fn code_upgrade_applied_after_delay() {
 			// 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;
-			Paras::schedule_code_upgrade(para_id, new_code.clone(), 1, &Configuration::config());
+			Paras::schedule_code_upgrade(
+				para_id,
+				new_code.clone(),
+				1,
+				&Configuration::config(),
+				SetGoAhead::Yes,
+			);
 			// Include votes for super-majority.
 			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
 
@@ -504,6 +510,127 @@ fn code_upgrade_applied_after_delay() {
 	});
 }
 
+#[test]
+fn code_upgrade_applied_without_setting_go_ahead_signal() {
+	let code_retention_period = 10;
+	let validation_upgrade_delay = 5;
+	let validation_upgrade_cooldown = 10;
+
+	let original_code = ValidationCode(vec![1, 2, 3]);
+	let paras = vec![(
+		0u32.into(),
+		ParaGenesisArgs {
+			para_kind: ParaKind::Parachain,
+			genesis_head: dummy_head_data(),
+			validation_code: original_code.clone(),
+		},
+	)];
+
+	let genesis_config = MockGenesisConfig {
+		paras: GenesisConfig { paras, ..Default::default() },
+		configuration: crate::configuration::GenesisConfig {
+			config: HostConfiguration {
+				code_retention_period,
+				validation_upgrade_delay,
+				validation_upgrade_cooldown,
+				..Default::default()
+			},
+		},
+		..Default::default()
+	};
+
+	new_test_ext(genesis_config).execute_with(|| {
+		check_code_is_stored(&original_code);
+
+		let para_id = ParaId::from(0);
+		let new_code = ValidationCode(vec![4, 5, 6]);
+
+		// Wait for at least one session change to set active validators.
+		const EXPECTED_SESSION: SessionIndex = 1;
+		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)
+		};
+
+		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);
+
+			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);
+		}
+
+		run_to_block(next_possible_upgrade_at + 1, None);
+
+		{
+			assert!(UpgradeRestrictionSignal::<Test>::get(&para_id).is_none());
+			assert!(UpgradeCooldowns::<Test>::get().is_empty());
+		}
+	});
+}
+
 #[test]
 fn code_upgrade_applied_after_delay_even_when_late() {
 	let code_retention_period = 10;
@@ -546,7 +673,13 @@ fn code_upgrade_applied_after_delay_even_when_late() {
 			// 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;
-			Paras::schedule_code_upgrade(para_id, new_code.clone(), 1, &Configuration::config());
+			Paras::schedule_code_upgrade(
+				para_id,
+				new_code.clone(),
+				1,
+				&Configuration::config(),
+				SetGoAhead::Yes,
+			);
 			// Include votes for super-majority.
 			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
 
@@ -624,7 +757,13 @@ fn submit_code_change_when_not_allowed_is_err() {
 		const EXPECTED_SESSION: SessionIndex = 1;
 		run_to_block(1, Some(vec![1]));
 
-		Paras::schedule_code_upgrade(para_id, new_code.clone(), 1, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			new_code.clone(),
+			1,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		// Include votes for super-majority.
 		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
 
@@ -636,7 +775,13 @@ fn submit_code_change_when_not_allowed_is_err() {
 		// ignore it. Note that this is only true from perspective of this module.
 		run_to_block(2, None);
 		assert!(!Paras::can_upgrade_validation_code(para_id));
-		Paras::schedule_code_upgrade(para_id, newer_code.clone(), 2, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			newer_code.clone(),
+			2,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		assert_eq!(
 			FutureCodeUpgrades::<Test>::get(&para_id),
 			Some(1 + validation_upgrade_delay), /* did not change since the same assertion from
@@ -694,7 +839,13 @@ fn upgrade_restriction_elapsed_doesnt_mean_can_upgrade() {
 		const EXPECTED_SESSION: SessionIndex = 1;
 		run_to_block(1, Some(vec![1]));
 
-		Paras::schedule_code_upgrade(para_id, new_code.clone(), 0, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			new_code.clone(),
+			0,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		// Include votes for super-majority.
 		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
 
@@ -713,7 +864,13 @@ fn upgrade_restriction_elapsed_doesnt_mean_can_upgrade() {
 		assert!(!Paras::can_upgrade_validation_code(para_id));
 
 		// And scheduling another upgrade does not do anything. `expected_at` is still the same.
-		Paras::schedule_code_upgrade(para_id, newer_code.clone(), 30, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			newer_code.clone(),
+			30,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		assert_eq!(FutureCodeUpgrades::<Test>::get(&para_id), Some(0 + validation_upgrade_delay));
 	});
 }
@@ -765,7 +922,13 @@ fn full_parachain_cleanup_storage() {
 		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(), 1, &Configuration::config());
+			Paras::schedule_code_upgrade(
+				para_id,
+				new_code.clone(),
+				1,
+				&Configuration::config(),
+				SetGoAhead::Yes,
+			);
 			// Include votes for super-majority.
 			submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
 
@@ -860,6 +1023,7 @@ fn cannot_offboard_ongoing_pvf_check() {
 			new_code.clone(),
 			RELAY_PARENT,
 			&Configuration::config(),
+			SetGoAhead::Yes,
 		);
 		assert!(!Paras::pvfs_require_precheck().is_empty());
 
@@ -1012,7 +1176,13 @@ fn code_hash_at_returns_up_to_end_of_code_retention_period() {
 		let para_id = ParaId::from(0);
 		let old_code: ValidationCode = vec![1, 2, 3].into();
 		let new_code: ValidationCode = vec![4, 5, 6].into();
-		Paras::schedule_code_upgrade(para_id, new_code.clone(), 0, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			new_code.clone(),
+			0,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		// Include votes for super-majority.
 		submit_super_majority_pvf_votes(&new_code, EXPECTED_SESSION, true);
 
@@ -1120,6 +1290,7 @@ fn pvf_check_coalescing_onboarding_and_upgrade() {
 			validation_code.clone(),
 			RELAY_PARENT,
 			&Configuration::config(),
+			SetGoAhead::Yes,
 		);
 		assert!(!Paras::pvfs_require_precheck().is_empty());
 
@@ -1224,7 +1395,13 @@ fn pvf_check_upgrade_reject() {
 		// Expected current session index.
 		const EXPECTED_SESSION: SessionIndex = 1;
 
-		Paras::schedule_code_upgrade(a, new_code.clone(), RELAY_PARENT, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			a,
+			new_code.clone(),
+			RELAY_PARENT,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		check_code_is_stored(&new_code);
 
 		// 1/3 of validators vote against `new_code`. PVF should not be rejected yet.
@@ -1404,7 +1581,13 @@ fn include_pvf_check_statement_refunds_weight() {
 		// Expected current session index.
 		const EXPECTED_SESSION: SessionIndex = 1;
 
-		Paras::schedule_code_upgrade(a, new_code.clone(), RELAY_PARENT, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			a,
+			new_code.clone(),
+			RELAY_PARENT,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 
 		let mut stmts = IntoIterator::into_iter([0, 1, 2, 3])
 			.map(|i| {
@@ -1499,7 +1682,13 @@ fn poke_unused_validation_code_doesnt_remove_code_with_users() {
 
 		// Then we add a user to the code, say by upgrading.
 		run_to_block(2, None);
-		Paras::schedule_code_upgrade(para_id, validation_code.clone(), 1, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			validation_code.clone(),
+			1,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		Paras::note_new_head(para_id, HeadData::default(), 1);
 
 		// Finally we poke the code, which should not remove it from the storage.
@@ -1564,7 +1753,13 @@ fn add_trusted_validation_code_insta_approval() {
 
 		// Then some parachain upgrades it's code with the relay-parent 1.
 		run_to_block(2, None);
-		Paras::schedule_code_upgrade(para_id, validation_code.clone(), 1, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			validation_code.clone(),
+			1,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		Paras::note_new_head(para_id, HeadData::default(), 1);
 
 		// Verify that the code upgrade has `expected_at` set to `26`.
@@ -1600,7 +1795,13 @@ fn add_trusted_validation_code_enacts_existing_pvf_vote() {
 	new_test_ext(genesis_config).execute_with(|| {
 		// First, some parachain upgrades it's code with the relay-parent 1.
 		run_to_block(2, None);
-		Paras::schedule_code_upgrade(para_id, validation_code.clone(), 1, &Configuration::config());
+		Paras::schedule_code_upgrade(
+			para_id,
+			validation_code.clone(),
+			1,
+			&Configuration::config(),
+			SetGoAhead::Yes,
+		);
 		Paras::note_new_head(para_id, HeadData::default(), 1);
 
 		// No upgrade should be scheduled at this point. PVF pre-checking vote should run for