support upgrade hooks to directly pass data

frame/bags-list/src/migrations.rs

@@ -24,6 +24,8 @@ use frame_support::traits::OnRuntimeUpgrade;
 
 #[cfg(feature = "try-runtime")]
 use frame_support::ensure;
+#[cfg(feature = "try-runtime")]
+use sp_std::vec::Vec;
 
 /// A struct that does not migration, but only checks that the counter prefix exists and is correct.
 pub struct CheckCounterPrefix<T: crate::Config<I>, I: 'static>(sp_std::marker::PhantomData<(T, I)>);
@@ -33,7 +35,7 @@ impl<T: crate::Config<I>, I: 'static> OnRuntimeUpgrade for CheckCounterPrefix<T,
 	}
 
 	#[cfg(feature = "try-runtime")]
-	fn pre_upgrade() -> Result<(), &'static str> {
+	fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
 		// The old explicit storage item.
 		#[frame_support::storage_alias]
 		type CounterForListNodes<T: crate::Config<I>, I: 'static> =
@@ -51,7 +53,7 @@ impl<T: crate::Config<I>, I: 'static> OnRuntimeUpgrade for CheckCounterPrefix<T,
 			crate::ListNodes::<T, I>::count()
 		);
 
-		Ok(())
+		Ok(Vec::new())
 	}
 }
 
@@ -80,26 +82,21 @@ mod old {
 	#[frame_support::storage_alias]
 	pub type CounterForListNodes<T: crate::Config<I>, I: 'static> =
 		StorageValue<crate::Pallet<T, I>, u32, ValueQuery>;
-
-	#[frame_support::storage_alias]
-	pub type TempStorage<T: crate::Config<I>, I: 'static> =
-		StorageValue<crate::Pallet<T, I>, u32, ValueQuery>;
 }
 
 /// A struct that migrates all bags lists to contain a score value.
 pub struct AddScore<T: crate::Config<I>, I: 'static = ()>(sp_std::marker::PhantomData<(T, I)>);
 impl<T: crate::Config<I>, I: 'static> OnRuntimeUpgrade for AddScore<T, I> {
 	#[cfg(feature = "try-runtime")]
-	fn pre_upgrade() -> Result<(), &'static str> {
+	fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
 		// The list node data should be corrupt at this point, so this is zero.
 		ensure!(crate::ListNodes::<T, I>::iter().count() == 0, "list node data is not corrupt");
 		// We can use the helper `old::ListNode` to get the existing data.
 		let iter_node_count: u32 = old::ListNodes::<T, I>::iter().count() as u32;
 		let tracked_node_count: u32 = old::CounterForListNodes::<T, I>::get();
 		crate::log!(info, "number of nodes before: {:?} {:?}", iter_node_count, tracked_node_count);
 		ensure!(iter_node_count == tracked_node_count, "Node count is wrong.");
-		old::TempStorage::<T, I>::put(iter_node_count);
-		Ok(())
+		Ok(iter_node_count.encode())
 	}
 
 	fn on_runtime_upgrade() -> frame_support::weights::Weight {
@@ -122,9 +119,10 @@ impl<T: crate::Config<I>, I: 'static> OnRuntimeUpgrade for AddScore<T, I> {
 	}
 
 	#[cfg(feature = "try-runtime")]
-	fn post_upgrade() -> Result<(), &'static str> {
-		let node_count_before = old::TempStorage::<T, I>::take();
-		// Now, the list node data is not corrupt anymore.
+	fn post_upgrade(node_count_before: Vec<u8>) -> Result<(), &'static str> {
+		let node_count_before: u32 = Decode::decode(&mut node_count_before.as_slice())
+			.expect("the state parameter should be something that was generated by pre_upgrade");
+		// Now the list node data is not corrupt anymore.
 		let iter_node_count_after: u32 = crate::ListNodes::<T, I>::iter().count() as u32;
 		let tracked_node_count_after: u32 = crate::ListNodes::<T, I>::count();
 		crate::log!(

frame/executive/src/lib.rs

@@ -287,10 +287,15 @@ where
 	///
 	/// This should only be used for testing.
 	pub fn try_runtime_upgrade() -> Result<frame_support::weights::Weight, &'static str> {
-		<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+		// ensure both `pre_upgrade` and `post_upgrade` won't change the storage root
+		let state = {
+			let _guard = frame_support::StorageNoopGuard::default();
+			<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::pre_upgrade().unwrap()
+		};
 		let weight = Self::execute_on_runtime_upgrade();
-		<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::post_upgrade().unwrap();
-
+		let _guard = frame_support::StorageNoopGuard::default();
+		<(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::post_upgrade(state)
+			.unwrap();
 		Ok(weight)
 	}
 }

frame/nomination-pools/src/migration.rs

@@ -18,7 +18,7 @@
 use super::*;
 use crate::log;
 use frame_support::traits::OnRuntimeUpgrade;
-use sp_std::collections::btree_map::BTreeMap;
+use sp_std::{collections::btree_map::BTreeMap, vec::Vec};
 
 pub mod v1 {
 	use super::*;
@@ -97,7 +97,7 @@ pub mod v1 {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn post_upgrade() -> Result<(), &'static str> {
+		fn post_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
 			// new version must be set.
 			assert_eq!(Pallet::<T>::on_chain_storage_version(), 1);
 			Pallet::<T>::try_state(frame_system::Pallet::<T>::block_number())?;
@@ -347,7 +347,7 @@ pub mod v2 {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn pre_upgrade() -> Result<(), &'static str> {
+		fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
 			// all reward accounts must have more than ED.
 			RewardPools::<T>::iter().for_each(|(id, _)| {
 				assert!(
@@ -356,11 +356,11 @@ pub mod v2 {
 				)
 			});
 
-			Ok(())
+			Ok(Vec::new())
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn post_upgrade() -> Result<(), &'static str> {
+		fn post_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
 			// new version must be set.
 			assert_eq!(Pallet::<T>::on_chain_storage_version(), 2);
 
@@ -430,16 +430,16 @@ pub mod v3 {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn pre_upgrade() -> Result<(), &'static str> {
+		fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
 			ensure!(
 				Pallet::<T>::current_storage_version() > Pallet::<T>::on_chain_storage_version(),
 				"the on_chain version is equal or more than the current one"
 			);
-			Ok(())
+			Ok(Vec::new())
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn post_upgrade() -> Result<(), &'static str> {
+		fn post_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
 			ensure!(
 				Metadata::<T>::iter_keys().all(|id| BondedPools::<T>::contains_key(&id)),
 				"not all of the stale metadata has been removed"

frame/staking/src/migrations.rs

@@ -61,6 +61,10 @@ pub mod v10 {
 
 pub mod v9 {
 	use super::*;
+	#[cfg(feature = "try-runtime")]
+	use frame_support::codec::{Decode, Encode};
+	#[cfg(feature = "try-runtime")]
+	use sp_std::vec::Vec;
 
 	/// Migration implementation that injects all validators into sorted list.
 	///
@@ -99,23 +103,22 @@ pub mod v9 {
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn pre_upgrade() -> Result<(), &'static str> {
-			use frame_support::traits::OnRuntimeUpgradeHelpersExt;
+		fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
 			frame_support::ensure!(
 				StorageVersion::<T>::get() == crate::Releases::V8_0_0,
 				"must upgrade linearly"
 			);
 
 			let prev_count = T::VoterList::count();
-			Self::set_temp_storage(prev_count, "prev");
-			Ok(())
+			Ok(prev_count.encode())
 		}
 
 		#[cfg(feature = "try-runtime")]
-		fn post_upgrade() -> Result<(), &'static str> {
-			use frame_support::traits::OnRuntimeUpgradeHelpersExt;
+		fn post_upgrade(prev_count: Vec<u8>) -> Result<(), &'static str> {
+			let prev_count: u32 = Decode::decode(&mut prev_count.as_slice()).expect(
+				"the state parameter should be something that was generated by pre_upgrade",
+			);
 			let post_count = T::VoterList::count();
-			let prev_count = Self::get_temp_storage::<u32>("prev").unwrap();
 			let validators = Validators::<T>::count();
 			assert!(post_count == prev_count + validators);
 

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

@@ -160,7 +160,7 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			}
 
 			#[cfg(feature = "try-runtime")]
-			fn pre_upgrade() -> Result<(), &'static str> {
+			fn pre_upgrade() -> Result<#frame_support::sp_std::vec::Vec<u8>, &'static str> {
 				<
 					Self
 					as
@@ -169,12 +169,12 @@ pub fn expand_hooks(def: &mut Def) -> proc_macro2::TokenStream {
 			}
 
 			#[cfg(feature = "try-runtime")]
-			fn post_upgrade() -> Result<(), &'static str> {
+			fn post_upgrade(state: #frame_support::sp_std::vec::Vec<u8>) -> Result<(), &'static str> {
 				<
 					Self
 					as
 					#frame_support::traits::Hooks<<T as #frame_system::Config>::BlockNumber>
-				>::post_upgrade()
+				>::post_upgrade(state)
 			}
 		}
 

frame/support/src/dispatch.rs

@@ -1609,12 +1609,12 @@ macro_rules! decl_module {
 			}
 
 			#[cfg(feature = "try-runtime")]
-			fn pre_upgrade() -> Result<(), &'static str> {
-				Ok(())
+			fn pre_upgrade() -> Result<$crate::sp_std::vec::Vec<u8>, &'static str> {
+				Ok($crate::sp_std::vec::Vec::new())
 			}
 
 			#[cfg(feature = "try-runtime")]
-			fn post_upgrade() -> Result<(), &'static str> {
+			fn post_upgrade(_: $crate::sp_std::vec::Vec<u8>) -> Result<(), &'static str> {
 				Ok(())
 			}
 		}
@@ -1647,12 +1647,12 @@ macro_rules! decl_module {
 			}
 
 			#[cfg(feature = "try-runtime")]
-			fn pre_upgrade() -> Result<(), &'static str> {
-				Ok(())
+			fn pre_upgrade() -> Result<$crate::sp_std::vec::Vec<u8>, &'static str> {
+				Ok($crate::sp_std::vec::Vec::new())
 			}
 
 			#[cfg(feature = "try-runtime")]
-			fn post_upgrade() -> Result<(), &'static str> {
+			fn post_upgrade(_: $crate::sp_std::vec::Vec<u8>) -> Result<(), &'static str> {
 				Ok(())
 			}
 		}

frame/support/src/lib.rs

@@ -93,7 +93,7 @@ pub mod unsigned {
 	};
 }
 
-#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
+#[cfg(any(feature = "std", feature = "runtime-benchmarks", feature = "try-runtime", test))]
 pub use self::storage::storage_noop_guard::StorageNoopGuard;
 pub use self::{
 	dispatch::{Callable, Parameter},

frame/support/src/storage/storage_noop_guard.rs

@@ -16,7 +16,7 @@
 // limitations under the License.
 
 // Feature gated since it can panic.
-#![cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
+#![cfg(any(feature = "std", feature = "runtime-benchmarks", feature = "try-runtime", test))]
 
 //! Contains the [`crate::StorageNoopGuard`] for conveniently asserting
 //! that no storage mutation has been made by a whole code block.

frame/support/src/traits.rs

@@ -108,4 +108,4 @@ pub use voting::{
 #[cfg(feature = "try-runtime")]
 mod try_runtime;
 #[cfg(feature = "try-runtime")]
-pub use try_runtime::{OnRuntimeUpgradeHelpersExt, Select as TryStateSelect, TryState};
+pub use try_runtime::{Select as TryStateSelect, TryState};