dapp-staking mbm

ermalkaleci · ermalkaleci · commit c167ef9e4332 · 2024-09-03T16:01:05.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pallets/dapp-staking/Cargo.toml b/pallets/dapp-staking/Cargo.toml
@@ -29,6 +29,7 @@ frame-benchmarking = { workspace = true, optional = true }
 
 [dev-dependencies]
 pallet-balances = { workspace = true }
+pallet-migrations = { workspace = true }
 
 [features]
 default = ["std"]
diff --git a/pallets/dapp-staking/src/benchmarking/mod.rs b/pallets/dapp-staking/src/benchmarking/mod.rs
@@ -21,7 +21,7 @@ use super::{Pallet as DappStaking, *};
 use astar_primitives::Balance;
 use frame_benchmarking::v2::*;
 
-use frame_support::assert_ok;
+use frame_support::{assert_ok, migrations::SteppedMigration, weights::WeightMeter};
 use frame_system::{Pallet as System, RawOrigin};
 use sp_std::prelude::*;
 
@@ -1070,6 +1070,73 @@ mod benchmarks {
         );
     }
 
+    /// Benchmark a single step of mbm migration.
+    #[benchmark]
+    fn step() {
+        let alice: T::AccountId = account("alice", 0, 1);
+
+        Ledger::<T>::set(
+            &alice,
+            AccountLedger {
+                locked: 1000,
+                unlocking: vec![
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 5,
+                    },
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 20,
+                    },
+                ]
+                .try_into()
+                .unwrap(),
+                staked: Default::default(),
+                staked_future: None,
+                contract_stake_count: 0,
+            },
+        );
+        CurrentEraInfo::<T>::put(EraInfo {
+            total_locked: 1000,
+            unlocking: 200,
+            current_stake_amount: Default::default(),
+            next_stake_amount: Default::default(),
+        });
+
+        System::<T>::set_block_number(10u32.into());
+        let mut meter = WeightMeter::new();
+
+        #[block]
+        {
+            crate::migration::LazyMigration::<T, weights::SubstrateWeight<T>>::step(
+                None, &mut meter,
+            )
+            .unwrap();
+        }
+
+        assert_eq!(
+            Ledger::<T>::get(&alice),
+            AccountLedger {
+                locked: 1000,
+                unlocking: vec![
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 5, // already unlocked
+                    },
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 30, // double remaining blocks
+                    },
+                ]
+                .try_into()
+                .unwrap(),
+                staked: Default::default(),
+                staked_future: None,
+                contract_stake_count: 0,
+            }
+        );
+    }
+
     impl_benchmark_test_suite!(
         Pallet,
         crate::benchmarking::tests::new_test_ext(),
diff --git a/pallets/dapp-staking/src/migration.rs b/pallets/dapp-staking/src/migration.rs
@@ -17,7 +17,12 @@
 // along with Astar. If not, see <http://www.gnu.org/licenses/>.
 
 use super::*;
-use frame_support::{storage_alias, traits::UncheckedOnRuntimeUpgrade};
+use frame_support::{
+    migrations::{MigrationId, SteppedMigration, SteppedMigrationError},
+    storage_alias,
+    traits::UncheckedOnRuntimeUpgrade,
+    weights::WeightMeter,
+};
 
 #[cfg(feature = "try-runtime")]
 use sp_std::vec::Vec;
@@ -412,3 +417,80 @@ pub mod v6 {
         pub period: PeriodNumber,
     }
 }
+
+const PALLET_MIGRATIONS_ID: &[u8; 16] = b"dapp-staking-mbm";
+
+pub struct LazyMigration<T, W: WeightInfo>(PhantomData<(T, W)>);
+
+impl<T: Config, W: WeightInfo> SteppedMigration for LazyMigration<T, W> {
+    type Cursor = <T as frame_system::Config>::AccountId;
+    // Without the explicit length here the construction of the ID would not be infallible.
+    type Identifier = MigrationId<16>;
+
+    /// The identifier of this migration. Which should be globally unique.
+    fn id() -> Self::Identifier {
+        MigrationId {
+            pallet_id: *PALLET_MIGRATIONS_ID,
+            version_from: 0,
+            version_to: 1,
+        }
+    }
+
+    fn step(
+        mut cursor: Option<Self::Cursor>,
+        meter: &mut WeightMeter,
+    ) -> Result<Option<Self::Cursor>, SteppedMigrationError> {
+        let required = W::step();
+        // If there is not enough weight for a single step, return an error. This case can be
+        // problematic if it is the first migration that ran in this block. But there is nothing
+        // that we can do about it here.
+        if meter.remaining().any_lt(required) {
+            return Err(SteppedMigrationError::InsufficientWeight { required });
+        }
+
+        let mut count = 0u32;
+        let current_block_number =
+            frame_system::Pallet::<T>::block_number().saturated_into::<u32>();
+
+        // We loop here to do as much progress as possible per step.
+        loop {
+            if meter.try_consume(required).is_err() {
+                break;
+            }
+
+            let mut iter = if let Some(last_key) = cursor {
+                // If a cursor is provided, start iterating from the stored value
+                // corresponding to the last key processed in the previous step.
+                // Note that this only works if the old and the new map use the same way to hash
+                // storage keys.
+                Ledger::<T>::iter_from(Ledger::<T>::hashed_key_for(last_key))
+            } else {
+                // If no cursor is provided, start iterating from the beginning.
+                Ledger::<T>::iter()
+            };
+
+            // If there's a next item in the iterator, perform the migration.
+            if let Some((ref last_key, mut ledger)) = iter.next() {
+                for chunk in ledger.unlocking.iter_mut() {
+                    let remaining_blocks = chunk.unlock_block.saturating_sub(current_block_number);
+                    chunk.unlock_block.saturating_accrue(remaining_blocks);
+                }
+
+                // Override ledger
+                Ledger::<T>::insert(last_key, ledger);
+
+                // inc counter
+                count.saturating_inc();
+
+                // Return the processed key as the new cursor.
+                cursor = Some(last_key.clone())
+            } else {
+                // Signal that the migration is complete (no more items to process).
+                cursor = None;
+                break;
+            }
+        }
+        log::debug!(target: LOG_TARGET, "migrated {count:?} entries");
+        Ok(cursor)
+    }
+}
diff --git a/pallets/dapp-staking/src/test/migrations.rs b/pallets/dapp-staking/src/test/migrations.rs
@@ -0,0 +1,85 @@
+// This file is part of Astar.
+
+// Copyright (C) Stake Technologies Pte.Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+// Astar 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.
+
+// Astar 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 Astar. If not, see <http://www.gnu.org/licenses/>.
+
+#![cfg(all(test, not(feature = "runtime-benchmarks")))]
+
+use crate::test::mock::*;
+use crate::{AccountLedger, CurrentEraInfo, EraInfo, Ledger, UnlockingChunk};
+use frame_support::traits::OnRuntimeUpgrade;
+
+#[test]
+fn lazy_migrations() {
+    ExtBuilder::default().build_and_execute(|| {
+        Ledger::<Test>::set(
+            &1,
+            AccountLedger {
+                locked: 1000,
+                unlocking: vec![
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 5,
+                    },
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 20,
+                    },
+                ]
+                .try_into()
+                .unwrap(),
+                staked: Default::default(),
+                staked_future: None,
+                contract_stake_count: 0,
+            },
+        );
+        CurrentEraInfo::<Test>::put(EraInfo {
+            total_locked: 1000,
+            unlocking: 200,
+            current_stake_amount: Default::default(),
+            next_stake_amount: Default::default(),
+        });
+
+        // go to block before migration
+        run_to_block(9);
+
+        // onboard MBMs
+        AllPalletsWithSystem::on_runtime_upgrade();
+        run_to_block(10);
+
+        assert_eq!(
+            Ledger::<Test>::get(&1),
+            AccountLedger {
+                locked: 1000,
+                unlocking: vec![
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 5, // already unlocked
+                    },
+                    UnlockingChunk {
+                        amount: 100,
+                        unlock_block: 30, // double remaining blocks
+                    },
+                ]
+                .try_into()
+                .unwrap(),
+                staked: Default::default(),
+                staked_future: None,
+                contract_stake_count: 0,
+            }
+        );
+    })
+}
diff --git a/pallets/dapp-staking/src/test/mock.rs b/pallets/dapp-staking/src/test/mock.rs
@@ -23,7 +23,9 @@ use crate::{
 };
 
 use frame_support::{
-    construct_runtime, ord_parameter_types, parameter_types,
+    construct_runtime, derive_impl,
+    migrations::MultiStepMigrator,
+    ord_parameter_types, parameter_types,
     traits::{fungible::Mutate as FunMutate, ConstBool, ConstU128, ConstU32, EitherOfDiverse},
     weights::Weight,
 };
@@ -54,6 +56,7 @@ construct_runtime!(
         System: frame_system,
         Balances: pallet_balances,
         DappStaking: pallet_dapp_staking,
+        MultiBlockMigrations: pallet_migrations,
     }
 );
 
@@ -89,7 +92,7 @@ impl frame_system::Config for Test {
     type MaxConsumers = frame_support::traits::ConstU32<16>;
     type RuntimeTask = RuntimeTask;
     type SingleBlockMigrations = ();
-    type MultiBlockMigrator = ();
+    type MultiBlockMigrator = MultiBlockMigrations;
     type PreInherents = ();
     type PostInherents = ();
     type PostTransactions = ();
@@ -111,6 +114,21 @@ impl pallet_balances::Config for Test {
     type WeightInfo = ();
 }
 
+parameter_types! {
+    pub const MaxServiceWeight: Weight = Weight::from_parts(1_000_000_000, 1_000_000);
+}
+
+#[derive_impl(pallet_migrations::config_preludes::TestDefaultConfig)]
+impl pallet_migrations::Config for Test {
+    #[cfg(not(feature = "runtime-benchmarks"))]
+    type Migrations =
+        (crate::migration::LazyMigration<Test, crate::weights::SubstrateWeight<Test>>,);
+    #[cfg(feature = "runtime-benchmarks")]
+    type Migrations = pallet_migrations::mock_helpers::MockedMigrations;
+    type MigrationStatusHandler = ();
+    type MaxServiceWeight = MaxServiceWeight;
+}
+
 pub struct DummyPriceProvider;
 impl PriceProvider for DummyPriceProvider {
     fn average_price() -> FixedU128 {
@@ -397,6 +415,9 @@ pub(crate) fn run_to_block(n: BlockNumber) {
         System::set_block_number(System::block_number() + 1);
         // This is performed outside of dapps staking but we expect it before on_initialize
 
+        // Done by Executive:
+        <Test as frame_system::Config>::MultiBlockMigrator::step();
+
         let pre_snapshot = MemorySnapshot::new();
         DappStaking::on_initialize(System::block_number());
         assert_block_bump(&pre_snapshot);
diff --git a/pallets/dapp-staking/src/test/mod.rs b/pallets/dapp-staking/src/test/mod.rs
@@ -16,6 +16,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Astar. If not, see <http://www.gnu.org/licenses/>.
 
+mod migrations;
 pub(crate) mod mock;
 mod testing_utils;
 mod tests;
diff --git a/pallets/dapp-staking/src/weights.rs b/pallets/dapp-staking/src/weights.rs
@@ -73,6 +73,7 @@ pub trait WeightInfo {
 	fn on_initialize_build_and_earn_to_build_and_earn() -> Weight;
 	fn dapp_tier_assignment(x: u32, ) -> Weight;
 	fn on_idle_cleanup() -> Weight;
+	fn step() -> Weight;
 }
 
 /// Weights for pallet_dapp_staking using the Substrate node and recommended hardware.
@@ -477,6 +478,10 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().reads(2_u64))
 			.saturating_add(T::DbWeight::get().writes(3_u64))
 	}
+
+	fn step() -> Weight {
+		Weight::default()
+	}
 }
 
 // For backwards compatibility and tests
@@ -880,4 +885,8 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().reads(2_u64))
 			.saturating_add(RocksDbWeight::get().writes(3_u64))
 	}
+
+	fn step() -> Weight {
+		Weight::default()
+	}
 }
diff --git a/pallets/vesting-mbm/src/lib.rs b/pallets/vesting-mbm/src/lib.rs
@@ -42,10 +42,7 @@ const PALLET_MIGRATIONS_ID: &[u8; 18] = b"pallet-vesting-mbm";
 
 pub struct LazyMigration<T, W: weights::WeightInfo>(core::marker::PhantomData<(T, W)>);
 
-impl<T, W: weights::WeightInfo> SteppedMigration for LazyMigration<T, W>
-where
-    T: frame_system::Config + pallet_vesting::Config,
-{
+impl<T: pallet_vesting::Config, W: weights::WeightInfo> SteppedMigration for LazyMigration<T, W> {
     type Cursor = <T as frame_system::Config>::AccountId;
     // Without the explicit length here the construction of the ID would not be infallible.
     type Identifier = MigrationId<18>;
diff --git a/runtime/shibuya/src/lib.rs b/runtime/shibuya/src/lib.rs

Original file line number	Diff line number	Diff line change
`@@ -73,6 +73,7 @@ pub trait WeightInfo {`
`73`	`73`	`fn on_initialize_build_and_earn_to_build_and_earn() -> Weight;`
`74`	`74`	`fn dapp_tier_assignment(x: u32, ) -> Weight;`
`75`	`75`	`fn on_idle_cleanup() -> Weight;`
	`76`	`+ fn step() -> Weight;`
`76`	`77`	`}`
`77`	`78`
`78`	`79`	`/// Weights for pallet_dapp_staking using the Substrate node and recommended hardware.`
`@@ -477,6 +478,10 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {`
`477`	`478`	`.saturating_add(T::DbWeight::get().reads(2_u64))`
`478`	`479`	`.saturating_add(T::DbWeight::get().writes(3_u64))`
`479`	`480`	`}`
	`481`	`+`
	`482`	`+ fn step() -> Weight {`
	`483`	`+ Weight::default()`
	`484`	`+ }`
`480`	`485`	`}`
`481`	`486`
`482`	`487`	`// For backwards compatibility and tests`
`@@ -880,4 +885,8 @@ impl WeightInfo for () {`
`880`	`885`	`.saturating_add(RocksDbWeight::get().reads(2_u64))`
`881`	`886`	`.saturating_add(RocksDbWeight::get().writes(3_u64))`
`882`	`887`	`}`
	`888`	`+`
	`889`	`+ fn step() -> Weight {`
	`890`	`+ Weight::default()`
	`891`	`+ }`
`883`	`892`	`}`