pallet-migrations: try-runtime support, and manage author mapping blake2 migration

pallets/author-mapping/src/lib.rs

@@ -37,6 +37,8 @@ mod mock;
 #[cfg(test)]
 mod tests;
 
+pub mod migrations;
+
 #[pallet]
 pub mod pallet {
 	use crate::WeightInfo;
@@ -237,7 +239,7 @@ pub mod pallet {
 	#[pallet::getter(fn account_and_deposit_of)]
 	/// We maintain a mapping from the AuthorIds used in the consensus layer
 	/// to the AccountIds runtime (including this staking pallet).
-	type MappingWithDeposit<T: Config> = StorageMap<
+	pub type MappingWithDeposit<T: Config> = StorageMap<
 		_,
 		Blake2_128Concat,
 		T::AuthorId,
@@ -283,119 +285,4 @@ pub mod pallet {
 			Self::account_and_deposit_of(author_id).map(|info| info.account)
 		}
 	}
-
-	#[pallet::hooks]
-	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-		fn on_runtime_upgrade() -> Weight {
-			use frame_support::storage::migration::{remove_storage_prefix, storage_key_iter};
-			use sp_std::{convert::TryInto, vec::Vec};
-
-			let pallet_prefix: &[u8] = b"AuthorMapping";
-			let storage_item_prefix: &[u8] = b"MappingWithDeposit";
-
-			// Read all the data into memory.
-			// https://crates.parity.io/frame_support/storage/migration/fn.storage_key_iter.html
-			let stored_data: Vec<_> = storage_key_iter::<
-				T::AuthorId,
-				RegistrationInfo<T::AccountId, BalanceOf<T>>,
-				Twox64Concat,
-			>(pallet_prefix, storage_item_prefix)
-			.collect();
-
-			let migrated_count: Weight = stored_data
-				.len()
-				.try_into()
-				.expect("There are between 0 and 2**64 mappings stored.");
-
-			// Now remove the old storage
-			// https://crates.parity.io/frame_support/storage/migration/fn.remove_storage_prefix.html
-			remove_storage_prefix(pallet_prefix, storage_item_prefix, &[]);
-
-			// Assert that old storage is empty
-			assert!(storage_key_iter::<
-				T::AuthorId,
-				RegistrationInfo<T::AccountId, BalanceOf<T>>,
-				Twox64Concat,
-			>(pallet_prefix, storage_item_prefix)
-			.next()
-			.is_none());
-
-			// Write the mappings back to storage with the new secure hasher
-			for (author_id, account_id) in stored_data {
-				MappingWithDeposit::<T>::insert(author_id, account_id);
-			}
-
-			// Return the weight used. For each migrated mapping there is a red to get it into
-			// memory, a write to clear the old stored value, and a write to re-store it.
-			let db_weights = T::DbWeight::get();
-			migrated_count.saturating_mul(2 * db_weights.write + db_weights.read)
-		}
-
-		#[cfg(feature = "try-runtime")]
-		fn pre_upgrade() -> Result<(), &'static str> {
-			use frame_support::storage::migration::{storage_iter, storage_key_iter};
-			use frame_support::traits::OnRuntimeUpgradeHelpersExt;
-
-			let pallet_prefix: &[u8] = b"AuthorMapping";
-			let storage_item_prefix: &[u8] = b"MappingWithDeposit";
-
-			// We want to test that:
-			// There are no entries in the new storage beforehand
-			// The same number of mappings exist before and after
-			// As long as there are some mappings stored, one representative key maps to the
-			// same value after the migration.
-			// There are no entries in the old storage afterward
-
-			// Assert new storage is empty
-			// Because the pallet and item prefixes are the same, the old storage is still at this
-			// key. However, the values can't be decoded so the assertion passes.
-			assert!(MappingWithDeposit::<T>::iter().next().is_none());
-
-			// Check number of entries, and set it aside in temp storage
-			let mapping_count = storage_iter::<RegistrationInfo<T::AccountId, BalanceOf<T>>>(
-				pallet_prefix,
-				storage_item_prefix,
-			)
-			.count() as u64;
-			Self::set_temp_storage(mapping_count, "mapping_count");
-
-			// Read an example pair from old storage and set it aside in temp storage
-			if mapping_count > 0 {
-				let example_pair = storage_key_iter::<
-					T::AuthorId,
-					RegistrationInfo<T::AccountId, BalanceOf<T>>,
-					Twox64Concat,
-				>(pallet_prefix, storage_item_prefix)
-				.next()
-				.expect("We already confirmed that there was at least one item stored");
-
-				Self::set_temp_storage(example_pair, "example_pair");
-			}
-
-			Ok(())
-		}
-
-		#[cfg(feature = "try-runtime")]
-		fn post_upgrade() -> Result<(), &'static str> {
-			use frame_support::traits::OnRuntimeUpgradeHelpersExt;
-
-			// Check number of entries matches what was set aside in pre_upgrade
-			let old_mapping_count: u64 = Self::get_temp_storage("mapping_count")
-				.expect("We stored a mapping count; it should be there; qed");
-			let new_mapping_count = MappingWithDeposit::<T>::iter().count() as u64;
-			assert_eq!(old_mapping_count, new_mapping_count);
-
-			// Check that our example pair is still well-mapped after the migration
-			if new_mapping_count > 0 {
-				let (account, original_info): (
-					T::AuthorId,
-					RegistrationInfo<T::AccountId, BalanceOf<T>>,
-				) = Self::get_temp_storage("example_pair").expect("qed");
-				let migrated_info = MappingWithDeposit::<T>::get(account).expect("qed");
-				assert_eq!(original_info, migrated_info);
-			}
-
-			Ok(())
-		}
-	}
 }

pallets/author-mapping/src/migrations.rs

@@ -0,0 +1,136 @@
+// Copyright 2019-2021 PureStake Inc.
+// This file is part of Moonbeam.
+
+// Moonbeam is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Moonbeam is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
+
+use crate::{BalanceOf, Config, MappingWithDeposit, RegistrationInfo};
+use frame_support::{
+	pallet_prelude::PhantomData,
+	storage::migration::{remove_storage_prefix, storage_key_iter},
+	traits::{Get, OnRuntimeUpgrade},
+	weights::Weight,
+	Twox64Concat,
+};
+use sp_std::convert::TryInto;
+//TODO sometimes this is unused, sometimes its necessary
+use sp_std::vec::Vec;
+
+/// Migrates the AuthorMapping's storage map fro mthe insecure Twox64 hasher to the secure
+/// BlakeTwo hasher.
+pub struct TwoXToBlake<T>(PhantomData<T>);
+impl<T: Config> OnRuntimeUpgrade for TwoXToBlake<T> {
+	fn on_runtime_upgrade() -> Weight {
+		let pallet_prefix: &[u8] = b"AuthorMapping";
+		let storage_item_prefix: &[u8] = b"MappingWithDeposit";
+
+		// Read all the data into memory.
+		// https://crates.parity.io/frame_support/storage/migration/fn.storage_key_iter.html
+		let stored_data: Vec<_> = storage_key_iter::<
+			T::AuthorId,
+			RegistrationInfo<T::AccountId, BalanceOf<T>>,
+			Twox64Concat,
+		>(pallet_prefix, storage_item_prefix)
+		.collect();
+
+		let migrated_count: Weight = stored_data
+			.len()
+			.try_into()
+			.expect("There are between 0 and 2**64 mappings stored.");
+
+		// Now remove the old storage
+		// https://crates.parity.io/frame_support/storage/migration/fn.remove_storage_prefix.html
+		remove_storage_prefix(pallet_prefix, storage_item_prefix, &[]);
+
+		// Assert that old storage is empty
+		assert!(storage_key_iter::<
+			T::AuthorId,
+			RegistrationInfo<T::AccountId, BalanceOf<T>>,
+			Twox64Concat,
+		>(pallet_prefix, storage_item_prefix)
+		.next()
+		.is_none());
+
+		// Write the mappings back to storage with the new secure hasher
+		for (author_id, account_id) in stored_data {
+			MappingWithDeposit::<T>::insert(author_id, account_id);
+		}
+
+		// Return the weight used. For each migrated mapping there is a red to get it into
+		// memory, a write to clear the old stored value, and a write to re-store it.
+		let db_weights = T::DbWeight::get();
+		migrated_count.saturating_mul(2 * db_weights.write + db_weights.read)
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<(), &'static str> {
+		let pallet_prefix: &[u8] = b"AuthorMapping";
+		let storage_item_prefix: &[u8] = b"MappingWithDeposit";
+
+		// We want to test that:
+		// There are no entries in the new storage beforehand
+		// The same number of mappings exist before and after
+		// As long as there are some mappings stored, one representative key maps to the
+		// same value after the migration.
+		// There are no entries in the old storage afterward
+
+		// Assert new storage is empty
+		// Because the pallet and item prefixes are the same, the old storage is still at this
+		// key. However, the values can't be decoded so the assertion passes.
+		assert!(MappingWithDeposit::<T>::iter().next().is_none());
+
+		// Check number of entries, and set it aside in temp storage
+		let mapping_count = storage_iter::<RegistrationInfo<T::AccountId, BalanceOf<T>>>(
+			pallet_prefix,
+			storage_item_prefix,
+		)
+		.count() as u64;
+		Self::set_temp_storage(mapping_count, "mapping_count");
+
+		// Read an example pair from old storage and set it aside in temp storage
+		if mapping_count > 0 {
+			let example_pair = storage_key_iter::<
+				T::AuthorId,
+				RegistrationInfo<T::AccountId, BalanceOf<T>>,
+				Twox64Concat,
+			>(pallet_prefix, storage_item_prefix)
+			.next()
+			.expect("We already confirmed that there was at least one item stored");
+
+			Self::set_temp_storage(example_pair, "example_pair");
+		}
+
+		Ok(())
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade() -> Result<(), &'static str> {
+		// Check number of entries matches what was set aside in pre_upgrade
+		let old_mapping_count: u64 = Self::get_temp_storage("mapping_count")
+			.expect("We stored a mapping count; it should be there; qed");
+		let new_mapping_count = MappingWithDeposit::<T>::iter().count() as u64;
+		assert_eq!(old_mapping_count, new_mapping_count);
+
+		// Check that our example pair is still well-mapped after the migration
+		if new_mapping_count > 0 {
+			let (account, original_info): (
+				T::AuthorId,
+				RegistrationInfo<T::AccountId, BalanceOf<T>>,
+			) = Self::get_temp_storage("example_pair").expect("qed");
+			let migrated_info = MappingWithDeposit::<T>::get(account).expect("qed");
+			assert_eq!(original_info, migrated_info);
+		}
+
+		Ok(())
+	}
+}

pallets/migrations/Cargo.toml

@@ -26,3 +26,6 @@ std = [
 	"sp-std/std",
 	"sp-runtime/std",
 ]
+try-runtime = [
+	"frame-support/try-runtime",
+]

pallets/migrations/src/lib.rs

@@ -44,6 +44,18 @@ pub trait Migration {
 	/// constraints will lead to a bricked chain upon a runtime upgrade because the parachain will
 	/// not be able to produce a block that the relay chain will accept.
 	fn migrate(&self, available_weight: Weight) -> Weight;
+
+	/// Run a standard pre-runtime test. This works the same way as in a normal runtime upgrade.
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade(&self) -> Result<(), &'static str> {
+		Ok(())
+	}
+
+	/// Run a standard post-runtime test. This works the same way as in a normal runtime upgrade.
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(&self) -> Result<(), &'static str> {
+		Ok(())
+	}
 }
 
 #[pallet]
@@ -103,6 +115,88 @@ pub mod pallet {
 
 			weight
 		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<(), &'static str> {
+			use frame_support::traits::OnRuntimeUpgradeHelpersExt;
+
+			let mut failed = false;
+			for migration in &T::MigrationsList::get() {
+				let migration_name = migration.friendly_name();
+				let migration_name_as_bytes = migration_name.as_bytes();
+
+				let migration_done = <MigrationState<T>>::get(migration_name_as_bytes);
+				if migration_done {
+					continue;
+				}
+				log::trace!("invoking pre_upgrade() on migration {}", migration_name);
+
+				// dump the migration name to temp storage so post_upgrade will know which
+				// migrations were performed (as opposed to skipped)
+				Self::set_temp_storage(true, migration_name);
+
+				let result = migration.pre_upgrade();
+				match result {
+					Ok(()) => {
+						log::info!("migration {} pre_upgrade() => Ok()", migration_name);
+					}
+					Err(msg) => {
+						log::error!("migration {} pre_upgrade() => Err({})", migration_name, msg);
+						failed = true;
+					}
+				}
+			}
+
+			if failed {
+				Err("One or more pre_upgrade tests failed; see output above.")
+			} else {
+				Ok(())
+			}
+		}
+
+		/// Run a standard post-runtime test. This works the same way as in a normal runtime upgrade.
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade() -> Result<(), &'static str> {
+			use frame_support::traits::OnRuntimeUpgradeHelpersExt;
+
+			// TODO: my desire to DRY all the things feels like this code is very repetitive...
+
+			let mut failed = false;
+			for migration in &T::MigrationsList::get() {
+				let migration_name = migration.friendly_name();
+
+				// we can't query MigrationState because on_runtime_upgrade() would have
+				// unconditionally set it to true, so we read a hint from temp storage which was
+				// left for us by pre_upgrade()
+				match Self::get_temp_storage::<bool>(migration_name) {
+					Some(value) => assert!(true == value, "our dummy value might as well be true"),
+					None => continue,
+				}
+
+				log::trace!("invoking post_upgrade() on migration {}", migration_name);
+
+				let result = migration.post_upgrade();
+				match result {
+					Ok(()) => {
+						log::info!("migration {} post_upgrade() => Ok()", migration_name);
+					}
+					Err(msg) => {
+						log::error!(
+							"migration {} post_upgrade() => Err({})",
+							migration_name,
+							msg
+						);
+						failed = true;
+					}
+				}
+			}
+
+			if failed {
+				Err("One or more post_upgrade tests failed; see output above.")
+			} else {
+				Ok(())
+			}
+		}
 	}
 
 	#[pallet::storage]