Initialise on-chain StorageVersion for pallets added after genesis

substrate/frame/executive/src/lib.rs

@@ -121,8 +121,8 @@ use frame_support::{
 	dispatch::{DispatchClass, DispatchInfo, GetDispatchInfo, PostDispatchInfo},
 	pallet_prelude::InvalidTransaction,
 	traits::{
-		EnsureInherentsAreFirst, ExecuteBlock, OffchainWorker, OnFinalize, OnIdle, OnInitialize,
-		OnRuntimeUpgrade,
+		BeforeAllRuntimeMigrations, EnsureInherentsAreFirst, ExecuteBlock, OffchainWorker,
+		OnFinalize, OnIdle, OnInitialize, OnRuntimeUpgrade,
 	},
 	weights::Weight,
 };
@@ -188,6 +188,7 @@ impl<
 		Context: Default,
 		UnsignedValidator,
 		AllPalletsWithSystem: OnRuntimeUpgrade
+			+ BeforeAllRuntimeMigrations
 			+ OnInitialize<BlockNumberFor<System>>
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
@@ -225,6 +226,7 @@ impl<
 		Context: Default,
 		UnsignedValidator,
 		AllPalletsWithSystem: OnRuntimeUpgrade
+			+ BeforeAllRuntimeMigrations
 			+ OnInitialize<BlockNumberFor<System>>
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
@@ -366,7 +368,9 @@ where
 			)?;
 		}
 
-		let weight =
+		let before_all_weight =
+			<AllPalletsWithSystem as BeforeAllRuntimeMigrations>::before_all_runtime_migrations();
+		let try_on_runtime_upgrade_weight =
 			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::try_on_runtime_upgrade(
 				checks.pre_and_post(),
 			)?;
@@ -381,7 +385,7 @@ where
 			)?;
 		}
 
-		Ok(weight)
+		Ok(before_all_weight.saturating_add(try_on_runtime_upgrade_weight))
 	}
 }
 
@@ -394,6 +398,7 @@ impl<
 		Context: Default,
 		UnsignedValidator,
 		AllPalletsWithSystem: OnRuntimeUpgrade
+			+ BeforeAllRuntimeMigrations
 			+ OnInitialize<BlockNumberFor<System>>
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
@@ -410,7 +415,10 @@ where
 {
 	/// Execute all `OnRuntimeUpgrade` of this runtime, and return the aggregate weight.
 	pub fn execute_on_runtime_upgrade() -> Weight {
+		let before_all_weight =
+			<AllPalletsWithSystem as BeforeAllRuntimeMigrations>::before_all_runtime_migrations();
 		<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::on_runtime_upgrade()
+			.saturating_add(before_all_weight)
 	}
 
 	/// Start the execution of a particular block.

substrate/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::__private::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::__private::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::__private::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::__private::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::__private::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::__private::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,
 				);
 
@@ -186,6 +201,27 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			}
 		}
 
+		impl<#type_impl_gen>
+			#frame_support::traits::BeforeAllRuntimeMigrations
+			for #pallet_ident<#type_use_gen> #where_clause
+		{
+			fn before_all_runtime_migrations() -> #frame_support::weights::Weight {
+				use #frame_support::traits::Get;
+				#frame_support::__private::sp_tracing::enter_span!(
+					#frame_support::__private::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)
+				}
+			}
+		}
+
 		impl<#type_impl_gen>
 			#frame_support::traits::OnRuntimeUpgrade
 			for #pallet_ident<#type_use_gen> #where_clause
@@ -196,11 +232,6 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 				);
 
 				// 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
 
 				<

substrate/frame/support/src/traits.rs

@@ -83,8 +83,8 @@ mod hooks;
 #[allow(deprecated)]
 pub use hooks::GenesisBuild;
 pub use hooks::{
-	BuildGenesisConfig, Hooks, IntegrityTest, OnFinalize, OnGenesis, OnIdle, OnInitialize,
-	OnRuntimeUpgrade, OnTimestampSet,
+	BeforeAllRuntimeMigrations, BuildGenesisConfig, Hooks, IntegrityTest, OnFinalize, OnGenesis,
+	OnIdle, OnInitialize, OnRuntimeUpgrade, OnTimestampSet,
 };
 
 pub mod schedule;

substrate/frame/support/src/traits/hooks.rs

@@ -99,6 +99,22 @@ pub trait OnGenesis {
 	fn on_genesis() {}
 }
 
+/// Implemented by pallets, allows defining logic to run prior to any [`OnRuntimeUpgrade`] logic.
+///
+/// This was introduced to allow pallets to initialize their on-chain storage version in the case
+/// that they are added to the runtime after genesis.
+///
+/// This hook is intended to be used internally in FRAME and not be exposed to FRAME developers.
+///
+/// It is defined as a seperate trait from [`OnRuntimeUpgrade`] precisely to not pollute the public
+/// API.
+pub trait BeforeAllRuntimeMigrations {
+	/// Something that should happen before runtime migrations are executed.
+	fn before_all_runtime_migrations() -> Weight {
+		Weight::zero()
+	}
+}
+
 /// See [`Hooks::on_runtime_upgrade`].
 pub trait OnRuntimeUpgrade {
 	/// See [`Hooks::on_runtime_upgrade`].
@@ -150,10 +166,24 @@ pub trait OnRuntimeUpgrade {
 	}
 }
 
+#[cfg_attr(all(not(feature = "tuples-96"), not(feature = "tuples-128")), impl_for_tuples(64))]
+#[cfg_attr(all(feature = "tuples-96", not(feature = "tuples-128")), impl_for_tuples(96))]
+#[cfg_attr(feature = "tuples-128", impl_for_tuples(128))]
+impl BeforeAllRuntimeMigrations for Tuple {
+	/// Implements the default behavior of
+	/// [`BeforeAllRuntimeMigrations::before_all_runtime_migrations`] for tuples.
+	fn before_all_runtime_migrations() -> Weight {
+		let mut weight = Weight::zero();
+		for_tuples!( #( weight = weight.saturating_add(Tuple::before_all_runtime_migrations()); )* );
+		weight
+	}
+}
+
 #[cfg_attr(all(not(feature = "tuples-96"), not(feature = "tuples-128")), impl_for_tuples(64))]
 #[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 {
+	/// Implements the default behavior of [`OnRuntimeUpgrade::on_runtime_upgrade`] for tuples.
 	fn on_runtime_upgrade() -> Weight {
 		let mut weight = Weight::zero();
 		for_tuples!( #( weight = weight.saturating_add(Tuple::on_runtime_upgrade()); )* );

substrate/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 {

substrate/frame/support/test/tests/pallet.rs

@@ -2329,6 +2329,55 @@ 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);
+
+		// [`frame_support::traits::BeforeAllRuntimeMigrations`] hook should initialize 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));
+
+		// [`frame_support::traits::BeforeAllRuntimeMigrations`] 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() {
@@ -2381,8 +2430,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 one less than the current version for `Example`, simulating a
+		// forgotten migration
+		StorageVersion::new(9).put::<Example2>();
+
 		// The version isn't changed, we should detect it.
 		assert!(
 			Executive::try_runtime_upgrade(UpgradeCheckSelect::PreAndPost).unwrap_err() ==