#![cfg_attr(not(feature = "std"), no_std)]
extern crate alloc;
pub use pallet::*;
pub mod weights;
const LOG_TARGET: &str = "runtime::state-trie-migration";
#[macro_export]
macro_rules! log {
($level:tt, $patter:expr $(, $values:expr)* $(,)?) => {
log::$level!(
target: crate::LOG_TARGET,
concat!("[{:?}] ๐ค ", $patter), frame_system::Pallet::<T>::block_number() $(, $values)*
)
};
}
#[frame_support::pallet]
pub mod pallet {
pub use crate::weights::WeightInfo;
use alloc::{vec, vec::Vec};
use core::ops::Deref;
use frame_support::{
dispatch::{DispatchErrorWithPostInfo, PostDispatchInfo},
ensure,
pallet_prelude::*,
traits::{
fungible::{hold::Balanced, Inspect, InspectHold, Mutate, MutateHold},
tokens::{Fortitude, Precision},
Get,
},
};
use frame_system::{self, pallet_prelude::*};
use sp_core::{
hexdisplay::HexDisplay, storage::well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX,
};
use sp_runtime::{
self,
traits::{Saturating, Zero},
};
pub(crate) type BalanceOf<T> =
<<T as Config>::Currency as Inspect<<T as frame_system::Config>::AccountId>>::Balance;
#[derive(
CloneNoBound,
Encode,
Decode,
scale_info::TypeInfo,
PartialEqNoBound,
EqNoBound,
MaxEncodedLen,
)]
#[scale_info(skip_type_params(MaxKeyLen))]
pub enum Progress<MaxKeyLen: Get<u32>> {
ToStart,
LastKey(BoundedVec<u8, MaxKeyLen>),
Complete,
}
pub type ProgressOf<T> = Progress<<T as Config>::MaxKeyLen>;
#[derive(Clone, Encode, Decode, scale_info::TypeInfo, PartialEq, Eq, MaxEncodedLen)]
#[scale_info(skip_type_params(T))]
pub struct MigrationTask<T: Config> {
pub(crate) progress_top: ProgressOf<T>,
pub(crate) progress_child: ProgressOf<T>,
#[codec(skip)]
pub(crate) dyn_top_items: u32,
#[codec(skip)]
pub(crate) dyn_child_items: u32,
#[codec(skip)]
pub(crate) dyn_size: u32,
pub(crate) size: u32,
pub(crate) top_items: u32,
pub(crate) child_items: u32,
#[codec(skip)]
pub(crate) _ph: core::marker::PhantomData<T>,
}
impl<Size: Get<u32>> core::fmt::Debug for Progress<Size> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Progress::ToStart => f.write_str("To start"),
Progress::LastKey(key) => write!(f, "Last: {:?}", HexDisplay::from(key.deref())),
Progress::Complete => f.write_str("Complete"),
}
}
}
impl<T: Config> core::fmt::Debug for MigrationTask<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("MigrationTask")
.field("top", &self.progress_top)
.field("child", &self.progress_child)
.field("dyn_top_items", &self.dyn_top_items)
.field("dyn_child_items", &self.dyn_child_items)
.field("dyn_size", &self.dyn_size)
.field("size", &self.size)
.field("top_items", &self.top_items)
.field("child_items", &self.child_items)
.finish()
}
}
impl<T: Config> Default for MigrationTask<T> {
fn default() -> Self {
Self {
progress_top: Progress::ToStart,
progress_child: Progress::ToStart,
dyn_child_items: Default::default(),
dyn_top_items: Default::default(),
dyn_size: Default::default(),
_ph: Default::default(),
size: Default::default(),
top_items: Default::default(),
child_items: Default::default(),
}
}
}
impl<T: Config> MigrationTask<T> {
pub(crate) fn finished(&self) -> bool {
matches!(self.progress_top, Progress::Complete)
}
fn exhausted(&self, limits: MigrationLimits) -> bool {
self.dyn_total_items() >= limits.item || self.dyn_size >= limits.size
}
pub(crate) fn dyn_total_items(&self) -> u32 {
self.dyn_child_items.saturating_add(self.dyn_top_items)
}
pub fn migrate_until_exhaustion(
&mut self,
limits: MigrationLimits,
) -> Result<(), Error<T>> {
log!(debug, "running migrations on top of {:?} until {:?}", self, limits);
if limits.item.is_zero() || limits.size.is_zero() {
log!(warn, "limits are zero. stopping");
return Ok(())
}
while !self.exhausted(limits) && !self.finished() {
if let Err(e) = self.migrate_tick() {
log!(error, "migrate_until_exhaustion failed: {:?}", e);
return Err(e)
}
}
self.size = self.size.saturating_add(self.dyn_size);
self.child_items = self.child_items.saturating_add(self.dyn_child_items);
self.top_items = self.top_items.saturating_add(self.dyn_top_items);
log!(debug, "finished with {:?}", self);
Ok(())
}
fn migrate_tick(&mut self) -> Result<(), Error<T>> {
match (&self.progress_top, &self.progress_child) {
(Progress::ToStart, _) => self.migrate_top(),
(Progress::LastKey(_), Progress::LastKey(_)) => {
self.migrate_child()
},
(Progress::LastKey(top_key), Progress::ToStart) => {
if !top_key.starts_with(DEFAULT_CHILD_STORAGE_KEY_PREFIX) {
self.migrate_top()
} else {
self.migrate_child()
}
},
(Progress::LastKey(_), Progress::Complete) => {
self.migrate_top()?;
self.progress_child = Progress::ToStart;
Ok(())
},
(Progress::Complete, _) => {
Ok(())
},
}
}
fn migrate_child(&mut self) -> Result<(), Error<T>> {
use sp_io::default_child_storage as child_io;
let (maybe_current_child, child_root) = match (&self.progress_child, &self.progress_top)
{
(Progress::LastKey(last_child), Progress::LastKey(last_top)) => {
let child_root = Pallet::<T>::transform_child_key_or_halt(last_top);
let maybe_current_child: Option<BoundedVec<u8, T::MaxKeyLen>> =
if let Some(next) = child_io::next_key(child_root, last_child) {
Some(next.try_into().map_err(|_| Error::<T>::KeyTooLong)?)
} else {
None
};
(maybe_current_child, child_root)
},
(Progress::ToStart, Progress::LastKey(last_top)) => {
let child_root = Pallet::<T>::transform_child_key_or_halt(last_top);
(Some(Default::default()), child_root)
},
_ => {
frame_support::defensive!("cannot migrate child key.");
return Ok(())
},
};
if let Some(current_child) = maybe_current_child.as_ref() {
let added_size = if let Some(data) = child_io::get(child_root, current_child) {
child_io::set(child_root, current_child, &data);
data.len() as u32
} else {
Zero::zero()
};
self.dyn_size = self.dyn_size.saturating_add(added_size);
self.dyn_child_items.saturating_inc();
}
log!(trace, "migrated a child key, next_child_key: {:?}", maybe_current_child);
self.progress_child = match maybe_current_child {
Some(last_child) => Progress::LastKey(last_child),
None => Progress::Complete,
};
Ok(())
}
fn migrate_top(&mut self) -> Result<(), Error<T>> {
let maybe_current_top = match &self.progress_top {
Progress::LastKey(last_top) => {
let maybe_top: Option<BoundedVec<u8, T::MaxKeyLen>> =
if let Some(next) = sp_io::storage::next_key(last_top) {
Some(next.try_into().map_err(|_| Error::<T>::KeyTooLong)?)
} else {
None
};
maybe_top
},
Progress::ToStart => Some(Default::default()),
Progress::Complete => {
frame_support::defensive!("cannot migrate top key.");
return Ok(())
},
};
if let Some(current_top) = maybe_current_top.as_ref() {
let added_size = if let Some(data) = sp_io::storage::get(current_top) {
sp_io::storage::set(current_top, &data);
data.len() as u32
} else {
Zero::zero()
};
self.dyn_size = self.dyn_size.saturating_add(added_size);
self.dyn_top_items.saturating_inc();
}
log!(trace, "migrated a top key, next_top_key = {:?}", maybe_current_top);
self.progress_top = match maybe_current_top {
Some(last_top) => Progress::LastKey(last_top),
None => Progress::Complete,
};
Ok(())
}
}
#[derive(
Clone,
Copy,
Encode,
Decode,
scale_info::TypeInfo,
Default,
Debug,
PartialEq,
Eq,
MaxEncodedLen,
)]
pub struct MigrationLimits {
pub size: u32,
pub item: u32,
}
#[derive(Clone, Copy, Encode, Decode, scale_info::TypeInfo, Debug, PartialEq, Eq)]
pub enum MigrationCompute {
Signed,
Auto,
}
#[pallet::event]
#[pallet::generate_deposit(pub(super) fn deposit_event)]
pub enum Event<T: Config> {
Migrated { top: u32, child: u32, compute: MigrationCompute },
Slashed { who: T::AccountId, amount: BalanceOf<T> },
AutoMigrationFinished,
Halted { error: Error<T> },
}
#[pallet::pallet]
pub struct Pallet<T>(_);
pub mod config_preludes {
use super::*;
use frame_support::derive_impl;
pub struct TestDefaultConfig;
#[derive_impl(frame_system::config_preludes::TestDefaultConfig, no_aggregated_types)]
impl frame_system::DefaultConfig for TestDefaultConfig {}
#[frame_support::register_default_impl(TestDefaultConfig)]
impl DefaultConfig for TestDefaultConfig {
#[inject_runtime_type]
type RuntimeEvent = ();
#[inject_runtime_type]
type RuntimeHoldReason = ();
}
}
#[pallet::composite_enum]
pub enum HoldReason {
#[codec(index = 0)]
SlashForMigrate,
}
#[pallet::config(with_default)]
pub trait Config: frame_system::Config {
#[pallet::no_default]
type ControlOrigin: EnsureOrigin<Self::RuntimeOrigin>;
#[pallet::no_default]
type SignedFilter: EnsureOrigin<Self::RuntimeOrigin, Success = Self::AccountId>;
#[pallet::no_default_bounds]
type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
#[pallet::no_default]
type Currency: InspectHold<Self::AccountId>
+ Mutate<Self::AccountId>
+ MutateHold<Self::AccountId, Reason = Self::RuntimeHoldReason>
+ Balanced<Self::AccountId>;
#[pallet::no_default_bounds]
type RuntimeHoldReason: From<HoldReason>;
#[pallet::constant]
#[pallet::no_default]
type MaxKeyLen: Get<u32>;
#[pallet::no_default]
type SignedDepositPerItem: Get<BalanceOf<Self>>;
#[pallet::no_default]
type SignedDepositBase: Get<BalanceOf<Self>>;
#[pallet::no_default]
type WeightInfo: WeightInfo;
}
#[pallet::storage]
#[pallet::getter(fn migration_process)]
pub type MigrationProcess<T> = StorageValue<_, MigrationTask<T>, ValueQuery>;
#[pallet::storage]
#[pallet::getter(fn auto_limits)]
pub type AutoLimits<T> = StorageValue<_, Option<MigrationLimits>, ValueQuery>;
#[pallet::storage]
#[pallet::getter(fn signed_migration_max_limits)]
pub type SignedMigrationMaxLimits<T> = StorageValue<_, MigrationLimits, OptionQuery>;
#[pallet::error]
#[derive(Clone, PartialEq)]
pub enum Error<T> {
MaxSignedLimits,
KeyTooLong,
NotEnoughFunds,
BadWitness,
SignedMigrationNotAllowed,
BadChildRoot,
}
#[pallet::call]
impl<T: Config> Pallet<T> {
#[pallet::call_index(0)]
#[pallet::weight(T::DbWeight::get().reads_writes(1, 1))]
pub fn control_auto_migration(
origin: OriginFor<T>,
maybe_config: Option<MigrationLimits>,
) -> DispatchResult {
T::ControlOrigin::ensure_origin(origin)?;
AutoLimits::<T>::put(maybe_config);
Ok(())
}
#[pallet::call_index(1)]
#[pallet::weight(
Pallet::<T>::dynamic_weight(limits.item, * real_size_upper)
+ T::WeightInfo::continue_migrate()
)]
pub fn continue_migrate(
origin: OriginFor<T>,
limits: MigrationLimits,
real_size_upper: u32,
witness_task: MigrationTask<T>,
) -> DispatchResultWithPostInfo {
let who = T::SignedFilter::ensure_origin(origin)?;
let max_limits =
Self::signed_migration_max_limits().ok_or(Error::<T>::SignedMigrationNotAllowed)?;
ensure!(
limits.size <= max_limits.size && limits.item <= max_limits.item,
Error::<T>::MaxSignedLimits,
);
let deposit = Self::calculate_deposit_for(limits.item);
ensure!(
T::Currency::can_hold(&HoldReason::SlashForMigrate.into(), &who, deposit),
Error::<T>::NotEnoughFunds
);
let mut task = Self::migration_process();
ensure!(
task == witness_task,
DispatchErrorWithPostInfo {
error: Error::<T>::BadWitness.into(),
post_info: PostDispatchInfo {
actual_weight: Some(T::WeightInfo::continue_migrate_wrong_witness()),
pays_fee: Pays::Yes
}
}
);
let migration = task.migrate_until_exhaustion(limits);
if real_size_upper < task.dyn_size {
Self::slash(who, deposit)?;
return Ok(().into())
}
Self::deposit_event(Event::<T>::Migrated {
top: task.dyn_top_items,
child: task.dyn_child_items,
compute: MigrationCompute::Signed,
});
let actual_weight = Some(
Pallet::<T>::dynamic_weight(limits.item, task.dyn_size)
.saturating_add(T::WeightInfo::continue_migrate()),
);
MigrationProcess::<T>::put(task);
let post_info = PostDispatchInfo { actual_weight, pays_fee: Pays::No };
if let Err(error) = migration {
Self::halt(error);
}
Ok(post_info)
}
#[pallet::call_index(2)]
#[pallet::weight(
T::WeightInfo::migrate_custom_top_success()
.max(T::WeightInfo::migrate_custom_top_fail())
.saturating_add(
Pallet::<T>::dynamic_weight(keys.len() as u32, *witness_size)
)
)]
pub fn migrate_custom_top(
origin: OriginFor<T>,
keys: Vec<Vec<u8>>,
witness_size: u32,
) -> DispatchResultWithPostInfo {
let who = T::SignedFilter::ensure_origin(origin)?;
let deposit = Self::calculate_deposit_for(keys.len() as u32);
ensure!(
T::Currency::can_hold(&HoldReason::SlashForMigrate.into(), &who, deposit),
Error::<T>::NotEnoughFunds
);
let mut dyn_size = 0u32;
for key in &keys {
if let Some(data) = sp_io::storage::get(key) {
dyn_size = dyn_size.saturating_add(data.len() as u32);
sp_io::storage::set(key, &data);
}
}
if dyn_size > witness_size {
Self::slash(who, deposit)?;
Ok(().into())
} else {
Self::deposit_event(Event::<T>::Migrated {
top: keys.len() as u32,
child: 0,
compute: MigrationCompute::Signed,
});
Ok(PostDispatchInfo {
actual_weight: Some(
T::WeightInfo::migrate_custom_top_success().saturating_add(
Pallet::<T>::dynamic_weight(keys.len() as u32, dyn_size),
),
),
pays_fee: Pays::Yes,
})
}
}
#[pallet::call_index(3)]
#[pallet::weight(
T::WeightInfo::migrate_custom_child_success()
.max(T::WeightInfo::migrate_custom_child_fail())
.saturating_add(
Pallet::<T>::dynamic_weight(child_keys.len() as u32, *total_size)
)
)]
pub fn migrate_custom_child(
origin: OriginFor<T>,
root: Vec<u8>,
child_keys: Vec<Vec<u8>>,
total_size: u32,
) -> DispatchResultWithPostInfo {
use sp_io::default_child_storage as child_io;
let who = T::SignedFilter::ensure_origin(origin)?;
let deposit = Self::calculate_deposit_for(child_keys.len() as u32);
ensure!(
T::Currency::can_hold(&HoldReason::SlashForMigrate.into(), &who, deposit),
Error::<T>::NotEnoughFunds
);
let mut dyn_size = 0u32;
let transformed_child_key = Self::transform_child_key(&root).ok_or("bad child key")?;
for child_key in &child_keys {
if let Some(data) = child_io::get(transformed_child_key, child_key) {
dyn_size = dyn_size.saturating_add(data.len() as u32);
child_io::set(transformed_child_key, child_key, &data);
}
}
if dyn_size != total_size {
Self::slash(who, deposit)?;
Ok(PostDispatchInfo {
actual_weight: Some(T::WeightInfo::migrate_custom_child_fail()),
pays_fee: Pays::Yes,
})
} else {
Self::deposit_event(Event::<T>::Migrated {
top: 0,
child: child_keys.len() as u32,
compute: MigrationCompute::Signed,
});
Ok(PostDispatchInfo {
actual_weight: Some(
T::WeightInfo::migrate_custom_child_success().saturating_add(
Pallet::<T>::dynamic_weight(child_keys.len() as u32, total_size),
),
),
pays_fee: Pays::Yes,
})
}
}
#[pallet::call_index(4)]
#[pallet::weight(T::DbWeight::get().reads_writes(1, 1))]
pub fn set_signed_max_limits(
origin: OriginFor<T>,
limits: MigrationLimits,
) -> DispatchResult {
let _ = T::ControlOrigin::ensure_origin(origin)?;
SignedMigrationMaxLimits::<T>::put(limits);
Ok(())
}
#[pallet::call_index(5)]
#[pallet::weight(T::DbWeight::get().reads_writes(1, 1))]
pub fn force_set_progress(
origin: OriginFor<T>,
progress_top: ProgressOf<T>,
progress_child: ProgressOf<T>,
) -> DispatchResult {
let _ = T::ControlOrigin::ensure_origin(origin)?;
MigrationProcess::<T>::mutate(|task| {
task.progress_top = progress_top;
task.progress_child = progress_child;
});
Ok(())
}
}
#[pallet::hooks]
impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
fn on_initialize(_: BlockNumberFor<T>) -> Weight {
if let Some(limits) = Self::auto_limits() {
let mut task = Self::migration_process();
if let Err(e) = task.migrate_until_exhaustion(limits) {
Self::halt(e);
}
let weight = Self::dynamic_weight(task.dyn_total_items(), task.dyn_size);
log!(
info,
"migrated {} top keys, {} child keys, and a total of {} bytes.",
task.dyn_top_items,
task.dyn_child_items,
task.dyn_size,
);
if task.finished() {
Self::deposit_event(Event::<T>::AutoMigrationFinished);
AutoLimits::<T>::kill();
} else {
Self::deposit_event(Event::<T>::Migrated {
top: task.dyn_top_items,
child: task.dyn_child_items,
compute: MigrationCompute::Auto,
});
}
MigrationProcess::<T>::put(task);
weight
} else {
T::DbWeight::get().reads(1)
}
}
}
impl<T: Config> Pallet<T> {
fn dynamic_weight(items: u32, size: u32) -> frame_support::pallet_prelude::Weight {
let items = items as u64;
<T as frame_system::Config>::DbWeight::get()
.reads_writes(1, 1)
.saturating_mul(items)
.saturating_add(T::WeightInfo::process_top_key(size))
}
fn halt(error: Error<T>) {
log!(error, "migration halted due to: {:?}", error);
AutoLimits::<T>::kill();
Self::deposit_event(Event::<T>::Halted { error });
}
fn transform_child_key(root: &Vec<u8>) -> Option<&[u8]> {
use sp_core::storage::{ChildType, PrefixedStorageKey};
match ChildType::from_prefixed_key(PrefixedStorageKey::new_ref(root)) {
Some((ChildType::ParentKeyId, root)) => Some(root),
_ => None,
}
}
fn transform_child_key_or_halt(root: &Vec<u8>) -> &[u8] {
let key = Self::transform_child_key(root);
if key.is_none() {
Self::halt(Error::<T>::BadChildRoot);
}
key.unwrap_or_default()
}
#[cfg(any(test, feature = "runtime-benchmarks"))]
pub(crate) fn childify(root: &'static str) -> Vec<u8> {
let mut string = DEFAULT_CHILD_STORAGE_KEY_PREFIX.to_vec();
string.extend_from_slice(root.as_ref());
string
}
pub(crate) fn calculate_deposit_for(keys_count: u32) -> BalanceOf<T> {
T::SignedDepositBase::get()
.saturating_add(T::SignedDepositPerItem::get().saturating_mul(keys_count.into()))
}
fn slash(who: T::AccountId, amount: BalanceOf<T>) -> Result<(), DispatchError> {
T::Currency::hold(&HoldReason::SlashForMigrate.into(), &who, amount)?;
let _burned = T::Currency::burn_all_held(
&HoldReason::SlashForMigrate.into(),
&who,
Precision::BestEffort,
Fortitude::Force,
)?;
debug_assert!(amount.saturating_sub(_burned).is_zero());
Self::deposit_event(Event::<T>::Slashed { who, amount });
Ok(())
}
}
}
#[cfg(feature = "runtime-benchmarks")]
mod benchmarks {
use super::{pallet::Pallet as StateTrieMigration, *};
use alloc::vec;
use frame_support::traits::fungible::{Inspect, Mutate};
const KEY: &[u8] = b"key";
fn set_balance_for_deposit<T: Config>(caller: &T::AccountId, item: u32) -> BalanceOf<T> {
let deposit = StateTrieMigration::<T>::calculate_deposit_for(item);
let stash = T::Currency::minimum_balance() * BalanceOf::<T>::from(1000u32) + deposit;
T::Currency::set_balance(caller, stash);
stash
}
frame_benchmarking::benchmarks! {
continue_migrate {
let null = MigrationLimits::default();
let caller = frame_benchmarking::whitelisted_caller();
let stash = set_balance_for_deposit::<T>(&caller, null.item);
SignedMigrationMaxLimits::<T>::put(MigrationLimits { size: 1024, item: 5 });
}: _(frame_system::RawOrigin::Signed(caller.clone()), null, 0, StateTrieMigration::<T>::migration_process())
verify {
assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
assert_eq!(T::Currency::balance(&caller), stash)
}
continue_migrate_wrong_witness {
let null = MigrationLimits::default();
let caller = frame_benchmarking::whitelisted_caller();
let bad_witness = MigrationTask { progress_top: Progress::LastKey(vec![1u8].try_into().unwrap()), ..Default::default() };
}: {
assert!(
StateTrieMigration::<T>::continue_migrate(
frame_system::RawOrigin::Signed(caller).into(),
null,
0,
bad_witness,
)
.is_err()
)
}
verify {
assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default())
}
migrate_custom_top_success {
let null = MigrationLimits::default();
let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
let stash = set_balance_for_deposit::<T>(&caller, null.item);
}: migrate_custom_top(frame_system::RawOrigin::Signed(caller.clone()), Default::default(), 0)
verify {
assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
assert_eq!(T::Currency::balance(&caller), stash)
}
migrate_custom_top_fail {
let null = MigrationLimits::default();
let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
let stash = set_balance_for_deposit::<T>(&caller, null.item);
sp_io::storage::set(b"foo", vec![1u8;33].as_ref());
}: {
assert!(
StateTrieMigration::<T>::migrate_custom_top(
frame_system::RawOrigin::Signed(caller.clone()).into(),
vec![b"foo".to_vec()],
1,
).is_ok()
);
frame_system::Pallet::<T>::assert_last_event(
<T as Config>::RuntimeEvent::from(crate::Event::Slashed {
who: caller.clone(),
amount: StateTrieMigration::<T>::calculate_deposit_for(1u32),
}).into(),
);
}
verify {
assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
assert!(T::Currency::balance(&caller) < stash)
}
migrate_custom_child_success {
let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
let stash = set_balance_for_deposit::<T>(&caller, 0);
}: migrate_custom_child(
frame_system::RawOrigin::Signed(caller.clone()),
StateTrieMigration::<T>::childify(Default::default()),
Default::default(),
0
)
verify {
assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
assert_eq!(T::Currency::balance(&caller), stash);
}
migrate_custom_child_fail {
let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
let stash = set_balance_for_deposit::<T>(&caller, 1);
sp_io::default_child_storage::set(b"top", b"foo", vec![1u8;33].as_ref());
}: {
assert!(
StateTrieMigration::<T>::migrate_custom_child(
frame_system::RawOrigin::Signed(caller.clone()).into(),
StateTrieMigration::<T>::childify("top"),
vec![b"foo".to_vec()],
1,
).is_ok()
)
}
verify {
assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
assert!(T::Currency::balance(&caller) < stash)
}
process_top_key {
let v in 1 .. (4 * 1024 * 1024);
let value = alloc::vec![1u8; v as usize];
sp_io::storage::set(KEY, &value);
}: {
let data = sp_io::storage::get(KEY).unwrap();
sp_io::storage::set(KEY, &data);
let _next = sp_io::storage::next_key(KEY);
assert_eq!(data, value);
}
impl_benchmark_test_suite!(
StateTrieMigration,
crate::mock::new_test_ext(sp_runtime::StateVersion::V0, true, None, None),
crate::mock::Test
);
}
}
#[cfg(test)]
mod mock {
use super::*;
use crate as pallet_state_trie_migration;
use alloc::{vec, vec::Vec};
use frame_support::{derive_impl, parameter_types, traits::Hooks, weights::Weight};
use frame_system::{EnsureRoot, EnsureSigned};
use sp_core::{
storage::{ChildInfo, StateVersion},
H256,
};
use sp_runtime::{traits::Header as _, BuildStorage, StorageChild};
type Block = frame_system::mocking::MockBlockU32<Test>;
frame_support::construct_runtime!(
pub enum Test
{
System: frame_system,
Balances: pallet_balances,
StateTrieMigration: pallet_state_trie_migration,
}
);
parameter_types! {
pub const SS58Prefix: u8 = 42;
}
#[derive_impl(frame_system::config_preludes::TestDefaultConfig)]
impl frame_system::Config for Test {
type Block = Block;
type AccountData = pallet_balances::AccountData<u64>;
}
parameter_types! {
pub const SignedDepositPerItem: u64 = 1;
pub const SignedDepositBase: u64 = 5;
pub const MigrationMaxKeyLen: u32 = 512;
}
#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)]
impl pallet_balances::Config for Test {
type ReserveIdentifier = [u8; 8];
type AccountStore = System;
}
pub struct StateMigrationTestWeight;
impl WeightInfo for StateMigrationTestWeight {
fn process_top_key(_: u32) -> Weight {
Weight::from_parts(1000000, 0)
}
fn continue_migrate() -> Weight {
Weight::from_parts(1000000, 0)
}
fn continue_migrate_wrong_witness() -> Weight {
Weight::from_parts(1000000, 0)
}
fn migrate_custom_top_fail() -> Weight {
Weight::from_parts(1000000, 0)
}
fn migrate_custom_top_success() -> Weight {
Weight::from_parts(1000000, 0)
}
fn migrate_custom_child_fail() -> Weight {
Weight::from_parts(1000000, 0)
}
fn migrate_custom_child_success() -> Weight {
Weight::from_parts(1000000, 0)
}
}
#[derive_impl(super::config_preludes::TestDefaultConfig)]
impl pallet_state_trie_migration::Config for Test {
type ControlOrigin = EnsureRoot<u64>;
type Currency = Balances;
type MaxKeyLen = MigrationMaxKeyLen;
type SignedDepositPerItem = SignedDepositPerItem;
type SignedDepositBase = SignedDepositBase;
type SignedFilter = EnsureSigned<Self::AccountId>;
type WeightInfo = StateMigrationTestWeight;
}
pub fn new_test_ext(
version: StateVersion,
with_pallets: bool,
custom_keys: Option<Vec<(Vec<u8>, Vec<u8>)>>,
custom_child: Option<Vec<(Vec<u8>, Vec<u8>, Vec<u8>)>>,
) -> sp_io::TestExternalities {
let minimum_size = sp_core::storage::TRIE_VALUE_NODE_THRESHOLD as usize + 1;
let mut custom_storage = sp_core::storage::Storage {
top: vec![
(b"key1".to_vec(), vec![1u8; minimum_size + 1]), (b"key2".to_vec(), vec![1u8; minimum_size + 2]), (b"key3".to_vec(), vec![1u8; minimum_size + 3]), (b"key4".to_vec(), vec![1u8; minimum_size + 4]), (b"key5".to_vec(), vec![1u8; minimum_size + 5]), (b"key6".to_vec(), vec![1u8; minimum_size + 6]), (b"key7".to_vec(), vec![1u8; minimum_size + 7]), (b"key8".to_vec(), vec![1u8; minimum_size + 8]), (b"key9".to_vec(), vec![1u8; minimum_size + 9]), (b"CODE".to_vec(), vec![1u8; minimum_size + 100]), ]
.into_iter()
.chain(custom_keys.unwrap_or_default())
.collect(),
children_default: vec![
(
b"chk1".to_vec(), StorageChild {
data: vec![
(b"key1".to_vec(), vec![1u8; 55]),
(b"key2".to_vec(), vec![2u8; 66]),
]
.into_iter()
.collect(),
child_info: ChildInfo::new_default(b"chk1"),
},
),
(
b"chk2".to_vec(),
StorageChild {
data: vec![
(b"key1".to_vec(), vec![1u8; 54]),
(b"key2".to_vec(), vec![2u8; 64]),
]
.into_iter()
.collect(),
child_info: ChildInfo::new_default(b"chk2"),
},
),
]
.into_iter()
.chain(
custom_child
.unwrap_or_default()
.into_iter()
.map(|(r, k, v)| {
(
r.clone(),
StorageChild {
data: vec![(k, v)].into_iter().collect(),
child_info: ChildInfo::new_default(&r),
},
)
})
.collect::<Vec<_>>(),
)
.collect(),
};
if with_pallets {
frame_system::GenesisConfig::<Test>::default()
.assimilate_storage(&mut custom_storage)
.unwrap();
pallet_balances::GenesisConfig::<Test> { balances: vec![(1, 1000)] }
.assimilate_storage(&mut custom_storage)
.unwrap();
}
sp_tracing::try_init_simple();
(custom_storage, version).into()
}
pub(crate) fn run_to_block(n: u32) -> (H256, Weight) {
let mut root = Default::default();
let mut weight_sum = Weight::zero();
log::trace!(target: LOG_TARGET, "running from {:?} to {:?}", System::block_number(), n);
while System::block_number() < n {
System::set_block_number(System::block_number() + 1);
System::on_initialize(System::block_number());
weight_sum += StateTrieMigration::on_initialize(System::block_number());
root = *System::finalize().state_root();
System::on_finalize(System::block_number());
}
(root, weight_sum)
}
}
#[cfg(test)]
mod test {
use super::{mock::*, *};
use frame_support::assert_ok;
use sp_runtime::{bounded_vec, traits::Bounded, StateVersion};
#[test]
fn fails_if_no_migration() {
let mut ext = new_test_ext(StateVersion::V0, false, None, None);
let root1 = ext.execute_with(|| run_to_block(30).0);
let mut ext2 = new_test_ext(StateVersion::V1, false, None, None);
let root2 = ext2.execute_with(|| run_to_block(30).0);
assert_ne!(root1, root2);
}
#[test]
fn halts_if_top_key_too_long() {
let bad_key = vec![1u8; MigrationMaxKeyLen::get() as usize + 1];
let bad_top_keys = vec![(bad_key.clone(), vec![])];
new_test_ext(StateVersion::V0, true, Some(bad_top_keys), None).execute_with(|| {
System::set_block_number(1);
assert_eq!(MigrationProcess::<Test>::get(), Default::default());
SignedMigrationMaxLimits::<Test>::put(MigrationLimits { size: 1 << 20, item: 50 });
frame_support::assert_ok!(StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(1),
MigrationLimits { item: 50, size: 1 << 20 },
Bounded::max_value(),
MigrationProcess::<Test>::get()
),);
System::assert_last_event(
crate::Event::Halted { error: Error::<Test>::KeyTooLong }.into(),
);
assert!(AutoLimits::<Test>::get().is_none());
let mut task = StateTrieMigration::migration_process();
let result = task.migrate_until_exhaustion(
StateTrieMigration::signed_migration_max_limits().unwrap(),
);
assert!(result.is_err());
});
}
#[test]
fn halts_if_child_key_too_long() {
let bad_key = vec![1u8; MigrationMaxKeyLen::get() as usize + 1];
let bad_child_keys = vec![(bad_key.clone(), vec![], vec![])];
new_test_ext(StateVersion::V0, true, None, Some(bad_child_keys)).execute_with(|| {
System::set_block_number(1);
assert_eq!(MigrationProcess::<Test>::get(), Default::default());
SignedMigrationMaxLimits::<Test>::put(MigrationLimits { size: 1 << 20, item: 50 });
frame_support::assert_ok!(StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(1),
MigrationLimits { item: 50, size: 1 << 20 },
Bounded::max_value(),
MigrationProcess::<Test>::get()
));
System::assert_last_event(
crate::Event::Halted { error: Error::<Test>::KeyTooLong }.into(),
);
assert!(AutoLimits::<Test>::get().is_none());
let mut task = StateTrieMigration::migration_process();
let result = task.migrate_until_exhaustion(
StateTrieMigration::signed_migration_max_limits().unwrap(),
);
assert!(result.is_err());
});
}
#[test]
fn detects_value_in_empty_top_key() {
let limit = MigrationLimits { item: 1, size: 1000 };
let initial_keys = Some(vec![(vec![], vec![66u8; 77])]);
let mut ext = new_test_ext(StateVersion::V0, false, initial_keys.clone(), None);
let root_upgraded = ext.execute_with(|| {
AutoLimits::<Test>::put(Some(limit));
let root = run_to_block(30).0;
assert!(StateTrieMigration::migration_process().finished());
root
});
let mut ext2 = new_test_ext(StateVersion::V1, false, initial_keys, None);
let root = ext2.execute_with(|| {
AutoLimits::<Test>::put(Some(limit));
run_to_block(30).0
});
assert_eq!(root, root_upgraded);
}
#[test]
fn detects_value_in_first_child_key() {
let limit = MigrationLimits { item: 1, size: 1000 };
let initial_child = Some(vec![(b"chk1".to_vec(), vec![], vec![66u8; 77])]);
let mut ext = new_test_ext(StateVersion::V0, false, None, initial_child.clone());
let root_upgraded = ext.execute_with(|| {
AutoLimits::<Test>::put(Some(limit));
let root = run_to_block(30).0;
assert!(StateTrieMigration::migration_process().finished());
root
});
let mut ext2 = new_test_ext(StateVersion::V1, false, None, initial_child);
let root = ext2.execute_with(|| {
AutoLimits::<Test>::put(Some(limit));
run_to_block(30).0
});
assert_eq!(root, root_upgraded);
}
#[test]
fn auto_migrate_works() {
let run_with_limits = |limit, from, until| {
let mut ext = new_test_ext(StateVersion::V0, false, None, None);
let root_upgraded = ext.execute_with(|| {
assert_eq!(AutoLimits::<Test>::get(), None);
assert_eq!(MigrationProcess::<Test>::get(), Default::default());
let _ = run_to_block(from);
assert_eq!(MigrationProcess::<Test>::get(), Default::default());
AutoLimits::<Test>::put(Some(limit));
let root = run_to_block(until).0;
assert!(matches!(
StateTrieMigration::migration_process(),
MigrationTask { progress_top: Progress::Complete, .. }
));
root
});
let mut ext2 = new_test_ext(StateVersion::V1, false, None, None);
let root = ext2.execute_with(|| {
let _ = run_to_block(from);
AutoLimits::<Test>::put(Some(limit));
run_to_block(until).0
});
assert_eq!(root, root_upgraded);
};
run_with_limits(MigrationLimits { item: 1, size: 1000 }, 10, 100);
run_with_limits(MigrationLimits { item: 5, size: 1000 }, 10, 100);
run_with_limits(MigrationLimits { item: 1000, size: 128 }, 10, 100);
run_with_limits(
MigrationLimits { item: Bounded::max_value(), size: Bounded::max_value() },
10,
100,
);
}
#[test]
fn signed_migrate_works() {
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
assert_eq!(MigrationProcess::<Test>::get(), Default::default());
SignedMigrationMaxLimits::<Test>::put(MigrationLimits { size: 1024, item: 5 });
frame_support::assert_err!(
StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(1),
MigrationLimits { item: 5, size: sp_runtime::traits::Bounded::max_value() },
Bounded::max_value(),
MigrationProcess::<Test>::get()
),
Error::<Test>::MaxSignedLimits,
);
frame_support::assert_err!(
StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(2),
MigrationLimits { item: 5, size: 100 },
100,
MigrationProcess::<Test>::get()
),
Error::<Test>::NotEnoughFunds,
);
frame_support::assert_err_ignore_postinfo!(
StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(1),
MigrationLimits { item: 5, size: 100 },
100,
MigrationTask {
progress_top: Progress::LastKey(bounded_vec![1u8]),
..Default::default()
}
),
Error::<Test>::BadWitness,
);
while !MigrationProcess::<Test>::get().finished() {
let mut task = StateTrieMigration::migration_process();
assert_ok!(task.migrate_until_exhaustion(
StateTrieMigration::signed_migration_max_limits().unwrap(),
));
frame_support::assert_ok!(StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(1),
StateTrieMigration::signed_migration_max_limits().unwrap(),
task.dyn_size,
MigrationProcess::<Test>::get()
));
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(Balances::free_balance(&1), 1000);
assert!(matches!(
StateTrieMigration::migration_process(),
MigrationTask { size: x, .. } if x > 0,
));
}
});
}
#[test]
fn continue_migrate_slashing_works() {
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
assert_eq!(MigrationProcess::<Test>::get(), Default::default());
SignedMigrationMaxLimits::<Test>::put(MigrationLimits { size: 1024, item: 5 });
let mut task = StateTrieMigration::migration_process();
assert_ok!(task.migrate_until_exhaustion(
StateTrieMigration::signed_migration_max_limits().unwrap(),
));
frame_support::assert_ok!(StateTrieMigration::continue_migrate(
RuntimeOrigin::signed(1),
StateTrieMigration::signed_migration_max_limits().unwrap(),
task.dyn_size - 1,
MigrationProcess::<Test>::get()
));
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(
Balances::free_balance(&1),
1000 - StateTrieMigration::calculate_deposit_for(5)
);
});
}
#[test]
fn custom_migrate_top_works() {
let correct_witness = 3 + sp_core::storage::TRIE_VALUE_NODE_THRESHOLD * 3 + 1 + 2 + 3;
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
frame_support::assert_ok!(StateTrieMigration::migrate_custom_top(
RuntimeOrigin::signed(1),
vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()],
correct_witness,
));
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(Balances::free_balance(&1), 1000);
});
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
frame_support::assert_ok!(StateTrieMigration::migrate_custom_top(
RuntimeOrigin::signed(1),
vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()],
correct_witness + 99,
));
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(Balances::free_balance(&1), 1000);
});
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
assert_eq!(Balances::free_balance(&1), 1000);
frame_support::assert_ok!(StateTrieMigration::migrate_custom_top(
RuntimeOrigin::signed(1),
vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()],
correct_witness - 1,
),);
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(
Balances::free_balance(&1),
1000 - StateTrieMigration::calculate_deposit_for(3)
);
});
}
#[test]
fn custom_migrate_child_works() {
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
frame_support::assert_ok!(StateTrieMigration::migrate_custom_child(
RuntimeOrigin::signed(1),
StateTrieMigration::childify("chk1"),
vec![b"key1".to_vec(), b"key2".to_vec()],
55 + 66,
));
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(Balances::free_balance(&1), 1000);
});
new_test_ext(StateVersion::V0, true, None, None).execute_with(|| {
assert_eq!(Balances::free_balance(&1), 1000);
frame_support::assert_ok!(StateTrieMigration::migrate_custom_child(
RuntimeOrigin::signed(1),
StateTrieMigration::childify("chk1"),
vec![b"key1".to_vec(), b"key2".to_vec()],
999999, ));
assert_eq!(Balances::reserved_balance(&1), 0);
assert_eq!(
Balances::free_balance(&1),
1000 - StateTrieMigration::calculate_deposit_for(2)
);
});
}
}
#[cfg(feature = "remote-test")]
pub(crate) mod remote_tests {
use crate::{AutoLimits, MigrationLimits, Pallet as StateTrieMigration, LOG_TARGET};
use codec::Encode;
use frame_support::{
traits::{Get, Hooks},
weights::Weight,
};
use frame_system::{pallet_prelude::BlockNumberFor, Pallet as System};
use remote_externalities::Mode;
use sp_core::H256;
use sp_runtime::{
traits::{Block as BlockT, HashingFor, Header as _, One, Zero},
DeserializeOwned,
};
use thousands::Separable;
#[allow(dead_code)]
fn run_to_block<Runtime: crate::Config<Hash = H256>>(
n: BlockNumberFor<Runtime>,
) -> (H256, Weight) {
let mut root = Default::default();
let mut weight_sum = Weight::zero();
while System::<Runtime>::block_number() < n {
System::<Runtime>::set_block_number(System::<Runtime>::block_number() + One::one());
System::<Runtime>::on_initialize(System::<Runtime>::block_number());
weight_sum +=
StateTrieMigration::<Runtime>::on_initialize(System::<Runtime>::block_number());
root = *System::<Runtime>::finalize().state_root();
System::<Runtime>::on_finalize(System::<Runtime>::block_number());
}
(root, weight_sum)
}
#[allow(dead_code)]
pub(crate) async fn run_with_limits<Runtime, Block>(
limits: MigrationLimits,
mode: Mode<Block::Hash>,
) where
Runtime: crate::Config<Hash = H256>,
Block: BlockT<Hash = H256> + DeserializeOwned,
Block::Header: serde::de::DeserializeOwned,
{
let mut ext = remote_externalities::Builder::<Block>::new()
.mode(mode)
.overwrite_state_version(sp_core::storage::StateVersion::V0)
.build()
.await
.unwrap();
let mut now = ext.execute_with(|| {
AutoLimits::<Runtime>::put(Some(limits));
frame_system::Pallet::<Runtime>::block_number()
});
let mut duration: BlockNumberFor<Runtime> = Zero::zero();
ext.state_version = sp_core::storage::StateVersion::V1;
let status =
substrate_state_trie_migration_rpc::migration_status(&ext.as_backend()).unwrap();
assert!(
status.top_remaining_to_migrate > 0,
"no node needs migrating, this probably means that state was initialized with `StateVersion::V1`",
);
log::info!(
target: LOG_TARGET,
"initial check: top_left: {}, child_left: {}, total_top {}, total_child {}",
status.top_remaining_to_migrate.separate_with_commas(),
status.child_remaining_to_migrate.separate_with_commas(),
status.total_top.separate_with_commas(),
status.total_child.separate_with_commas(),
);
loop {
let last_state_root = *ext.backend.root();
let ((finished, weight), proof) = ext.execute_and_prove(|| {
let weight = run_to_block::<Runtime>(now + One::one()).1;
if StateTrieMigration::<Runtime>::migration_process().finished() {
return (true, weight)
}
duration += One::one();
now += One::one();
(false, weight)
});
let compact_proof =
proof.clone().into_compact_proof::<HashingFor<Block>>(last_state_root).unwrap();
log::info!(
target: LOG_TARGET,
"proceeded to #{}, weight: [{} / {}], proof: [{} / {} / {}]",
now,
weight.separate_with_commas(),
<Runtime as frame_system::Config>::BlockWeights::get()
.max_block
.separate_with_commas(),
proof.encoded_size().separate_with_commas(),
compact_proof.encoded_size().separate_with_commas(),
zstd::stream::encode_all(&compact_proof.encode()[..], 0)
.unwrap()
.len()
.separate_with_commas(),
);
ext.commit_all().unwrap();
if finished {
break
}
}
ext.execute_with(|| {
log::info!(
target: LOG_TARGET,
"finished on_initialize migration in {} block, final state of the task: {:?}",
duration,
StateTrieMigration::<Runtime>::migration_process(),
)
});
let status =
substrate_state_trie_migration_rpc::migration_status(&ext.as_backend()).unwrap();
assert_eq!(status.top_remaining_to_migrate, 0);
assert_eq!(status.child_remaining_to_migrate, 0);
}
}
#[cfg(all(test, feature = "remote-test"))]
mod remote_tests_local {
use super::{
mock::{RuntimeCall as MockCall, *},
remote_tests::run_with_limits,
*,
};
use remote_externalities::{Mode, OfflineConfig, OnlineConfig, SnapshotConfig};
use sp_runtime::traits::Bounded;
use std::env::var as env_var;
type Extrinsic = sp_runtime::testing::TestXt<MockCall, ()>;
type Block = sp_runtime::testing::Block<Extrinsic>;
#[tokio::test]
async fn on_initialize_migration() {
let snap: SnapshotConfig = env_var("SNAP").expect("Need SNAP env var").into();
let ws_api = env_var("WS_API").expect("Need WS_API env var").into();
sp_tracing::try_init_simple();
let mode = Mode::OfflineOrElseOnline(
OfflineConfig { state_snapshot: snap.clone() },
OnlineConfig { transport: ws_api, state_snapshot: Some(snap), ..Default::default() },
);
run_with_limits::<Test, Block>(
MigrationLimits { item: 8 * 1024, size: 128 * 1024 * 1024 },
mode.clone(),
)
.await;
run_with_limits::<Test, Block>(
MigrationLimits { item: Bounded::max_value(), size: 64 * 1024 },
mode,
)
.await;
}
}