[Staking] Kill unnecessary storage bloat (#970)

Co-authored-by: Joshy Orndorff <JoshOrndorff@users.noreply.github.com>

pallets/parachain-staking/src/lib.rs

@@ -49,12 +49,12 @@
 #[cfg(any(test, feature = "runtime-benchmarks"))]
 mod benchmarks;
 mod inflation;
+pub mod migrations;
 #[cfg(test)]
 mod mock;
 mod set;
 #[cfg(test)]
 mod tests;
-
 pub mod weights;
 use weights::WeightInfo;
 
@@ -918,7 +918,7 @@ pub mod pallet {
 	#[pallet::storage]
 	#[pallet::getter(fn round)]
 	/// Current round index and next round scheduled transition
-	type Round<T: Config> = StorageValue<_, RoundInfo<T::BlockNumber>, ValueQuery>;
+	pub(crate) type Round<T: Config> = StorageValue<_, RoundInfo<T::BlockNumber>, ValueQuery>;
 
 	#[pallet::storage]
 	#[pallet::getter(fn nominator_state2)]
@@ -1706,18 +1706,18 @@ pub mod pallet {
 			));
 			Ok(().into())
 		}
-		fn pay_stakers(next: RoundIndex) {
-			// payout is next - duration rounds ago => next - duration > 0 else return early
+		fn pay_stakers(now: RoundIndex) {
+			// payout is now - duration rounds ago => now - duration > 0 else return early
 			let duration = T::RewardPaymentDelay::get();
-			if next <= duration {
+			if now <= duration {
 				return;
 			}
-			let round_to_payout = next - duration;
-			let total = <Points<T>>::get(round_to_payout);
+			let round_to_payout = now - duration;
+			let total = <Points<T>>::take(round_to_payout);
 			if total.is_zero() {
 				return;
 			}
-			let total_staked = <Staked<T>>::get(round_to_payout);
+			let total_staked = <Staked<T>>::take(round_to_payout);
 			let total_issuance = Self::compute_issuance(total_staked);
 			let mut left_issuance = total_issuance;
 			// reserve portion of issuance for parachain bond account
@@ -1932,7 +1932,7 @@ pub mod pallet {
 		}
 		/// Best as in most cumulatively supported in terms of stake
 		/// Returns [collator_count, nomination_count, total staked]
-		fn select_top_candidates(next: RoundIndex) -> (u32, u32, BalanceOf<T>) {
+		fn select_top_candidates(now: RoundIndex) -> (u32, u32, BalanceOf<T>) {
 			let (mut collator_count, mut nomination_count, mut total) =
 				(0u32, 0u32, BalanceOf::<T>::zero());
 			// choose the top TotalSelected qualified candidates, ordered by stake
@@ -1946,8 +1946,8 @@ pub mod pallet {
 				let amount = state.total_counted;
 				total += amount;
 				let exposure: CollatorSnapshot<T::AccountId, BalanceOf<T>> = state.into();
-				<AtStake<T>>::insert(next, account, exposure);
-				Self::deposit_event(Event::CollatorChosen(next, account.clone(), amount));
+				<AtStake<T>>::insert(now, account, exposure);
+				Self::deposit_event(Event::CollatorChosen(now, account.clone(), amount));
 			}
 			// insert canonical collator set
 			<SelectedCandidates<T>>::put(collators);

pallets/parachain-staking/src/migrations.rs

@@ -0,0 +1,74 @@
+// 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/>.
+
+//! # Migrations
+use crate::{Config, Points, Round, Staked};
+use frame_support::{
+	pallet_prelude::PhantomData,
+	traits::{Get, OnRuntimeUpgrade},
+	weights::Weight,
+};
+
+/// Migration to purge staking storage bloat for `Points` and `AtStake` storage items
+pub struct PurgeStaleStorage<T>(PhantomData<T>);
+impl<T: Config> OnRuntimeUpgrade for PurgeStaleStorage<T> {
+	fn on_runtime_upgrade() -> Weight {
+		log::info!(target: "PurgeStaleStorage", "running migration to remove storage bloat");
+		let current_round = <Round<T>>::get().current;
+		let payment_delay = T::RewardPaymentDelay::get();
+		let db_weight = T::DbWeight::get();
+		let (reads, mut writes) = (3u64, 0u64);
+		if current_round <= payment_delay {
+			// early enough so no storage bloat exists yet
+			// (only relevant for chains <= payment_delay rounds old)
+			return db_weight.reads(reads);
+		}
+		// already paid out at the beginning of current round
+		let most_recent_round_to_kill = current_round - payment_delay;
+		for i in 1..=most_recent_round_to_kill {
+			writes += 2u64;
+			<Staked<T>>::remove(i);
+			<Points<T>>::remove(i);
+		}
+		// 5% of the max block weight as safety margin for computation
+		db_weight.reads(reads) + db_weight.writes(writes) + 25_000_000_000
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<(), &'static str> {
+		// trivial migration
+		Ok(())
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade() -> Result<(), &'static str> {
+		// expect only the storage items for the last 2 rounds to be stored
+		let staked_count = Staked::<T>::iter().count() as u64;
+		let points_count = Points::<T>::iter().count() as u64;
+		let delay = T::RewardPaymentDelay::get();
+		assert_eq!(
+			staked_count, delay,
+			"Expected {} for `Staked` count, Found: {}",
+			delay, staked_count
+		);
+		assert_eq!(
+			points_count, delay,
+			"Expected {} for `Points` count, Found: {}",
+			delay, staked_count
+		);
+		Ok(())
+	}
+}

pallets/parachain-staking/src/tests.rs

@@ -3948,3 +3948,36 @@ fn nomination_events_convey_correct_position() {
 			);
 		});
 }
+
+// Migration Unit Test
+#[test]
+fn verify_purge_storage_migration_works() {
+	use crate::{Points, Round, RoundInfo, Staked};
+	use frame_support::traits::OnRuntimeUpgrade;
+	ExtBuilder::default().build().execute_with(|| {
+		// mutate storage similar to if 10 rounds had passed
+		for i in 1..=10 {
+			<Staked<Test>>::insert(i, 100);
+			<Points<Test>>::insert(i, 100);
+		}
+		// set the round information to the 10th round
+		// (we do not use roll_to because the payment logic uses `take` in the code)
+		<Round<Test>>::put(RoundInfo {
+			current: 10,
+			first: 45,
+			length: 5,
+		});
+		// execute the migration
+		crate::migrations::PurgeStaleStorage::<Test>::on_runtime_upgrade();
+		// verify that all inserted items are removed except last 2 rounds
+		for i in 1..=8 {
+			assert_eq!(<Staked<Test>>::get(i), 0);
+			assert_eq!(<Points<Test>>::get(i), 0);
+		}
+		// last 2 rounds are still stored (necessary for future payouts)
+		for i in 9..=10 {
+			assert_eq!(<Staked<Test>>::get(i), 100);
+			assert_eq!(<Points<Test>>::get(i), 100);
+		}
+	});
+}

runtime/common/Cargo.toml

@@ -13,6 +13,7 @@ sp-runtime = { git = "https://github.com/purestake/substrate", branch = "moonbea
 sp-std = { git = "https://github.com/purestake/substrate", branch = "moonbeam-polkadot-v0.9.11", default-features = false }
 frame-support = { git = "https://github.com/purestake/substrate", default-features = false, branch = "moonbeam-polkadot-v0.9.11" }
 pallet-author-mapping = { path = "../../pallets/author-mapping", default-features = false }
+parachain-staking = { path = "../../pallets/parachain-staking", default-features = false }
 pallet-collective = { git = "https://github.com/purestake/substrate", default-features = false, branch = "moonbeam-polkadot-v0.9.11" }
 frame-system = { git = "https://github.com/purestake/substrate", default-features = false, branch = "moonbeam-polkadot-v0.9.11" }
 log = "0.4"
@@ -24,4 +25,5 @@ std = [
 	"sp-std/std",
 	"frame-support/std",
 	"pallet-author-mapping/std",
+	"parachain-staking/std",
 ]

runtime/common/src/migrations.rs

@@ -24,11 +24,36 @@ use frame_support::{
 };
 use pallet_author_mapping::{migrations::TwoXToBlake, Config as AuthorMappingConfig};
 use pallet_migrations::Migration;
+use parachain_staking::{migrations::PurgeStaleStorage, Config as ParachainStakingConfig};
 use sp_std::{marker::PhantomData, prelude::*};
 
 /// This module acts as a registry where each migration is defined. Each migration should implement
 /// the "Migration" trait declared in the pallet-migrations crate.
 
+/// A moonbeam migration wrapping the similarly named migration in parachain-staking
+pub struct ParachainStakingPurgeStaleStorage<T>(PhantomData<T>);
+impl<T: ParachainStakingConfig> Migration for ParachainStakingPurgeStaleStorage<T> {
+	fn friendly_name(&self) -> &str {
+		"MM_Parachain_Staking_PurgeStaleStorage"
+	}
+
+	fn migrate(&self, _available_weight: Weight) -> Weight {
+		PurgeStaleStorage::<T>::on_runtime_upgrade()
+	}
+
+	/// 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> {
+		PurgeStaleStorage::<T>::pre_upgrade()
+	}
+
+	/// 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> {
+		PurgeStaleStorage::<T>::post_upgrade()
+	}
+}
+
 /// A moonbeam migration wrapping the similarly named migration in pallet-author-mapping
 pub struct AuthorMappingTwoXToBlake<T>(PhantomData<T>);
 impl<T: AuthorMappingConfig> Migration for AuthorMappingTwoXToBlake<T> {
@@ -94,7 +119,7 @@ pub struct CommonMigrations<Runtime, Council, Tech>(PhantomData<(Runtime, Counci
 impl<Runtime, Council, Tech> Get<Vec<Box<dyn Migration>>>
 	for CommonMigrations<Runtime, Council, Tech>
 where
-	Runtime: pallet_author_mapping::Config,
+	Runtime: pallet_author_mapping::Config + parachain_staking::Config,
 	Council: GetStorageVersion + PalletInfoAccess + 'static,
 	Tech: GetStorageVersion + PalletInfoAccess + 'static,
 {
@@ -106,6 +131,9 @@ where
 		// let migration_collectives =
 		//	MigrateCollectivePallets::<Runtime, Council, Tech>(Default::default());
 
+		let migration_parachain_staking_purge_stale_storage =
+			ParachainStakingPurgeStaleStorage::<Runtime>(Default::default());
+
 		// TODO: this is a lot of allocation to do upon every get() call. this *should* be avoided
 		// except when pallet_migrations undergoes a runtime upgrade -- but TODO: review
 
@@ -114,6 +142,7 @@ where
 			// Box::new(migration_author_mapping_twox_to_blake),
 			// completed in runtime 900
 			// Box::new(migration_collectives),
+			Box::new(migration_parachain_staking_purge_stale_storage),
 		]
 	}
 }