support upgrade hooks to directly pass data (#12185)

* update interfaces of OnRuntimeUpgrade & Hooks

Signed-off-by: linning <linningde25@gmail.com>

* remove try-runtime for PreStateDigest

Signed-off-by: linning <linningde25@gmail.com>

* remove the Default bound of PreStateDigest

Signed-off-by: linning <linningde25@gmail.com>

* remove try-runtime for PreStateDigest & pre_upgrade

Signed-off-by: linning <linningde25@gmail.com>

* remove tmp storage between upgrade hooks

Signed-off-by: linning <linningde25@gmail.com>

* ensure hooks are storage noop

Signed-off-by: linning <linningde25@gmail.com>

* remove OnRuntimeUpgradeHelpersExt

Signed-off-by: linning <linningde25@gmail.com>

* cargo check & fmt

Signed-off-by: linning <linningde25@gmail.com>

* rename PreStateDigest to PreUpgradeState

Signed-off-by: linning <linningde25@gmail.com>

* replace associate type with codec & vec

Signed-off-by: linning <linningde25@gmail.com>

* add helper strcut to help encode/decode tuple

Signed-off-by: linning <linningde25@gmail.com>

* update comment

Signed-off-by: linning <linningde25@gmail.com>

* fix

Signed-off-by: linning <linningde25@gmail.com>

* add test

Signed-off-by: linning <linningde25@gmail.com>

* address comment

Signed-off-by: linning <linningde25@gmail.com>

* fix doc

Signed-off-by: linning <linningde25@gmail.com>

* fix ci

Signed-off-by: linning <linningde25@gmail.com>

* address comment

Signed-off-by: linning <linningde25@gmail.com>

* add more test cases

Signed-off-by: linning <linningde25@gmail.com>

* make clippy happy

Signed-off-by: linning <linningde25@gmail.com>

* fmt

Signed-off-by: linning <linningde25@gmail.com>

* update comment

Signed-off-by: linning <linningde25@gmail.com>

* fmt

Signed-off-by: linning <linningde25@gmail.com>

Signed-off-by: linning <linningde25@gmail.com>

substrate/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,17 +82,13 @@ 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.
@@ -98,8 +96,7 @@ impl<T: crate::Config<I>, I: 'static> OnRuntimeUpgrade for AddScore<T, I> {
 		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!(

substrate/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)
 	}
 }

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

substrate/frame/staking/src/migrations.rs

@@ -154,6 +154,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.
 	///
@@ -192,23 +196,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);
 

substrate/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)
 			}
 		}
 

substrate/frame/support/src/dispatch.rs

@@ -2158,12 +2158,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(())
 			}
 		}
@@ -2196,12 +2196,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(())
 			}
 		}

substrate/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},

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

substrate/frame/support/src/traits.rs

@@ -109,4 +109,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};

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

@@ -22,6 +22,9 @@ use impl_trait_for_tuples::impl_for_tuples;
 use sp_runtime::traits::AtLeast32BitUnsigned;
 use sp_std::prelude::*;
 
+#[cfg(feature = "try-runtime")]
+use codec::{Decode, Encode};
+
 /// The block initialization trait.
 ///
 /// Implementing this lets you express what should happen for your pallet when the block is
@@ -135,17 +138,25 @@ pub trait OnRuntimeUpgrade {
 
 	/// Execute some pre-checks prior to a runtime upgrade.
 	///
+	/// Return a `Vec<u8>` that can contain arbitrary encoded data (usually some pre-upgrade state),
+	/// which will be passed to `post_upgrade` after upgrading for post-check. An empty vector
+	/// should be returned if there is no such need.
+	///
 	/// This hook is never meant to be executed on-chain but is meant to be used by testing tools.
 	#[cfg(feature = "try-runtime")]
-	fn pre_upgrade() -> Result<(), &'static str> {
-		Ok(())
+	fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+		Ok(Vec::new())
 	}
 
 	/// Execute some post-checks after a runtime upgrade.
 	///
+	/// The `state` parameter is the `Vec<u8>` returned by `pre_upgrade` before upgrading, which
+	/// can be used for post-check. NOTE: if `pre_upgrade` is not implemented an empty vector will
+	/// be passed in, in such case `post_upgrade` should ignore it.
+	///
 	/// This hook is never meant to be executed on-chain but is meant to be used by testing tools.
 	#[cfg(feature = "try-runtime")]
-	fn post_upgrade() -> Result<(), &'static str> {
+	fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
 		Ok(())
 	}
 }
@@ -161,17 +172,21 @@ impl OnRuntimeUpgrade for Tuple {
 	}
 
 	#[cfg(feature = "try-runtime")]
-	fn pre_upgrade() -> Result<(), &'static str> {
-		let mut result = Ok(());
-		for_tuples!( #( result = result.and(Tuple::pre_upgrade()); )* );
-		result
+	fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+		let mut state: Vec<Vec<u8>> = Vec::default();
+		for_tuples!( #( state.push(Tuple::pre_upgrade()?); )* );
+		Ok(state.encode())
 	}
 
 	#[cfg(feature = "try-runtime")]
-	fn post_upgrade() -> Result<(), &'static str> {
-		let mut result = Ok(());
-		for_tuples!( #( result = result.and(Tuple::post_upgrade()); )* );
-		result
+	fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
+		let state: Vec<Vec<u8>> = Decode::decode(&mut state.as_slice())
+			.expect("the state parameter should be the same as pre_upgrade generated");
+		let mut state_iter = state.into_iter();
+		for_tuples!( #( Tuple::post_upgrade(
+			state_iter.next().expect("the state parameter should be the same as pre_upgrade generated")
+		)?; )* );
+		Ok(())
 	}
 }
 
@@ -243,17 +258,25 @@ pub trait Hooks<BlockNumber> {
 
 	/// Execute some pre-checks prior to a runtime upgrade.
 	///
+	/// Return a `Vec<u8>` that can contain arbitrary encoded data (usually some pre-upgrade state),
+	/// which will be passed to `post_upgrade` after upgrading for post-check. An empty vector
+	/// should be returned if there is no such need.
+	///
 	/// This hook is never meant to be executed on-chain but is meant to be used by testing tools.
 	#[cfg(feature = "try-runtime")]
-	fn pre_upgrade() -> Result<(), &'static str> {
-		Ok(())
+	fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+		Ok(Vec::new())
 	}
 
 	/// Execute some post-checks after a runtime upgrade.
 	///
+	/// The `state` parameter is the `Vec<u8>` returned by `pre_upgrade` before upgrading, which
+	/// can be used for post-check. NOTE: if `pre_upgrade` is not implemented an empty vector will
+	/// be passed in, in such case `post_upgrade` should ignore it.
+	///
 	/// This hook is never meant to be executed on-chain but is meant to be used by testing tools.
 	#[cfg(feature = "try-runtime")]
-	fn post_upgrade() -> Result<(), &'static str> {
+	fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
 		Ok(())
 	}
 
@@ -390,4 +413,94 @@ mod tests {
 			ON_IDLE_INVOCATION_ORDER.clear();
 		}
 	}
+
+	#[cfg(feature = "try-runtime")]
+	#[test]
+	fn on_runtime_upgrade_tuple() {
+		struct Test1;
+		struct Test2;
+		struct Test3;
+
+		impl OnRuntimeUpgrade for Test1 {
+			fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+				Ok("Test1".encode())
+			}
+			fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
+				let s: String = Decode::decode(&mut state.as_slice()).unwrap();
+				assert_eq!(s, "Test1");
+				Ok(())
+			}
+		}
+		impl OnRuntimeUpgrade for Test2 {
+			fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+				Ok(100u32.encode())
+			}
+			fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
+				let s: u32 = Decode::decode(&mut state.as_slice()).unwrap();
+				assert_eq!(s, 100);
+				Ok(())
+			}
+		}
+		impl OnRuntimeUpgrade for Test3 {
+			fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+				Ok(true.encode())
+			}
+			fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
+				let s: bool = Decode::decode(&mut state.as_slice()).unwrap();
+				assert_eq!(s, true);
+				Ok(())
+			}
+		}
+
+		type TestEmpty = ();
+		let origin_state = <TestEmpty as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+		let states: Vec<Vec<u8>> = Decode::decode(&mut origin_state.as_slice()).unwrap();
+		assert!(states.is_empty());
+		<TestEmpty as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+
+		type Test1Tuple = (Test1,);
+		let origin_state = <Test1Tuple as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+		let states: Vec<Vec<u8>> = Decode::decode(&mut origin_state.as_slice()).unwrap();
+		assert_eq!(states.len(), 1);
+		assert_eq!(
+			<String as Decode>::decode(&mut states[0].as_slice()).unwrap(),
+			"Test1".to_owned()
+		);
+		<Test1Tuple as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+
+		type Test123 = (Test1, Test2, Test3);
+		let origin_state = <Test123 as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+		let states: Vec<Vec<u8>> = Decode::decode(&mut origin_state.as_slice()).unwrap();
+		assert_eq!(
+			<String as Decode>::decode(&mut states[0].as_slice()).unwrap(),
+			"Test1".to_owned()
+		);
+		assert_eq!(<u32 as Decode>::decode(&mut states[1].as_slice()).unwrap(), 100u32);
+		assert_eq!(<bool as Decode>::decode(&mut states[2].as_slice()).unwrap(), true);
+		<Test123 as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+
+		type Test321 = (Test3, Test2, Test1);
+		let origin_state = <Test321 as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+		let states: Vec<Vec<u8>> = Decode::decode(&mut origin_state.as_slice()).unwrap();
+		assert_eq!(<bool as Decode>::decode(&mut states[0].as_slice()).unwrap(), true);
+		assert_eq!(<u32 as Decode>::decode(&mut states[1].as_slice()).unwrap(), 100u32);
+		assert_eq!(
+			<String as Decode>::decode(&mut states[2].as_slice()).unwrap(),
+			"Test1".to_owned()
+		);
+		<Test321 as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+
+		type TestNested123 = (Test1, (Test2, Test3));
+		let origin_state = <TestNested123 as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+		let states: Vec<Vec<u8>> = Decode::decode(&mut origin_state.as_slice()).unwrap();
+		assert_eq!(
+			<String as Decode>::decode(&mut states[0].as_slice()).unwrap(),
+			"Test1".to_owned()
+		);
+		// nested state for (Test2, Test3)
+		let nested_states: Vec<Vec<u8>> = Decode::decode(&mut states[1].as_slice()).unwrap();
+		assert_eq!(<u32 as Decode>::decode(&mut nested_states[0].as_slice()).unwrap(), 100u32);
+		assert_eq!(<bool as Decode>::decode(&mut nested_states[1].as_slice()).unwrap(), true);
+		<TestNested123 as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+	}
 }

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

@@ -17,46 +17,10 @@
 
 //! Try-runtime specific traits and types.
 
-use super::*;
 use impl_trait_for_tuples::impl_for_tuples;
 use sp_arithmetic::traits::AtLeast32BitUnsigned;
 use sp_std::prelude::*;
 
-/// Prefix to be used (optionally) for implementing [`OnRuntimeUpgradeHelpersExt::storage_key`].
-const ON_RUNTIME_UPGRADE_PREFIX: &[u8] = b"__ON_RUNTIME_UPGRADE__";
-
-/// Some helper functions for [`OnRuntimeUpgrade`] during `try-runtime` testing.
-pub trait OnRuntimeUpgradeHelpersExt {
-	/// Generate a storage key unique to this runtime upgrade.
-	///
-	/// This can be used to communicate data from pre-upgrade to post-upgrade state and check
-	/// them. See [`Self::set_temp_storage`] and [`Self::get_temp_storage`].
-	fn storage_key(ident: &str) -> [u8; 32] {
-		crate::storage::storage_prefix(ON_RUNTIME_UPGRADE_PREFIX, ident.as_bytes())
-	}
-
-	/// Get temporary storage data written by [`Self::set_temp_storage`].
-	///
-	/// Returns `None` if either the data is unavailable or un-decodable.
-	///
-	/// A `at` storage identifier must be provided to indicate where the storage is being read from.
-	fn get_temp_storage<T: codec::Decode>(at: &str) -> Option<T> {
-		sp_io::storage::get(&Self::storage_key(at))
-			.and_then(|bytes| codec::Decode::decode(&mut &*bytes).ok())
-	}
-
-	/// Write some temporary data to a specific storage that can be read (potentially in
-	/// post-upgrade hook) via [`Self::get_temp_storage`].
-	///
-	/// A `at` storage identifier must be provided to indicate where the storage is being written
-	/// to.
-	fn set_temp_storage<T: codec::Encode>(data: T, at: &str) {
-		sp_io::storage::set(&Self::storage_key(at), &data.encode());
-	}
-}
-
-impl<U: OnRuntimeUpgrade> OnRuntimeUpgradeHelpersExt for U {}
-
 // Which state tests to execute.
 #[derive(codec::Encode, codec::Decode, Clone)]
 pub enum Select {
@@ -68,7 +32,7 @@ pub enum Select {
 	RoundRobin(u32),
 	/// Run only pallets who's name matches the given list.
 	///
-	/// Pallet names are obtained from [`PalletInfoAccess`].
+	/// Pallet names are obtained from [`super::PalletInfoAccess`].
 	Only(Vec<Vec<u8>>),
 }
 

substrate/utils/frame/try-runtime/cli/src/lib.rs

@@ -132,20 +132,20 @@
 //! added, given the right flag:
 //!
 //! ```ignore
+//! 
 //! #[cfg(feature = try-runtime)]
-//! fn pre_upgrade() -> Result<(), &'static str> {}
+//! fn pre_upgrade() -> Result<Vec<u8>, &'static str> {}
 //!
 //! #[cfg(feature = try-runtime)]
-//! fn post_upgrade() -> Result<(), &'static str> {}
+//! fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {}
 //! ```
 //!
 //! (The pallet macro syntax will support this simply as a part of `#[pallet::hooks]`).
 //!
 //! These hooks allow you to execute some code, only within the `on-runtime-upgrade` command, before
-//! and after the migration. If any data needs to be temporarily stored between the pre/post
-//! migration hooks, `OnRuntimeUpgradeHelpersExt` can help with that. Note that you should be
-//! mindful with any mutable storage ops in the pre/post migration checks, as you almost certainly
-//! will not want to mutate any of the storage that is to be migrated.
+//! and after the migration. Moreover, `pre_upgrade` can return a `Vec<u8>` that contains arbitrary
+//! encoded data (usually some pre-upgrade state) which will be passed to `post_upgrade` after
+//! upgrading and used for post checking.
 //!
 //! #### Logging
 //!