Add migrate_fee_rate migration.

Author: shamil-gadelshin

pallets/swap/src/lib.rs

@@ -12,6 +12,7 @@ pub use pallet::*;
 #[cfg(feature = "runtime-benchmarks")]
 pub mod benchmarking;
 
+mod migrations;
 #[cfg(test)]
 pub(crate) mod mock;
 

pallets/swap/src/migrations/fee_rate_migration.rs

@@ -0,0 +1,49 @@
+use crate::HasMigrationRun;
+use crate::{Config, pallet};
+use frame_support::traits::Get;
+use frame_support::weights::Weight;
+use scale_info::prelude::string::String;
+use sp_std::collections::btree_map::BTreeMap;
+
+pub const MAX_FEE_RATE: u16 = 1310; // 2 %
+
+pub fn migrate_fee_rate<T: Config>() -> Weight {
+    let migration_name = b"migrate_fee_rate".to_vec();
+
+    // Initialize the weight with one read operation.
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Check if the migration has already run
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    let keypairs = pallet::FeeRate::<T>::iter().collect::<BTreeMap<_, _>>();
+    weight = weight.saturating_add(T::DbWeight::get().reads(keypairs.len() as u64));
+
+    for (netuid, rate) in keypairs {
+        if rate > MAX_FEE_RATE {
+            pallet::FeeRate::<T>::mutate(netuid, |rate| *rate = MAX_FEE_RATE);
+            weight = weight.saturating_add(T::DbWeight::get().writes(1));
+        }
+    }
+
+    // Mark the migration as completed
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    weight
+}

pallets/swap/src/migrations/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod fee_rate_migration;

pallets/swap/src/mock.rs

@@ -4,6 +4,7 @@ use core::num::NonZeroU64;
 
 use frame_support::construct_runtime;
 use frame_support::pallet_prelude::*;
+use frame_support::weights::constants::RocksDbWeight;
 use frame_support::{
     PalletId, parameter_types,
     traits::{ConstU32, Everything},
@@ -50,7 +51,7 @@ impl system::Config for Test {
     type BaseCallFilter = Everything;
     type BlockWeights = ();
     type BlockLength = ();
-    type DbWeight = ();
+    type DbWeight = RocksDbWeight;
     type RuntimeOrigin = RuntimeOrigin;
     type RuntimeCall = RuntimeCall;
     type Hash = H256;

pallets/swap/src/pallet/hooks.rs

@@ -0,0 +1,16 @@
+use frame_support::pallet_macros::pallet_section;
+
+/// A [`pallet_section`] that defines the events for a pallet.
+/// This can later be imported into the pallet using [`import_section`].
+#[pallet_section]
+mod hooks {
+    // ================
+    // ==== Hooks =====
+    // ================
+    #[pallet::hooks]
+    impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+        fn on_runtime_upgrade() -> frame_support::weights::Weight {
+            migrations::fee_rate_migration::migrate_fee_rate::<T>()
+        }
+    }
+}

pallets/swap/src/pallet/mod.rs

@@ -16,17 +16,22 @@ use crate::{
 
 pub use pallet::*;
 
+mod hooks;
 mod impls;
 mod swap_step;
 #[cfg(test)]
 mod tests;
 
+use crate::migrations;
+
+#[import_section(hooks::hooks)]
 #[allow(clippy::module_inception)]
 #[frame_support::pallet]
 #[allow(clippy::expect_used)]
 mod pallet {
     use super::*;
     use frame_system::{ensure_root, ensure_signed};
+    use sp_std::vec::Vec;
 
     #[pallet::pallet]
     pub struct Pallet<T>(_);
@@ -78,6 +83,11 @@ mod pallet {
         33 // ~0.05 %
     }
 
+    /// Storage for migration run status
+    #[pallet::storage]
+    #[pallet::unbounded]
+    pub type HasMigrationRun<T: Config> = StorageMap<_, Identity, Vec<u8>, bool, ValueQuery>;
+
     /// The fee rate applied to swaps per subnet, normalized value between 0 and u16::MAX
     #[pallet::storage]
     pub type FeeRate<T> = StorageMap<_, Twox64Concat, NetUid, u16, ValueQuery, DefaultFeeRate>;

pallets/swap/src/pallet/tests.rs

@@ -14,6 +14,7 @@ use subtensor_runtime_common::NetUid;
 use subtensor_swap_interface::Order as OrderT;
 
 use super::*;
+use crate::migrations::fee_rate_migration::MAX_FEE_RATE;
 use crate::pallet::swap_step::*;
 use crate::{SqrtPrice, mock::*};
 
@@ -2913,3 +2914,45 @@ fn adjust_protocol_liquidity_uses_and_sets_scrap_reservoirs() {
         );
     });
 }
+
+#[test]
+fn test_migrate_fee_rate() {
+    new_test_ext().execute_with(|| {
+        let migration_name = b"migrate_fee_rate".to_vec();
+
+        assert!(
+            !HasMigrationRun::<Test>::get(migration_name.clone()),
+            "HasMigrationRun should be false before migration"
+        );
+
+        let netuid1 = NetUid::from(1);
+        let netuid2 = NetUid::from(2);
+        let value1 = 1000u16;
+        let value2 = 2000u16;
+
+        FeeRate::<Test>::insert(netuid1, value1);
+        FeeRate::<Test>::insert(netuid2, value2);
+
+        // run migration
+        let weight = migrations::fee_rate_migration::migrate_fee_rate::<Test>();
+        assert!(!weight.is_zero(), "migration weight should be > 0");
+
+        // check results
+
+        assert_eq!(FeeRate::<Test>::get(netuid1), value1);
+        assert_eq!(FeeRate::<Test>::get(netuid2), MAX_FEE_RATE);
+
+        // running the migration again should do nothing
+        FeeRate::<Test>::insert(netuid1, value1);
+        FeeRate::<Test>::insert(netuid2, value2);
+
+        let _weight2 = migrations::fee_rate_migration::migrate_fee_rate::<Test>();
+
+        assert!(
+            HasMigrationRun::<Test>::get(migration_name.clone()),
+            "HasMigrationRun remains true on second run"
+        );
+        assert_eq!(FeeRate::<Test>::get(netuid1), value1);
+        assert_eq!(FeeRate::<Test>::get(netuid2), value2);
+    });
+}