Initialize on-chain StorageVersion for pallets added after genesis

frame/executive/src/lib.rs

@@ -364,7 +364,9 @@ where
 			)?;
 		}
 
-		let weight =
+		let before_all_weight =
+			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::before_all();
+		let try_on_runtime_upgrade_weight =
 			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::try_on_runtime_upgrade(
 				checks.pre_and_post(),
 			)?;
@@ -379,7 +381,7 @@ where
 			)?;
 		}
 
-		Ok(weight)
+		Ok(before_all_weight.saturating_add(try_on_runtime_upgrade_weight))
 	}
 }
 
@@ -408,7 +410,10 @@ where
 {
 	/// Execute all `OnRuntimeUpgrade` of this runtime, and return the aggregate weight.
 	pub fn execute_on_runtime_upgrade() -> Weight {
+		let before_all_weight =
+			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::before_all();
 		<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::on_runtime_upgrade()
+			.saturating_add(before_all_weight)
 	}
 
 	/// Start the execution of a particular block.

frame/support/procedural/src/pallet/expand/hooks.rs

@@ -34,6 +34,38 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 	let type_use_gen = &def.type_use_generics(span);
 	let pallet_ident = &def.pallet_struct.pallet;
 	let frame_system = &def.frame_system;
+	let pallet_name = quote::quote! {
+		<
+			<T as #frame_system::Config>::PalletInfo
+			as
+			#frame_support::traits::PalletInfo
+		>::name::<Self>().unwrap_or("<unknown pallet name>")
+	};
+
+	let initialize_on_chain_storage_version = if let Some(current_version) =
+		&def.pallet_struct.storage_version
+	{
+		quote::quote! {
+			#frame_support::log::info!(
+				target: #frame_support::LOG_TARGET,
+				"🐥 Pallet {:?} has no on-chain storage version. Initializing the on-chain storage version to match the storage version defined in the pallet: {:?}",
+				#pallet_name,
+				#current_version
+			);
+			#current_version.put::<Self>();
+		}
+	} else {
+		quote::quote! {
+			let default_version = #frame_support::traits::StorageVersion::new(0);
+			#frame_support::log::info!(
+				target: #frame_support::LOG_TARGET,
+				"🐥 Pallet {:?} has no on-chain storage version. The pallet has not defined storage version, so the on-chain version is being initialized to {:?}.",
+				#pallet_name,
+				default_version
+			);
+			default_version.put::<Self>();
+		}
+	};
 
 	let log_runtime_upgrade = if has_runtime_upgrade {
 		// a migration is defined here.
@@ -42,7 +74,7 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 				target: #frame_support::LOG_TARGET,
 				"⚠️ {} declares internal migrations (which *might* execute). \
 				 On-chain `{:?}` vs current storage version `{:?}`",
-				pallet_name,
+				#pallet_name,
 				<Self as #frame_support::traits::GetStorageVersion>::on_chain_storage_version(),
 				<Self as #frame_support::traits::GetStorageVersion>::current_storage_version(),
 			);
@@ -53,21 +85,16 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			#frame_support::log::debug!(
 				target: #frame_support::LOG_TARGET,
 				"✅ no migration for {}",
-				pallet_name,
+				#pallet_name,
 			);
 		}
 	};
 
 	let log_try_state = quote::quote! {
-		let pallet_name = <
-			<T as #frame_system::Config>::PalletInfo
-			as
-			#frame_support::traits::PalletInfo
-		>::name::<Self>().expect("No name found for the pallet! This usually means that the pallet wasn't added to `construct_runtime!`.");
 		#frame_support::log::debug!(
 			target: #frame_support::LOG_TARGET,
 			"🩺 try-state pallet {:?}",
-			pallet_name,
+			#pallet_name,
 		);
 	};
 
@@ -91,16 +118,10 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			let current_version = <Self as #frame_support::traits::GetStorageVersion>::current_storage_version();
 
 			if on_chain_version != current_version {
-				let pallet_name = <
-					<T as #frame_system::Config>::PalletInfo
-					as
-					#frame_support::traits::PalletInfo
-				>::name::<Self>().unwrap_or("<unknown pallet name>");
-
 				#frame_support::log::error!(
 					target: #frame_support::LOG_TARGET,
 					"{}: On chain storage version {:?} doesn't match current storage version {:?}.",
-					pallet_name,
+					#pallet_name,
 					on_chain_version,
 					current_version,
 				);
@@ -113,17 +134,11 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			let on_chain_version = <Self as #frame_support::traits::GetStorageVersion>::on_chain_storage_version();
 
 			if on_chain_version != #frame_support::traits::StorageVersion::new(0) {
-				let pallet_name = <
-					<T as #frame_system::Config>::PalletInfo
-					as
-					#frame_support::traits::PalletInfo
-				>::name::<Self>().unwrap_or("<unknown pallet name>");
-
 				#frame_support::log::error!(
 					target: #frame_support::LOG_TARGET,
 					"{}: On chain storage version {:?} is set to non zero, \
 					 while the pallet is missing the `#[pallet::storage_version(VERSION)]` attribute.",
-					pallet_name,
+					#pallet_name,
 					on_chain_version,
 				);
 
@@ -190,17 +205,28 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			#frame_support::traits::OnRuntimeUpgrade
 			for #pallet_ident<#type_use_gen> #where_clause
 		{
+			fn before_all() -> #frame_support::weights::Weight {
+				use #frame_support::traits::Get;
+				#frame_support::sp_tracing::enter_span!(
+					#frame_support::sp_tracing::trace_span!("before_all")
+				);
+
+				// Check if the on-chain version has been set yet
+				let exists = #frame_support::traits::StorageVersion::exists::<Self>();
+				if !exists {
+					#initialize_on_chain_storage_version
+					<T as #frame_system::Config>::DbWeight::get().reads_writes(1, 1)
+				} else {
+					<T as #frame_system::Config>::DbWeight::get().reads(1)
+				}
+			}
+
 			fn on_runtime_upgrade() -> #frame_support::weights::Weight {
 				#frame_support::sp_tracing::enter_span!(
 					#frame_support::sp_tracing::trace_span!("on_runtime_update")
 				);
 
 				// log info about the upgrade.
-				let pallet_name = <
-					<T as #frame_system::Config>::PalletInfo
-					as
-					#frame_support::traits::PalletInfo
-				>::name::<Self>().unwrap_or("<unknown pallet name>");
 				#log_runtime_upgrade
 
 				<

frame/support/src/traits/hooks.rs

@@ -101,6 +101,36 @@ pub trait OnGenesis {
 
 /// See [`Hooks::on_runtime_upgrade`].
 pub trait OnRuntimeUpgrade {
+	/// Advanced lifecycle hook called for all pallets and custom migrations before any other
+	/// [`OnRuntimeUpgrade`] hooks.
+	///
+	/// Sometimes it is useful to have some logic execute prior to any other [`OnRuntimeUpgrade`]
+	/// hooks being called, which [`OnRuntimeUpgrade::before_all`] allows for.
+	///
+	/// When the Executive pallet executes runtime upgrades, hooks are called in this order:
+	/// 1. [`OnRuntimeUpgrade::before_all`] for all custom migrations
+	/// 2. [`OnRuntimeUpgrade::before_all`] for all pallet migrations
+	/// 3. For each custom migration:
+	/// 	1. [`OnRuntimeUpgrade::pre_upgrade`]
+	/// 	2. [`OnRuntimeUpgrade::on_runtime_upgrade`]
+	/// 	3. [`OnRuntimeUpgrade::post_upgrade`]
+	/// 4. For each pallet:
+	/// 	1. [`OnRuntimeUpgrade::pre_upgrade`]
+	/// 	2. [`OnRuntimeUpgrade::on_runtime_upgrade`]
+	/// 	3. [`OnRuntimeUpgrade::post_upgrade`]
+	///
+	/// An example of when this hook is useful is initializing the on-chain storage version of a
+	/// pallet when it is added to the runtime in a way which will not break
+	/// [`OnRuntimeUpgrade::post_upgrade`] checks.
+	///
+	/// This is an advanced hook unlikely to be needed in most [`OnRuntimeUpgrade`] implementations,
+	/// you can probably ignore it.
+	///
+	/// Returns the non-negotiable weight consumed by the checks.
+	fn before_all() -> Weight {
+		Weight::zero()
+	}
+
 	/// See [`Hooks::on_runtime_upgrade`].
 	fn on_runtime_upgrade() -> Weight {
 		Weight::zero()
@@ -145,6 +175,12 @@ pub trait OnRuntimeUpgrade {
 #[cfg_attr(all(feature = "tuples-96", not(feature = "tuples-128")), impl_for_tuples(96))]
 #[cfg_attr(feature = "tuples-128", impl_for_tuples(128))]
 impl OnRuntimeUpgrade for Tuple {
+	fn before_all() -> Weight {
+		let mut weight = Weight::zero();
+		for_tuples!( #( weight = weight.saturating_add(Tuple::before_all()); )* );
+		weight
+	}
+
 	fn on_runtime_upgrade() -> Weight {
 		let mut weight = Weight::zero();
 		for_tuples!( #( weight = weight.saturating_add(Tuple::on_runtime_upgrade()); )* );

frame/support/src/traits/metadata.rs

@@ -218,6 +218,23 @@ impl StorageVersion {
 
 		crate::storage::unhashed::get_or_default(&key)
 	}
+
+	/// Returns if the storage key for the given pallet exists in storage.
+	///
+	/// See [`STORAGE_VERSION_STORAGE_KEY_POSTFIX`] on how this key is built.
+	///
+	/// # Panics
+	///
+	/// This function will panic iff `Pallet` can not be found by `PalletInfo`.
+	/// In a runtime that is put together using
+	/// [`construct_runtime!`](crate::construct_runtime) this should never happen.
+	///
+	/// It will also panic if this function isn't executed in an externalities
+	/// provided environment.
+	pub fn exists<P: PalletInfoAccess>() -> bool {
+		let key = Self::storage_key::<P>();
+		crate::storage::unhashed::exists(&key)
+	}
 }
 
 impl PartialEq<u16> for StorageVersion {

frame/support/test/tests/pallet.rs

@@ -2110,6 +2110,54 @@ fn test_storage_alias() {
 	})
 }
 
+#[test]
+fn pallet_on_chain_storage_version_initializes_correctly() {
+	type Executive = frame_executive::Executive<
+		Runtime,
+		Block,
+		frame_system::ChainContext<Runtime>,
+		Runtime,
+		AllPalletsWithSystem,
+	>;
+
+	// Simple example of a pallet with current version 10 being added to the runtime for the first
+	// time.
+	TestExternalities::default().execute_with(|| {
+		let current_version = Example::current_storage_version();
+
+		// Check the on-chain storage version initially is not set and does not match the current
+		// version.
+		let on_chain_version_before = StorageVersion::get::<Example>();
+		assert_eq!(StorageVersion::exists::<Example>(), false);
+		assert_ne!(on_chain_version_before, current_version);
+
+		// OnRuntimeUpgrade calls pallet `before_all` hook which initializes the storage version.
+		Executive::execute_on_runtime_upgrade();
+
+		// Check that the storage version was initialized to the current version
+		let on_chain_version_after = StorageVersion::get::<Example>();
+		assert_eq!(on_chain_version_after, current_version);
+	});
+
+	// Pallet with no current storage version should have the on-chain version initialized to 0.
+	TestExternalities::default().execute_with(|| {
+		// Example4 current_storage_version is NoStorageVersionSet.
+
+		// Check the storage version is not set and implicitly 0.
+		let on_chain_version_before = StorageVersion::get::<Example4>();
+		assert_eq!(StorageVersion::exists::<Example4>(), false);
+		assert_eq!(on_chain_version_before, StorageVersion::new(0));
+
+		// OnRuntimeUpgrade calls pallet `before_all` hook which initializes the storage version.
+		Executive::execute_on_runtime_upgrade();
+
+		// Check that the storage version now exists and was initialized to 0.
+		let on_chain_version_after = StorageVersion::get::<Example4>();
+		assert_eq!(StorageVersion::exists::<Example4>(), true);
+		assert_eq!(on_chain_version_after, StorageVersion::new(0));
+	});
+}
+
 #[cfg(feature = "try-runtime")]
 #[test]
 fn post_runtime_upgrade_detects_storage_version_issues() {
@@ -2164,6 +2212,10 @@ fn post_runtime_upgrade_detects_storage_version_issues() {
 	TestExternalities::default().execute_with(|| {
 		// Call `on_genesis` to put the storage version of `Example` into the storage.
 		Example::on_genesis();
+
+		// Set the on-chain version to something different to the current version
+		StorageVersion::new(100).put::<Example>();
+
 		// The version isn't changed, we should detect it.
 		assert!(
 			Executive::try_runtime_upgrade(UpgradeCheckSelect::PreAndPost).unwrap_err() ==