Initialise on-chain StorageVersion for pallets added after genesis (#1297)

Original PR https://github.com/paritytech/substrate/pull/14641

---

Closes https://github.com/paritytech/polkadot-sdk/issues/109

### Problem
Quoting from the above issue:

> When adding a pallet to chain after genesis we currently don't set the
StorageVersion. So, when calling on_chain_storage_version it returns 0
while the pallet is maybe already at storage version 9 when it was added
to the chain. This could lead to issues when running migrations.

### Solution

- Create a new trait `BeforeAllRuntimeMigrations` with a single method
`fn before_all_runtime_migrations() -> Weight` trait with a noop default
implementation
- Modify `Executive` to call
`BeforeAllRuntimeMigrations::before_all_runtime_migrations` for all
pallets before running any other hooks
- Implement `BeforeAllRuntimeMigrations` in the pallet proc macro to
initialize the on-chain version to the current pallet version if the
pallet has no storage set (indicating it has been recently added to the
runtime and needs to have its version initialised).

### Other changes in this PR

- Abstracted repeated boilerplate to access the `pallet_name` in the
pallet expand proc macro.

### FAQ

#### Why create a new hook instead of adding this logic to the pallet
`pre_upgrade`?

`Executive` currently runs `COnRuntimeUpgrade` (custom migrations)
before `AllPalletsWithSystem` migrations. We need versions to be
initialized before the `COnRuntimeUpgrade` migrations are run, because
`COnRuntimeUpgrade` migrations may use the on-chain version for critical
logic. e.g. `VersionedRuntimeUpgrade` uses it to decide whether or not
to execute.

We cannot reorder `COnRuntimeUpgrade` and `AllPalletsWithSystem` so
`AllPalletsWithSystem` runs first, because `AllPalletsWithSystem` have
some logic in their `post_upgrade` hooks to verify that the on-chain
version and current pallet version match. A common use case of
`COnRuntimeUpgrade` migrations is to perform a migration which will
result in the versions matching, so if they were reordered these
`post_upgrade` checks would fail.

#### Why init the on-chain version for pallets without a current storage
version?

We must init the on-chain version for pallets even if they don't have a
defined storage version so if there is a future version bump, the
on-chain version is not automatically set to that new version without a
proper migration.

e.g. bad scenario:

1. A pallet with no 'current version' is added to the runtime
2. Later, the pallet is upgraded with the 'current version' getting set
to 1 and a migration is added to Executive Migrations to migrate the
storage from 0 to 1
    a. Runtime upgrade occurs
    b. `before_all` hook initializes the on-chain version to 1
c. `on_runtime_upgrade` of the migration executes, and sees the on-chain
version is already 1 therefore think storage is already migrated and
does not execute the storage migration
Now, on-chain version is 1 but storage is still at version 0.

By always initializing the on-chain version when the pallet is added to
the runtime we avoid that scenario.

---------

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
Co-authored-by: Bastian Köcher <git@kchr.de>

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,
+				"🐥 New pallet {:?} detected in the runtime. 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,
+				"🐥 New pallet {:?} detected in the runtime. The pallet has no 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,7 +85,7 @@ 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,
 			);
 		}
 	};
@@ -78,16 +110,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,
 				);
@@ -100,17 +126,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,
 				);
 
@@ -173,6 +193,32 @@ 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, PalletInfoAccess};
+				use #frame_support::__private::hashing::twox_128;
+				use #frame_support::storage::unhashed::contains_prefixed_key;
+				#frame_support::__private::sp_tracing::enter_span!(
+					#frame_support::__private::sp_tracing::trace_span!("before_all")
+				);
+
+				// Check if the pallet has any keys set, including the storage version. If there are
+				// no keys set, the pallet was just added to the runtime and needs to have its
+				// version initialized.
+				let pallet_hashed_prefix = <Self as PalletInfoAccess>::name_hash();
+				let exists = contains_prefixed_key(&pallet_hashed_prefix);
+				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
@@ -183,11 +229,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
 
 				<
@@ -258,16 +299,10 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 				n: #frame_system::pallet_prelude::BlockNumberFor::<T>,
 				_s: #frame_support::traits::TryStateSelect
 			) -> Result<(), #frame_support::sp_runtime::TryRuntimeError> {
-				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::info!(
 					target: #frame_support::LOG_TARGET,
 					"🩺 Running {:?} try-state checks",
-					pallet_name,
+					#pallet_name,
 				);
 				<
 					Self as #frame_support::traits::Hooks<
@@ -277,7 +312,7 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 					#frame_support::__private::log::error!(
 						target: #frame_support::LOG_TARGET,
 						"❌ {:?} try_state checks failed: {:?}",
-						pallet_name,
+						#pallet_name,
 						err
 					);
 

substrate/frame/support/src/lib.rs

@@ -47,7 +47,7 @@ pub mod __private {
 	pub use sp_core::{OpaqueMetadata, Void};
 	pub use sp_core_hashing_proc_macro;
 	pub use sp_inherents;
-	pub use sp_io::{self, storage::root as storage_root};
+	pub use sp_io::{self, hashing, storage::root as storage_root};
 	pub use sp_metadata_ir as metadata_ir;
 	#[cfg(feature = "std")]
 	pub use sp_runtime::{bounded_btree_map, bounded_vec};

substrate/frame/support/src/migrations.rs

@@ -119,7 +119,7 @@ impl<
 		let on_chain_version = Pallet::on_chain_storage_version();
 		if on_chain_version == FROM {
 			log::info!(
-				"🚚 Pallet {:?} migrating storage version from {:?} to {:?}.",
+				"🚚 Pallet {:?} VersionedMigration migrating storage version from {:?} to {:?}.",
 				Pallet::name(),
 				FROM,
 				TO
@@ -134,7 +134,7 @@ impl<
 			weight.saturating_add(DbWeight::get().reads_writes(1, 1))
 		} else {
 			log::warn!(
-				"🚚 Pallet {:?} migration {}->{} can be removed; on-chain is already at {:?}.",
+				"🚚 Pallet {:?} VersionedMigration migration {}->{} can be removed; on-chain is already at {:?}.",
 				Pallet::name(),
 				FROM,
 				TO,

substrate/frame/support/src/traits.rs

@@ -84,8 +84,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,19 @@ pub trait OnGenesis {
 	fn on_genesis() {}
 }
 
+/// Implemented by pallets, allows defining logic to run prior to any [`OnRuntimeUpgrade`] logic.
+///
+/// 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 +163,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

@@ -222,6 +222,23 @@ impl StorageVersion {
 
 		crate::storage::unhashed::get_or_default(&key)
 	}
+
+	/// Returns if the storage version 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

@@ -21,7 +21,7 @@ use frame_support::{
 	dispatch_context::with_context,
 	pallet_prelude::{StorageInfoTrait, ValueQuery},
 	parameter_types,
-	storage::unhashed,
+	storage::{unhashed, unhashed::contains_prefixed_key},
 	traits::{
 		ConstU32, GetCallIndex, GetCallName, GetStorageVersion, OnFinalize, OnGenesis,
 		OnInitialize, OnRuntimeUpgrade, PalletError, PalletInfoAccess, StorageVersion,
@@ -2330,6 +2330,58 @@ 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 pallet has no storage items set.
+		let pallet_hashed_prefix = twox_128(Example::name().as_bytes());
+		let exists = contains_prefixed_key(&pallet_hashed_prefix);
+		assert_eq!(exists, false);
+
+		// [`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 pallet has no storage items set.
+		let pallet_hashed_prefix = twox_128(Example4::name().as_bytes());
+		let exists = contains_prefixed_key(&pallet_hashed_prefix);
+		assert_eq!(exists, false);
+
+		// Confirm the storage version is implicitly 0.
+		let on_chain_version_before = StorageVersion::get::<Example4>();
+		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() {
@@ -2382,8 +2434,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() ==