Skip to content

Commit

Permalink
keep track of how many keys have been migrated
Browse files Browse the repository at this point in the history
  • Loading branch information
@TarekkMA
TarekkMA committed Sep 16, 2024
1 parent db60dea commit 381cfa4
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 26 deletions.
23 changes: 14 additions & 9 deletions pallets/moonbeam-lazy-migrations/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ pub mod pallet {

#[pallet::storage]
pub(crate) type StateMigrationStatusValue<T: Config> =
StorageValue<_, StateMigrationStatus, ValueQuery>;
StorageValue<_, (StateMigrationStatus, u64), ValueQuery>;

pub(crate) type StorageKey = BoundedVec<u8, ConstU32<1_024>>;

Expand Down Expand Up @@ -147,6 +147,7 @@ pub mod pallet {
struct StateMigrationResult {
last_key: Option<StorageKey>,
error: Option<&'static str>,
migrated: u64,
reads: u64,
writes: u64,
}
Expand All @@ -172,7 +173,7 @@ pub mod pallet {
return read_write_ops;
}

let status = StateMigrationStatusValue::<T>::get();
let (status, mut migrated_keys) = StateMigrationStatusValue::<T>::get();
read_write_ops.add_one_read();

let next_key = match &status {
Expand All @@ -183,14 +184,14 @@ pub mod pallet {
match next_key_result {
NextKeyResult::NextKey(next_key) => next_key,
NextKeyResult::NoMoreKeys => {
StateMigrationStatusValue::<T>::put(StateMigrationStatus::Complete);
StateMigrationStatusValue::<T>::put((StateMigrationStatus::Complete, migrated_keys));
read_write_ops.add_one_write();
return read_write_ops;
}
NextKeyResult::Error(e) => {
StateMigrationStatusValue::<T>::put(StateMigrationStatus::Error(
StateMigrationStatusValue::<T>::put((StateMigrationStatus::Error(
e.as_bytes().to_vec().try_into().unwrap_or_default(),
));
), migrated_keys));
read_write_ops.add_one_write();
return read_write_ops;
}
Expand All @@ -202,23 +203,24 @@ pub mod pallet {
};

let res = Pallet::<T>::migrate_keys(next_key, migration_limit);
migrated_keys += res.migrated;
read_write_ops.add_reads(res.reads);
read_write_ops.add_writes(res.writes);

match (res.last_key, res.error) {
(None, None) => {
StateMigrationStatusValue::<T>::put(StateMigrationStatus::Complete);
StateMigrationStatusValue::<T>::put((StateMigrationStatus::Complete, migrated_keys));
read_write_ops.add_one_write();
}
// maybe we should store the previous key in the storage as well
(_, Some(e)) => {
StateMigrationStatusValue::<T>::put(StateMigrationStatus::Error(
StateMigrationStatusValue::<T>::put((StateMigrationStatus::Error(
e.as_bytes().to_vec().try_into().unwrap_or_default(),
));
), migrated_keys));
read_write_ops.add_one_write();
}
(Some(key), None) => {
StateMigrationStatusValue::<T>::put(StateMigrationStatus::Started(key));
StateMigrationStatusValue::<T>::put((StateMigrationStatus::Started(key), migrated_keys));
read_write_ops.add_one_write();
}
}
Expand Down Expand Up @@ -276,6 +278,7 @@ pub mod pallet {
return StateMigrationResult {
last_key: None,
error: None,
migrated,
reads: migrated + next_key_reads,
writes,
};
Expand All @@ -284,6 +287,7 @@ pub mod pallet {
return StateMigrationResult {
last_key: Some(key),
error: Some(e),
migrated,
reads: migrated + next_key_reads,
writes,
};
Expand All @@ -295,6 +299,7 @@ pub mod pallet {
StateMigrationResult {
last_key: Some(key),
error: None,
migrated,
reads: migrated + next_key_reads,
writes,
}
Expand Down
35 changes: 18 additions & 17 deletions pallets/moonbeam-lazy-migrations/src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,7 @@ fn test_state_migration_baseline() {
ExtBuilder::default().build().execute_with(|| {
assert_eq!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::NotStarted
(StateMigrationStatus::NotStarted, 0)
);

let (keys, data) = count_keys_and_data_without_code();
Expand All @@ -366,15 +366,15 @@ fn test_state_migration_baseline() {

assert_eq!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::Complete
(StateMigrationStatus::Complete, keys)
);
})
}

#[test]
fn test_state_migration_cannot_fit_any_item() {
ExtBuilder::default().build().execute_with(|| {
StateMigrationStatusValue::<Test>::put(StateMigrationStatus::Complete);
StateMigrationStatusValue::<Test>::put((StateMigrationStatus::NotStarted, 0));

let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(0));

Expand All @@ -385,7 +385,7 @@ fn test_state_migration_cannot_fit_any_item() {
#[test]
fn test_state_migration_when_complete() {
ExtBuilder::default().build().execute_with(|| {
StateMigrationStatusValue::<Test>::put(StateMigrationStatus::Complete);
StateMigrationStatusValue::<Test>::put((StateMigrationStatus::Complete, 0));

let weight = LazyMigrations::on_idle(0, Weight::max_value());

Expand All @@ -397,9 +397,9 @@ fn test_state_migration_when_complete() {
#[test]
fn test_state_migration_when_errored() {
ExtBuilder::default().build().execute_with(|| {
StateMigrationStatusValue::<Test>::put(StateMigrationStatus::Error(
StateMigrationStatusValue::<Test>::put((StateMigrationStatus::Error(
"Error".as_bytes().to_vec().try_into().unwrap_or_default(),
));
), 1));

let weight = LazyMigrations::on_idle(0, Weight::max_value());

Expand All @@ -413,7 +413,7 @@ fn test_state_migration_can_only_fit_one_item() {
ExtBuilder::default().build().execute_with(|| {
assert_eq!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::NotStarted
(StateMigrationStatus::NotStarted, 0)
);

let data = sp_io::storage::get(Default::default());
Expand All @@ -425,7 +425,7 @@ fn test_state_migration_can_only_fit_one_item() {

assert!(matches!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::Started(_)
(StateMigrationStatus::Started(_), 1)
));

let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(3));
Expand All @@ -440,7 +440,7 @@ fn test_state_migration_can_only_fit_three_item() {
ExtBuilder::default().build().execute_with(|| {
assert_eq!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::NotStarted
(StateMigrationStatus::NotStarted, 0)
);

let weight = LazyMigrations::on_idle(0, rem_weight_for_entries(3));
Expand All @@ -454,7 +454,7 @@ fn test_state_migration_can_only_fit_three_item() {

assert!(matches!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::Started(_)
(StateMigrationStatus::Started(_), 3)
));
})
}
Expand All @@ -464,7 +464,7 @@ fn test_state_migration_can_fit_exactly_all_item() {
ExtBuilder::default().build().execute_with(|| {
assert_eq!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::NotStarted
(StateMigrationStatus::NotStarted, 0)
);

let (keys, data) = count_keys_and_data_without_code();
Expand All @@ -476,7 +476,7 @@ fn test_state_migration_can_fit_exactly_all_item() {

assert!(matches!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::Started(_),
(StateMigrationStatus::Started(_), keys),
));

// after calling on_idle status is added to the storage so we need to account for that
Expand All @@ -491,7 +491,7 @@ fn test_state_migration_can_fit_exactly_all_item() {

assert!(matches!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::Complete,
(StateMigrationStatus::Complete, new_keys),
));
})
}
Expand All @@ -501,14 +501,14 @@ fn test_state_migration_will_migrate_10_000_items() {
ExtBuilder::default().build().execute_with(|| {
assert_eq!(
StateMigrationStatusValue::<Test>::get(),
StateMigrationStatus::NotStarted
(StateMigrationStatus::NotStarted, 0)
);

for i in 0..100 {
mock_contract_with_entries(i as u8, i as u64, 100);
}

StateMigrationStatusValue::<Test>::put(StateMigrationStatus::NotStarted);
StateMigrationStatusValue::<Test>::put((StateMigrationStatus::NotStarted, 0));

let (keys, data) = count_keys_and_data_without_code();

Expand Down Expand Up @@ -545,16 +545,17 @@ fn test_state_migration_will_migrate_10_000_items() {

let status = StateMigrationStatusValue::<Test>::get();
if i < needed_on_idle_calls {
let migrated_so_far = i * entries_per_on_idle;
assert!(
matches!(status, StateMigrationStatus::Started(_)),
matches!(status, (StateMigrationStatus::Started(_), migrated_so_far)),
"Status: {:?} at call: #{} doesn't match Started",
status,
i,
);
assert!(weight.all_gte(weight_for(1, 0)));
} else {
assert!(
matches!(status, StateMigrationStatus::Complete),
matches!(status, (StateMigrationStatus::Complete, keys)),
"Status: {:?} at call: {} doesn't match Complete",
status,
i,
Expand Down

0 comments on commit 381cfa4

Please sign in to comment.