Do not write old state to storage in pre-migration check (#807)

pallets/aleph/src/migrations/v0_to_v1.rs

@@ -1,15 +1,16 @@
-#[cfg(feature = "try-runtime")]
-use frame_support::ensure;
 use frame_support::{
     log, storage_alias,
     traits::{Get, OnRuntimeUpgrade, PalletInfoAccess, StorageVersion},
     weights::Weight,
 };
-#[cfg(feature = "try-runtime")]
-use pallets_support::ensure_storage_version;
-use pallets_support::StorageMigration;
 use primitives::SessionIndex;
 use sp_std::vec::Vec;
+#[cfg(feature = "try-runtime")]
+use {
+    codec::{Decode, Encode},
+    frame_support::ensure,
+    pallets_support::ensure_storage_version,
+};
 
 use crate::Config;
 
@@ -24,9 +25,11 @@ type Validators<T> = StorageValue<Aleph, Accounts<T>>;
 /// Flattening double `Option<>` storage.
 pub struct Migration<T, P>(sp_std::marker::PhantomData<(T, P)>);
 
-impl<T: Config, P: PalletInfoAccess> StorageMigration for Migration<T, P> {
-    #[cfg(feature = "try-runtime")]
-    const MIGRATION_STORAGE_PREFIX: &'static [u8] = b"PALLET_ALEPH::V0_TO_V1_MIGRATION";
+#[cfg(feature = "try-runtime")]
+#[derive(Decode, Encode)]
+struct MigrationChecksState<T: Config> {
+    session: Option<Option<SessionIndex>>,
+    validators: Option<Option<Accounts<T>>>,
 }
 
 impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
@@ -87,39 +90,50 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
 
         ensure_storage_version::<P>(0)?;
 
-        Self::store_temp("session", SessionForValidatorsChange::get());
-        Self::store_temp("validators", Validators::<T>::get());
+        let session = SessionForValidatorsChange::get();
+        let validators = Validators::<T>::get();
 
-        Ok(Vec::new())
+        Ok(MigrationChecksState::<T> {
+            session,
+            validators,
+        }
+        .encode())
     }
 
     #[cfg(feature = "try-runtime")]
-    fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
+    fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
         ensure_storage_version::<P>(1)?;
 
+        let MigrationChecksState {
+            session: old_session,
+            validators: old_validators,
+        } = <MigrationChecksState<T>>::decode(&mut &*state)
+            .map_err(|_| "Failed to decode old state")?;
+
         let new_session = SessionForValidatorsChange::get();
-        let old_session = Self::read_temp::<Option<Option<SessionIndex>>>("session");
 
         match old_session {
             Some(Some(session)) => ensure!(
                 Some(session) == new_session,
                 "Mismatch on `SessionForValidatorsChange`",
             ),
             _ => ensure!(
-                None == new_session,
+                new_session.is_none(),
                 "New `SessionForValidatorsChange` should be `None`"
             ),
         };
 
         let new_validators = Validators::<T>::get();
-        let old_validators = Self::read_temp::<Option<Option<Accounts<T>>>>("validators");
 
         match old_validators {
             Some(Some(validators)) => ensure!(
                 Some(validators) == new_validators,
                 "Mismatch on `Validators`",
             ),
-            _ => ensure!(None == new_validators, "New `Validators` should be `None`"),
+            _ => ensure!(
+                new_validators.is_none(),
+                "New `Validators` should be `None`"
+            ),
         };
 
         Ok(())

pallets/aleph/src/migrations/v1_to_v2.rs

@@ -1,15 +1,10 @@
-#[cfg(feature = "try-runtime")]
-use frame_support::ensure;
 use frame_support::{
     log, storage_alias,
     traits::{Get, OnRuntimeUpgrade, PalletInfoAccess, StorageVersion},
     weights::Weight,
 };
 #[cfg(feature = "try-runtime")]
-use pallets_support::ensure_storage_version;
-use pallets_support::StorageMigration;
-#[cfg(feature = "try-runtime")]
-use sp_std::vec::Vec;
+use {frame_support::ensure, pallets_support::ensure_storage_version, sp_std::vec::Vec};
 
 use crate::Config;
 
@@ -32,11 +27,6 @@ type Validators = StorageValue<Aleph, ()>;
 ///   - Validators
 pub struct Migration<T, P>(sp_std::marker::PhantomData<(T, P)>);
 
-impl<T: Config, P: PalletInfoAccess> StorageMigration for Migration<T, P> {
-    #[cfg(feature = "try-runtime")]
-    const MIGRATION_STORAGE_PREFIX: &'static [u8] = b"PALLET_ALEPH::V1_TO_V2_MIGRATION";
-}
-
 impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
     fn on_runtime_upgrade() -> Weight {
         let mut writes = 0;

pallets/elections/src/migrations/v0_to_v1.rs

@@ -1,14 +1,15 @@
-#[cfg(feature = "try-runtime")]
-use frame_support::ensure;
 use frame_support::{
     log, storage_alias,
     traits::{Get, OnRuntimeUpgrade, PalletInfoAccess, StorageVersion},
     weights::Weight,
 };
-#[cfg(feature = "try-runtime")]
-use pallets_support::ensure_storage_version;
-use pallets_support::StorageMigration;
 use sp_std::vec::Vec;
+#[cfg(feature = "try-runtime")]
+use {
+    codec::{Decode, Encode},
+    frame_support::ensure,
+    pallets_support::ensure_storage_version,
+};
 
 use crate::{
     compute_validator_scaled_total_rewards,
@@ -39,11 +40,6 @@ type ErasMembers<T> = StorageValue<Elections, (Validators<T>, Validators<T>)>;
 /// - `ErasMembers` contain tuple of (content of `Members`, empty vector).
 pub struct Migration<T, P>(sp_std::marker::PhantomData<(T, P)>);
 
-impl<T: Config, P: PalletInfoAccess> StorageMigration for Migration<T, P> {
-    #[cfg(feature = "try-runtime")]
-    const MIGRATION_STORAGE_PREFIX: &'static [u8] = b"PALLET_ELECTIONS::V0_TO_V1_MIGRATION";
-}
-
 impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
     fn on_runtime_upgrade() -> Weight {
         log::info!(target: "pallet_elections", "Running migration from STORAGE_VERSION 0 to 1 for pallet elections");
@@ -85,15 +81,12 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
     #[cfg(feature = "try-runtime")]
     fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
         ensure_storage_version::<P>(0)?;
-
-        let members = Members::<T>::get().ok_or("No `Members` storage")?;
-        Self::store_temp("members", members);
-
-        Ok(Vec::new())
+        let members: Validators<T> = Members::<T>::get().ok_or("No `Members` storage")?;
+        Ok(members.encode())
     }
 
     #[cfg(feature = "try-runtime")]
-    fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
+    fn post_upgrade(old_members: Vec<u8>) -> Result<(), &'static str> {
         ensure_storage_version::<P>(1)?;
 
         let mps = MembersPerSession::get().ok_or("No `MembersPerSession` in the storage")?;
@@ -103,7 +96,8 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
             NonReservedMembers::<T>::get().ok_or("No `NonReservedMembers` in the storage")?;
         let eras_members = ErasMembers::<T>::get().ok_or("No `ErasMembers` in the storage")?;
 
-        let old_members = Self::read_temp::<Validators<T>>("members");
+        let old_members = <Validators<T>>::decode(&mut &*old_members)
+            .map_err(|_| "Failed to decode old members set")?;
 
         ensure!(
             reserved_members == old_members,

pallets/elections/src/migrations/v1_to_v2.rs

@@ -1,15 +1,15 @@
-#[cfg(feature = "try-runtime")]
-use frame_support::ensure;
 use frame_support::{
     log, storage_alias,
     traits::{Get, OnRuntimeUpgrade, PalletInfoAccess, StorageVersion},
     weights::Weight,
 };
 #[cfg(feature = "try-runtime")]
-use pallets_support::ensure_storage_version;
-use pallets_support::StorageMigration;
-#[cfg(feature = "try-runtime")]
-use sp_std::vec::Vec;
+use {
+    codec::{Decode, Encode},
+    frame_support::ensure,
+    pallets_support::ensure_storage_version,
+    sp_std::vec::Vec,
+};
 
 use crate::{migrations::Validators, Config, EraValidators};
 
@@ -42,9 +42,13 @@ type CurrentEraValidators<T> =
 /// - `ErasMembers` `(reserved, non_reserved)` -> `CurrentEraValidators` `ErasValidators { reserved, non_reserved}`
 pub struct Migration<T, P>(sp_std::marker::PhantomData<(T, P)>);
 
-impl<T: Config, P: PalletInfoAccess> StorageMigration for Migration<T, P> {
-    #[cfg(feature = "try-runtime")]
-    const MIGRATION_STORAGE_PREFIX: &'static [u8] = b"PALLET_ELECTIONS::V1_TO_V2_MIGRATION";
+#[cfg(feature = "try-runtime")]
+#[derive(Decode, Encode)]
+struct MigrationChecksState<T: Config> {
+    members_per_session: u32,
+    reserved_members: Validators<T>,
+    non_reserved_members: Validators<T>,
+    eras_members: (Validators<T>, Validators<T>),
 }
 
 impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
@@ -89,24 +93,23 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
 
         let members_per_session =
             MembersPerSession::get().ok_or("No `MembersPerSession` in the storage")?;
-        Self::store_temp("members_per_session", members_per_session);
-
         let reserved_members =
             ReservedMembers::<T>::get().ok_or("No `ReservedMembers` in the storage")?;
-        Self::store_temp("reserved_members", reserved_members);
-
         let non_reserved_members =
             NonReservedMembers::<T>::get().ok_or("No `NonReservedMembers` in the storage")?;
-        Self::store_temp("non_reserved_members", non_reserved_members);
-
         let eras_members = ErasMembers::<T>::get().ok_or("No `ErasMembers` in the storage")?;
-        Self::store_temp("eras_members", eras_members);
 
-        Ok(Vec::new())
+        Ok(MigrationChecksState::<T> {
+            members_per_session,
+            reserved_members,
+            non_reserved_members,
+            eras_members,
+        }
+        .encode())
     }
 
     #[cfg(feature = "try-runtime")]
-    fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
+    fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
         ensure_storage_version::<P>(2)?;
 
         let committee_size = CommitteeSize::get().ok_or("No `CommitteeSize` in the storage")?;
@@ -117,10 +120,13 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
         let current_era_validators =
             CurrentEraValidators::<T>::get().ok_or("No `CurrentEraValidators` in the storage")?;
 
-        let members_per_session = Self::read_temp::<u32>("members_per_session");
-        let reserved_members = Self::read_temp::<Validators<T>>("reserved_members");
-        let non_reserved_members = Self::read_temp::<Validators<T>>("non_reserved_members");
-        let eras_members = Self::read_temp::<(Validators<T>, Validators<T>)>("eras_members");
+        let MigrationChecksState {
+            members_per_session,
+            reserved_members,
+            non_reserved_members,
+            eras_members,
+        } = <MigrationChecksState<T>>::decode(&mut &*state)
+            .map_err(|_| "Failed to decode old state")?;
 
         ensure!(
             committee_size == members_per_session,

pallets/elections/src/migrations/v2_to_v3.rs

@@ -1,16 +1,16 @@
-#[cfg(feature = "try-runtime")]
-use frame_support::ensure;
 use frame_support::{
     log, storage_alias,
     traits::{Get, OnRuntimeUpgrade, PalletInfoAccess, StorageVersion},
     weights::Weight,
 };
-#[cfg(feature = "try-runtime")]
-use pallets_support::ensure_storage_version;
-use pallets_support::StorageMigration;
 use primitives::CommitteeSeats;
 #[cfg(feature = "try-runtime")]
-use sp_std::vec::Vec;
+use {
+    codec::{Decode, Encode},
+    frame_support::ensure,
+    pallets_support::ensure_storage_version,
+    sp_std::vec::Vec,
+};
 
 use crate::{migrations::Validators, Config, EraValidators};
 
@@ -31,9 +31,11 @@ type NextEraCommitteeSize = StorageValue<Elections, CommitteeSeats>;
 /// `CommitteeSeats`.
 pub struct Migration<T, P>(sp_std::marker::PhantomData<(T, P)>);
 
-impl<T: Config, P: PalletInfoAccess> StorageMigration for Migration<T, P> {
-    #[cfg(feature = "try-runtime")]
-    const MIGRATION_STORAGE_PREFIX: &'static [u8] = b"PALLET_ELECTIONS::V2_TO_V3_MIGRATION";
+#[cfg(feature = "try-runtime")]
+#[derive(Decode, Encode)]
+struct MigrationChecksState {
+    committee_size: Option<u32>,
+    next_era_committee_size: Option<u32>,
 }
 
 impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
@@ -108,16 +110,17 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
         ensure_storage_version::<P>(2)?;
 
         let committee_size = CommitteeSize::get();
-        Self::store_temp("committee_size", committee_size);
-
         let next_era_committee_size = NextEraCommitteeSize::get();
-        Self::store_temp("next_era_committee_size", next_era_committee_size);
 
-        Ok(Vec::new())
+        Ok(MigrationChecksState {
+            committee_size,
+            next_era_committee_size,
+        }
+        .encode())
     }
 
     #[cfg(feature = "try-runtime")]
-    fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
+    fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
         ensure_storage_version::<P>(3)?;
 
         let new_committee_size = CommitteeSize::get().ok_or("No `CommitteeSize` in the storage")?;
@@ -129,10 +132,11 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
         let next_era_reserved_validators = NextEraReservedValidators::<T>::get()
             .ok_or("No `NextEraReservedValidators` in the storage")?;
 
-        let old_committee_size =
-            Self::read_temp::<Option<u32>>("committee_size").unwrap_or_default();
-        let old_next_era_committee_size =
-            Self::read_temp::<Option<u32>>("next_era_committee_size").unwrap_or_default();
+        let MigrationChecksState {
+            committee_size: old_committee_size,
+            next_era_committee_size: old_next_era_committee_size,
+        } = <MigrationChecksState>::decode(&mut &*state)
+            .map_err(|_| "Failed to decode old state")?;
 
         let currently_reserved = current_era_validators.reserved.len();
         ensure!(
@@ -141,7 +145,9 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
         );
         ensure!(
             new_committee_size.non_reserved_seats
-                == old_committee_size.saturating_sub(currently_reserved as u32),
+                == old_committee_size
+                    .unwrap_or_default()
+                    .saturating_sub(currently_reserved as u32),
             "Mismatch between `CurrentEraValidators` and `CommitteeSize`"
         );
 
@@ -152,7 +158,9 @@ impl<T: Config, P: PalletInfoAccess> OnRuntimeUpgrade for Migration<T, P> {
         );
         ensure!(
             new_next_era_committee_size.non_reserved_seats
-                == old_next_era_committee_size.saturating_sub(next_reserved as u32),
+                == old_next_era_committee_size
+                    .unwrap_or_default()
+                    .saturating_sub(next_reserved as u32),
             "Mismatch between `NextEraReservedValidators` and `NextEraCommitteeSize`"
         );