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