feat: staking v5 migrations (#249)

* wip: initialize v5 migrations

* fix: try-runtime

run even if everything is up to date

* wip: staking migrations

* tests: add tests

* fix tests

* chore: 🧹

* add benchmarks

* fix weights

Author: weichweich

pallets/delegation/src/migrations.rs

@@ -58,7 +58,7 @@ impl<T: Config> VersionMigratorTrait<T> for DelegationStorageVersion {
 	fn pre_migrate(&self) -> Result<(), &str> {
 		match *self {
 			Self::V1 => v1::pre_migrate::<T>(),
-			Self::V2 => Err("Already latest v2 version."),
+			Self::V2 => Ok(()),
 		}
 	}
 
@@ -76,7 +76,7 @@ impl<T: Config> VersionMigratorTrait<T> for DelegationStorageVersion {
 	fn post_migrate(&self) -> Result<(), &str> {
 		match *self {
 			Self::V1 => v1::post_migrate::<T>(),
-			Self::V2 => Err("Migration from v2 should have never happened in the first place."),
+			Self::V2 => Ok(()),
 		}
 	}
 }
@@ -104,11 +104,6 @@ impl<T: Config> DelegationStorageMigrator<T> {
 	/// latest possible.
 	#[cfg(feature = "try-runtime")]
 	pub(crate) fn pre_migrate() -> Result<(), &'static str> {
-		ensure!(
-			StorageVersion::<T>::get() < DelegationStorageVersion::latest(),
-			"Already the latest storage version."
-		);
-
 		// Don't need to check for any other pre_migrate, as in try-runtime it is also
 		// called in the migrate() function. Same applies for post_migrate checks for
 		// each version migrator.

pallets/did/src/migrations.rs

@@ -58,7 +58,7 @@ impl<T: Config> VersionMigratorTrait<T> for DidStorageVersion {
 	fn pre_migrate(&self) -> Result<(), &str> {
 		match *self {
 			Self::V1 => v1::pre_migrate::<T>(),
-			Self::V2 => Err("Already latest v2 version."),
+			Self::V2 => Ok(()),
 		}
 	}
 
@@ -76,7 +76,7 @@ impl<T: Config> VersionMigratorTrait<T> for DidStorageVersion {
 	fn post_migrate(&self) -> Result<(), &str> {
 		match *self {
 			Self::V1 => v1::post_migrate::<T>(),
-			Self::V2 => Err("Migration from v2 should have never happened in the first place."),
+			Self::V2 => Ok(()),
 		}
 	}
 }
@@ -104,11 +104,6 @@ impl<T: Config> DidStorageMigrator<T> {
 	/// latest possible.
 	#[cfg(feature = "try-runtime")]
 	pub(crate) fn pre_migrate() -> Result<(), &'static str> {
-		ensure!(
-			StorageVersion::<T>::get() < DidStorageVersion::latest(),
-			"Already the latest storage version."
-		);
-
 		// Don't need to check for any other pre_migrate, as in try-runtime it is also
 		// called in the migrate() function. Same applies for post_migrate checks for
 		// each version migrator.

pallets/parachain-staking/src/default_weights.rs

@@ -23,7 +23,7 @@
 //! EXECUTION: Some(Wasm), WASM-EXECUTION: Compiled, CHAIN: Some("dev"), DB CACHE: 128
 
 // Executed Command:
-// /home/willi/mashnet-node/target/release/kilt-parachain
+// target/release/kilt-parachain
 // benchmark
 // --chain=dev
 // --steps=50
@@ -33,8 +33,8 @@
 // --execution=wasm
 // --wasm-execution=compiled
 // --heap-pages=4096
-// --output=../../pallets/parachain-staking/src/default_weights.rs
-// --template=../../.maintain/weight-template.hbs
+// --output=pallets/parachain-staking/src/default_weights.rs
+// --template=.maintain/weight-template.hbs
 
 
 #![cfg_attr(rustfmt, rustfmt_skip)]

pallets/parachain-staking/src/lib.rs

@@ -593,6 +593,7 @@ pub mod pallet {
 	/// It maps from an account to its information.
 	#[pallet::storage]
 	#[pallet::getter(fn candidate_pool)]
+	#[pallet::storage_prefix = "CollatorState"]
 	pub(crate) type CandidatePool<T: Config> = StorageMap<
 		_,
 		Twox64Concat,
@@ -613,6 +614,7 @@ pub mod pallet {
 	/// non collating candidates is not included in [TotalCollatorStake].
 	#[pallet::storage]
 	#[pallet::getter(fn total_collator_stake)]
+	#[pallet::storage_prefix = "Total"]
 	pub(crate) type TotalCollatorStake<T: Config> = StorageValue<_, TotalStake<BalanceOf<T>>, ValueQuery>;
 
 	/// The collator candidates with the highest amount of stake.
@@ -626,6 +628,7 @@ pub mod pallet {
 	/// that a collator can drop out of the collator set by reducing his stake.
 	#[pallet::storage]
 	#[pallet::getter(fn top_candidates)]
+	#[pallet::storage_prefix = "CollatorPool"]
 	pub(crate) type TopCandidates<T: Config> =
 		StorageValue<_, OrderedSet<Stake<T::AccountId, BalanceOf<T>>, T::MaxTopCandidates>, ValueQuery>;
 

pallets/parachain-staking/src/migrations.rs

@@ -29,6 +29,7 @@ use crate::*;
 mod v2;
 mod v3;
 mod v4;
+mod v5;
 
 /// A trait that allows version migrators to access the underlying pallet's
 /// context, e.g., its Config trait.
@@ -53,13 +54,14 @@ pub enum StakingStorageVersion {
 	V2_0_0, // New Reward calculation, MaxCollatorCandidateStake
 	V3_0_0, // Update InflationConfig
 	V4,     // Sort TopCandidates and parachain-stakings by amount
+	V5,     // Remove SelectedCandidates, Count Candidates
 }
 
 #[cfg(feature = "try-runtime")]
 impl StakingStorageVersion {
 	/// The latest storage version.
 	fn latest() -> Self {
-		Self::V4
+		Self::V5
 	}
 }
 
@@ -84,7 +86,8 @@ impl<T: Config> VersionMigratorTrait<T> for StakingStorageVersion {
 			Self::V1_0_0 => v2::pre_migrate::<T>(),
 			Self::V2_0_0 => v3::pre_migrate::<T>(),
 			Self::V3_0_0 => v4::pre_migrate::<T>(),
-			Self::V4 => Err("Already on latest version v4."),
+			Self::V4 => v5::pre_migrate::<T>(),
+			Self::V5 => Ok(()),
 		}
 	}
 
@@ -94,7 +97,8 @@ impl<T: Config> VersionMigratorTrait<T> for StakingStorageVersion {
 			Self::V1_0_0 => v2::migrate::<T>(),
 			Self::V2_0_0 => v3::migrate::<T>(),
 			Self::V3_0_0 => v4::migrate::<T>(),
-			Self::V4 => Weight::zero(),
+			Self::V4 => v5::migrate::<T>(),
+			Self::V5 => Weight::zero(),
 		}
 	}
 
@@ -106,7 +110,8 @@ impl<T: Config> VersionMigratorTrait<T> for StakingStorageVersion {
 			Self::V1_0_0 => v2::post_migrate::<T>(),
 			Self::V2_0_0 => v3::post_migrate::<T>(),
 			Self::V3_0_0 => v4::post_migrate::<T>(),
-			Self::V4 => Err("Migration from v4 should have never happened in the first place."),
+			Self::V4 => v5::post_migrate::<T>(),
+			Self::V5 => Ok(()),
 		}
 	}
 }
@@ -126,20 +131,16 @@ impl<T: Config> StakingStorageMigrator<T> {
 			StakingStorageVersion::V1_0_0 => Some(StakingStorageVersion::V2_0_0),
 			StakingStorageVersion::V2_0_0 => Some(StakingStorageVersion::V3_0_0),
 			// Migration happens naturally, no need to point to the latest version
-			StakingStorageVersion::V3_0_0 => None,
-			StakingStorageVersion::V4 => None,
+			StakingStorageVersion::V3_0_0 => Some(StakingStorageVersion::V4),
+			StakingStorageVersion::V4 => Some(StakingStorageVersion::V5),
+			StakingStorageVersion::V5 => None,
 		}
 	}
 
 	/// Checks whether the latest storage version deployed is lower than the
 	/// latest possible.
 	#[cfg(feature = "try-runtime")]
 	pub(crate) fn pre_migrate() -> Result<(), &'static str> {
-		ensure!(
-			StorageVersion::<T>::get() < StakingStorageVersion::latest(),
-			"Already the latest storage version."
-		);
-
 		// Don't need to check for any other pre_migrate, as in try-runtime it is also
 		// called in the migrate() function. Same applies for post_migrate checks for
 		// each version migrator.
@@ -198,35 +199,11 @@ mod tests {
 
 	use crate::mock::Test as TestRuntime;
 
-	#[test]
-	fn ok_from_v1_migration() {
-		let mut ext = mock::ExtBuilder::default()
-			.with_balances(vec![(1, 100), (2, 100)])
-			.with_collators(vec![(1, 100), (2, 100)])
-			.with_storage_version(StakingStorageVersion::V1_0_0)
-			.build();
-		ext.execute_with(|| {
-			#[cfg(feature = "try-runtime")]
-			assert!(
-				StakingStorageMigrator::<TestRuntime>::pre_migrate().is_ok(),
-				"Storage pre-migrate from v1 should not fail."
-			);
-
-			StakingStorageMigrator::<TestRuntime>::migrate();
-
-			#[cfg(feature = "try-runtime")]
-			assert!(
-				StakingStorageMigrator::<TestRuntime>::post_migrate().is_ok(),
-				"Storage post-migrate from v1 should not fail."
-			);
-		});
-	}
-
 	#[test]
 	fn ok_from_default_migration() {
 		let mut ext = mock::ExtBuilder::default()
-			.with_balances(vec![(1, 100), (2, 100)])
-			.with_collators(vec![(1, 100), (2, 100)])
+			.with_balances((0..15).into_iter().map(|n| (n, 120)).collect())
+			.with_collators((0..15).into_iter().map(|n| (n, 100)).collect())
 			.build();
 		ext.execute_with(|| {
 			#[cfg(feature = "try-runtime")]

pallets/parachain-staking/src/migrations/v5.rs

@@ -0,0 +1,66 @@
+// KILT Blockchain – https://botlabs.org
+// Copyright (C) 2019-2021 BOTLabs GmbH
+
+// The KILT Blockchain 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.
+
+// The KILT Blockchain 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 this program.  If not, see <https://www.gnu.org/licenses/>.
+
+// If you feel like getting in touch with us, you can do so at info@botlabs.org
+
+use crate::{migrations::StakingStorageVersion, CandidateCount, CandidatePool, Config, StorageVersion};
+use frame_support::{dispatch::Weight, storage::StorageValue, traits::Get};
+
+#[cfg(feature = "try-runtime")]
+pub(crate) fn pre_migrate<T: Config>() -> Result<(), &'static str> {
+	assert_eq!(StorageVersion::<T>::get(), StakingStorageVersion::V4);
+	Ok(())
+}
+
+pub(crate) fn migrate<T: Config>() -> Weight {
+	log::info!("Migrating staking to StakingStorageVersion::V5");
+
+	// Kill selected candidates list
+	old::SelectedCandidates::<T>::kill();
+
+	// count candidates
+	let counter: u32 = CandidatePool::<T>::iter().fold(0, |acc, _| acc.saturating_add(1));
+	CandidateCount::<T>::put(counter);
+
+	// update storage version
+	StorageVersion::<T>::put(StakingStorageVersion::V5);
+	log::info!("Completed staking migration to StakingStorageVersion::V5");
+
+	T::DbWeight::get().reads_writes(counter.saturating_add(2).into(), 3)
+}
+
+#[cfg(feature = "try-runtime")]
+pub(crate) fn post_migrate<T: Config>() -> Result<(), &'static str> {
+	assert_eq!(StorageVersion::<T>::get(), StakingStorageVersion::V5);
+	assert!(CandidateCount::<T>::get() > T::MinCollators::get());
+	Ok(())
+}
+
+pub(crate) mod old {
+	use super::*;
+	use frame_support::{decl_module, decl_storage};
+	use sp_std::prelude::*;
+
+	decl_module! {
+		pub struct OldPallet<T: Config> for enum Call where origin: T::Origin {}
+	}
+
+	decl_storage! {
+		trait Store for OldPallet<T: Config> as ParachainStaking {
+			pub(crate) SelectedCandidates: Vec<T::AccountId>;
+		}
+	}
+}

pallets/parachain-staking/src/mock.rs

@@ -273,11 +273,6 @@ impl ExtBuilder {
 		self
 	}
 
-	pub(crate) fn with_storage_version(mut self, storage_version: StakingStorageVersion) -> Self {
-		self.storage_version = storage_version;
-		self
-	}
-
 	pub(crate) fn build(self) -> sp_io::TestExternalities {
 		let mut t = frame_system::GenesisConfig::default()
 			.build_storage::<Test>()

runtimes/peregrine/src/lib.rs

@@ -104,7 +104,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	spec_name: create_runtime_str!("mashnet-node"),
 	impl_name: create_runtime_str!("mashnet-node"),
 	authoring_version: 4,
-	spec_version: 22,
+	spec_version: 23,
 	impl_version: 0,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 2,
@@ -1131,12 +1131,15 @@ impl_runtime_apis! {
 		}
 	}
 
-	// From the Polkadot repo: https://github.com/paritytech/polkadot/blob/master/runtime/polkadot/src/lib.rs#L1371
+	// From the Polkadot repo: https://github.com/paritytech/polkadot/blob/1876963f254f31f8cd2d7b8d5fb26cd38b7836ab/runtime/polkadot/src/lib.rs#L1413
 	#[cfg(feature = "try-runtime")]
 	impl frame_try_runtime::TryRuntime<Block> for Runtime {
 		fn on_runtime_upgrade() -> Result<(Weight, Weight), sp_runtime::RuntimeString> {
 			log::info!("try-runtime::on_runtime_upgrade for peregrine runtime.");
-			let weight = Executive::try_runtime_upgrade()?;
+			let weight = Executive::try_runtime_upgrade().map_err(|err|{
+				log::info!("try-runtime::on_runtime_upgrade failed with: {:?}", err);
+				err
+			})?;
 			Ok((weight, RuntimeBlockWeights::get().max_block))
 		}
 	}

runtimes/peregrine/src/weights/parachain_staking.rs

@@ -23,7 +23,7 @@
 //! EXECUTION: Some(Wasm), WASM-EXECUTION: Compiled, CHAIN: None, DB CACHE: 128
 
 // Executed Command:
-// /home/willi/mashnet-node/target/release/kilt-parachain
+// target/release/kilt-parachain
 // benchmark
 // --execution=wasm
 // --wasm-execution=Compiled
@@ -33,9 +33,9 @@
 // --steps=50
 // --repeat=20
 // --output
-// ../../runtimes/peregrine/src/weights/parachain_staking.rs
+// runtimes/peregrine/src/weights/parachain_staking.rs
 // --template
-// ../../.maintain/runtime-weight-template.hbs
+// .maintain/runtime-weight-template.hbs
 
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -214,4 +214,4 @@ impl<T: frame_system::Config> parachain_staking::WeightInfo for WeightInfo<T> {
 		(17_765_000_u64)
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-}
+}

runtimes/spiritnet/src/lib.rs

@@ -87,7 +87,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	spec_name: create_runtime_str!("kilt-spiritnet"),
 	impl_name: create_runtime_str!("kilt-spiritnet"),
 	authoring_version: 1,
-	spec_version: 22,
+	spec_version: 23,
 	impl_version: 0,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 1,