[Feature] Sequential migration execution for try-runtime (paritytech#…

…12319) * [Feature] Sequential migration execution for try-runtime * remove unused * guards * reinstate encode/decode * proper feature gate * proper test feature gate * Update frame/support/src/traits/hooks.rs Co-authored-by: Kian Paimani <[email protected]> * Update frame/support/src/traits/hooks.rs Co-authored-by: Kian Paimani <[email protected]> * fix tests * redo Tuple tests * Update frame/support/src/traits/hooks.rs Co-authored-by: Kian Paimani <[email protected]> * use parameter_types for testing * lint fix * Update frame/support/src/traits/hooks.rs Co-authored-by: Jegor Sidorenko <[email protected]> * Update frame/support/src/traits/hooks.rs Co-authored-by: Oliver Tale-Yazdi <[email protected]> * eloquent feature gate * redo tests * more fixes * properly handle pre/post errors * remove some tests and fix the others * add format import * import fix * more import fixes Co-authored-by: Kian Paimani <[email protected]> Co-authored-by: Jegor Sidorenko <[email protected]> Co-authored-by: Oliver Tale-Yazdi <[email protected]> Co-authored-by: parity-processbot <>
gear-tech · Feb 27, 2023 · 092a700 · 092a700
1 parent 6e13922
commit 092a700
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 71 deletions.
diff --git a/frame/executive/src/lib.rs b/frame/executive/src/lib.rs
@@ -287,15 +287,7 @@ where
  ///
  /// This should only be used for testing.
  pub fn try_runtime_upgrade() -> Result<frame_support::weights::Weight, &'static str> {
- // 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();
- let _guard = frame_support::StorageNoopGuard::default();
- <(COnRuntimeUpgrade, AllPalletsWithSystem) as OnRuntimeUpgrade>::post_upgrade(state)
- .unwrap();
  Ok(weight)
  }
 }

diff --git a/frame/support/src/traits/hooks.rs b/frame/support/src/traits/hooks.rs
@@ -22,7 +22,7 @@ use impl_trait_for_tuples::impl_for_tuples;
 use sp_runtime::traits::AtLeast32BitUnsigned;
 use sp_std::prelude::*;
 
-#[cfg(feature = "try-runtime")]
+#[cfg(all(feature = "try-runtime", test))]
 use codec::{Decode, Encode};
 
 /// The block initialization trait.
@@ -165,27 +165,50 @@ pub trait OnRuntimeUpgrade {
 #[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 {
+ #[cfg(not(feature = "try-runtime"))]
  fn on_runtime_upgrade() -> Weight {
  let mut weight = Weight::zero();
  for_tuples!( #( weight = weight.saturating_add(Tuple::on_runtime_upgrade()); )* );
  weight
  }
 
  #[cfg(feature = "try-runtime")]
+ /// We are executing pre- and post-checks sequentially in order to be able to test several
+ /// consecutive migrations for the same pallet without errors. Therefore pre and post upgrade
+ /// hooks for tuples are a noop.
+ fn on_runtime_upgrade() -> Weight {
+ use scale_info::prelude::format;
+
+ let mut weight = Weight::zero();
+ // migration index in the tuple, start with 1 for better readability
+ let mut i = 1;
+ for_tuples!( #(
+ let _guard = frame_support::StorageNoopGuard::default();
+ // we want to panic if any checks fail right here right now.
+ let state = Tuple::pre_upgrade().expect(&format!("PreUpgrade failed for migration #{}", i));
+ drop(_guard);
+
+ weight = weight.saturating_add(Tuple::on_runtime_upgrade());
+
+ let _guard = frame_support::StorageNoopGuard::default();
+ // we want to panic if any checks fail right here right now.
+ Tuple::post_upgrade(state).expect(&format!("PostUpgrade failed for migration #{}", i));
+ drop(_guard);
+
+ i += 1;
+ )* );
+ weight
+ }
+
+ #[cfg(feature = "try-runtime")]
+ /// noop
  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())
+ Ok(Vec::new())
  }
 
  #[cfg(feature = "try-runtime")]
- 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")
- )?; )* );
+ /// noop
+ fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
  Ok(())
  }
 }
@@ -342,7 +365,9 @@ mod tests {
 
  #[test]
  fn on_initialize_and_on_runtime_upgrade_weight_merge_works() {
+ use sp_io::TestExternalities;
  struct Test;
+
  impl OnInitialize<u8> for Test {
  fn on_initialize(_n: u8) -> Weight {
  Weight::from_ref_time(10)
@@ -354,8 +379,10 @@ mod tests {
  }
  }
 
- assert_eq!(<(Test, Test)>::on_initialize(0), Weight::from_ref_time(20));
- assert_eq!(<(Test, Test)>::on_runtime_upgrade(), Weight::from_ref_time(40));
+ TestExternalities::default().execute_with(|| {
+ assert_eq!(<(Test, Test)>::on_initialize(0), Weight::from_ref_time(20));
+ assert_eq!(<(Test, Test)>::on_runtime_upgrade(), Weight::from_ref_time(40));
+ });
  }
 
  #[test]
@@ -416,91 +443,107 @@ mod tests {
 
  #[cfg(feature = "try-runtime")]
  #[test]
+ #[allow(dead_code)]
  fn on_runtime_upgrade_tuple() {
+ use frame_support::parameter_types;
+ use sp_io::TestExternalities;
+
  struct Test1;
  struct Test2;
  struct Test3;
 
+ parameter_types! {
+ static Test1Assertions: u8 = 0;
+ static Test2Assertions: u8 = 0;
+ static Test3Assertions: u8 = 0;
+ static EnableSequentialTest: bool = false;
+ static SequentialAssertions: u8 = 0;
+ }
+
  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();
+ Test1Assertions::mutate(|val| *val += 1);
+ if EnableSequentialTest::get() {
+ SequentialAssertions::mutate(|val| *val += 1);
+ }
  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();
+ Test2Assertions::mutate(|val| *val += 1);
+ if EnableSequentialTest::get() {
+ assert_eq!(SequentialAssertions::get(), 1);
+ SequentialAssertions::mutate(|val| *val += 1);
+ }
  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();
+ Test3Assertions::mutate(|val| *val += 1);
+ if EnableSequentialTest::get() {
+ assert_eq!(SequentialAssertions::get(), 2);
+ SequentialAssertions::mutate(|val| *val += 1);
+ }
  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();
+ TestExternalities::default().execute_with(|| {
+ type TestEmpty = ();
+ let origin_state = <TestEmpty as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+ assert!(origin_state.is_empty());
+ <TestEmpty as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+
+ type Test1Tuple = (Test1,);
+ let origin_state = <Test1Tuple as OnRuntimeUpgrade>::pre_upgrade().unwrap();
+ assert!(origin_state.is_empty());
+ <Test1Tuple as OnRuntimeUpgrade>::post_upgrade(origin_state).unwrap();
+ assert_eq!(Test1Assertions::get(), 0);
+ <Test1Tuple as OnRuntimeUpgrade>::on_runtime_upgrade();
+ assert_eq!(Test1Assertions::take(), 1);
+
+ type Test321 = (Test3, Test2, Test1);
+ <Test321 as OnRuntimeUpgrade>::on_runtime_upgrade();
+ assert_eq!(Test1Assertions::take(), 1);
+ assert_eq!(Test2Assertions::take(), 1);
+ assert_eq!(Test3Assertions::take(), 1);
+
+ // enable sequential tests
+ EnableSequentialTest::mutate(|val| *val = true);
+
+ type Test123 = (Test1, Test2, Test3);
+ <Test123 as OnRuntimeUpgrade>::on_runtime_upgrade();
+ assert_eq!(Test1Assertions::take(), 1);
+ assert_eq!(Test2Assertions::take(), 1);
+ assert_eq!(Test3Assertions::take(), 1);
+
+ // reset assertions
+ SequentialAssertions::take();
+
+ type TestNested123 = (Test1, (Test2, Test3));
+ <TestNested123 as OnRuntimeUpgrade>::on_runtime_upgrade();
+ assert_eq!(Test1Assertions::take(), 1);
+ assert_eq!(Test2Assertions::take(), 1);
+ assert_eq!(Test3Assertions::take(), 1);
+ });
  }
 }