Skip to content

Commit

Permalink
[Feature] Sequential migration execution for try-runtime (paritytech#…
Browse files Browse the repository at this point in the history
…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 <5588131+kianenigma@users.noreply.github.com>

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* fix tests

* redo Tuple tests

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* use parameter_types for testing

* lint fix

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Jegor Sidorenko <5252494+jsidorenko@users.noreply.github.com>

* Update frame/support/src/traits/hooks.rs

Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* 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 <5588131+kianenigma@users.noreply.github.com>
Co-authored-by: Jegor Sidorenko <5252494+jsidorenko@users.noreply.github.com>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: parity-processbot <>
  • Loading branch information
4 people authored and ark0f committed Feb 27, 2023
1 parent 6e13922 commit 092a700
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 71 deletions.
8 changes: 0 additions & 8 deletions frame/executive/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
Expand Down
169 changes: 106 additions & 63 deletions frame/support/src/traits/hooks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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(())
}
}
Expand Down Expand Up @@ -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)
Expand All @@ -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]
Expand Down Expand Up @@ -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);
});
}
}

0 comments on commit 092a700

Please sign in to comment.