Add migration for RateLimitKey

ales-otf · ales-otf · commit 712d45947748 · 2025-10-20T22:37:27.000+03:00
diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs
@@ -2295,16 +2295,22 @@ impl<T: Config + pallet_balances::Config<Balance = u64>>
 #[derive(Encode, Decode, Clone, PartialEq, Eq, Debug, TypeInfo)]
 pub enum RateLimitKey<AccountId> {
     // The setting sn owner hotkey operation is rate limited per netuid
+    #[codec(index = 0)]
     SetSNOwnerHotkey(NetUid),
     // Generic rate limit for subnet-owner hyperparameter updates (per netuid)
+    #[codec(index = 1)]
     OwnerHyperparamUpdate(NetUid, Hyperparameter),
     // Subnet registration rate limit
+    #[codec(index = 2)]
     NetworkLastRegistered,
     // Last tx block limit per account ID
+    #[codec(index = 3)]
     LastTxBlock(AccountId),
     // Last tx block child key limit per account ID
+    #[codec(index = 4)]
     LastTxBlockChildKeyTake(AccountId),
     // Last tx block delegate key limit per account ID
+    #[codec(index = 5)]
     LastTxBlockDelegateTake(AccountId),
 }
 
diff --git a/pallets/subtensor/src/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs
@@ -138,6 +138,8 @@ mod hooks {
                 .saturating_add(migrations::migrate_fix_root_tao_and_alpha_in::migrate_fix_root_tao_and_alpha_in::<T>())
                 // Migrate last block rate limiting storage items
                 .saturating_add(migrations::migrate_rate_limiting_last_blocks::migrate_obsolete_rate_limiting_last_blocks_storage::<T>())
+                // Re-encode rate limit keys after introducing OwnerHyperparamUpdate variant
+                .saturating_add(migrations::migrate_rate_limit_keys::migrate_rate_limit_keys::<T>())
                 // Migrate remove network modality
                 .saturating_add(migrations::migrate_remove_network_modality::migrate_remove_network_modality::<T>())
                 // Migrate Immunity Period
diff --git a/pallets/subtensor/src/migrations/migrate_rate_limit_keys.rs b/pallets/subtensor/src/migrations/migrate_rate_limit_keys.rs
@@ -0,0 +1,239 @@
+use alloc::string::String;
+use codec::{Decode, Encode};
+use frame_support::traits::Get;
+use frame_support::weights::Weight;
+use sp_io::hashing::twox_128;
+use sp_io::storage;
+use sp_std::{collections::btree_set::BTreeSet, vec::Vec};
+use subtensor_runtime_common::NetUid;
+
+use crate::{
+    ChildKeys, Config, Delegates, HasMigrationRun, LastRateLimitedBlock, ParentKeys,
+    PendingChildKeys, RateLimitKey,
+};
+
+const MIGRATION_NAME: &[u8] = b"migrate_rate_limit_keys";
+
+#[allow(dead_code)]
+#[derive(Decode)]
+enum RateLimitKeyV0<AccountId> {
+    SetSNOwnerHotkey(NetUid),
+    NetworkLastRegistered,
+    LastTxBlock(AccountId),
+    LastTxBlockChildKeyTake(AccountId),
+    LastTxBlockDelegateTake(AccountId),
+}
+
+pub fn migrate_rate_limit_keys<T: Config>() -> Weight
+where
+    T::AccountId: Ord + Clone,
+{
+    let mut weight = T::DbWeight::get().reads(1);
+
+    if HasMigrationRun::<T>::get(MIGRATION_NAME) {
+        log::info!(
+            "Migration '{}' already executed - skipping",
+            String::from_utf8_lossy(MIGRATION_NAME)
+        );
+        return weight;
+    }
+
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(MIGRATION_NAME)
+    );
+
+    let (child_accounts, child_weight) = collect_child_related_accounts::<T>();
+    let (delegate_accounts, delegate_weight) = collect_delegate_accounts::<T>();
+    weight = weight.saturating_add(child_weight);
+    weight = weight.saturating_add(delegate_weight);
+
+    let prefix = storage_prefix("SubtensorModule", "LastRateLimitedBlock");
+    let mut cursor = prefix.clone();
+    let mut entries = Vec::new();
+
+    while let Some(next_key) = storage::next_key(&cursor) {
+        if !next_key.starts_with(&prefix) {
+            break;
+        }
+        if let Some(value) = storage::get(&next_key) {
+            entries.push((next_key.clone(), value));
+        }
+        cursor = next_key;
+    }
+
+    weight = weight.saturating_add(T::DbWeight::get().reads(entries.len() as u64));
+
+    let mut migrated_network = 0u64;
+    let mut migrated_last_tx = 0u64;
+    let mut migrated_child_take = 0u64;
+    let mut migrated_delegate_take = 0u64;
+
+    for (old_storage_key, value_bytes) in entries {
+        if value_bytes.is_empty() {
+            continue;
+        }
+
+        let encoded_key = &old_storage_key[prefix.len()..];
+        if encoded_key.is_empty() {
+            continue;
+        }
+
+        let Some(decoded_legacy) = decode_legacy::<T>(&encoded_key) else {
+            // Unknown entry – skip to avoid clobbering valid data.
+            continue;
+        };
+
+        let legacy_value = match decode_value(&value_bytes) {
+            Some(v) => v,
+            None => continue,
+        };
+
+        let Some(modern_key) =
+            legacy_to_modern(decoded_legacy, &child_accounts, &delegate_accounts)
+        else {
+            continue;
+        };
+        let new_storage_key = LastRateLimitedBlock::<T>::hashed_key_for(&modern_key);
+        weight = weight.saturating_add(T::DbWeight::get().reads(1));
+
+        let merged_value = storage::get(&new_storage_key)
+            .and_then(|data| decode_value(&data))
+            .map_or(legacy_value, |current| {
+                core::cmp::max(current, legacy_value)
+            });
+
+        storage::set(&new_storage_key, &merged_value.encode());
+        if new_storage_key != old_storage_key {
+            storage::clear(&old_storage_key);
+            weight = weight.saturating_add(T::DbWeight::get().writes(1));
+        }
+
+        weight = weight.saturating_add(T::DbWeight::get().writes(1));
+        match &modern_key {
+            RateLimitKey::NetworkLastRegistered => {
+                migrated_network = migrated_network.saturating_add(1);
+            }
+            RateLimitKey::LastTxBlock(_) => {
+                migrated_last_tx = migrated_last_tx.saturating_add(1);
+            }
+            RateLimitKey::LastTxBlockChildKeyTake(_) => {
+                migrated_child_take = migrated_child_take.saturating_add(1);
+            }
+            RateLimitKey::LastTxBlockDelegateTake(_) => {
+                migrated_delegate_take = migrated_delegate_take.saturating_add(1);
+            }
+            _ => {}
+        }
+    }
+
+    HasMigrationRun::<T>::insert(MIGRATION_NAME, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{}' completed. network={}, last_tx={}, child_take={}, delegate_take={}",
+        String::from_utf8_lossy(MIGRATION_NAME),
+        migrated_network,
+        migrated_last_tx,
+        migrated_child_take,
+        migrated_delegate_take
+    );
+
+    weight
+}
+
+fn storage_prefix(pallet: &str, storage: &str) -> Vec<u8> {
+    let pallet_hash = twox_128(pallet.as_bytes());
+    let storage_hash = twox_128(storage.as_bytes());
+    [pallet_hash, storage_hash].concat()
+}
+
+fn decode_legacy<T: Config>(bytes: &[u8]) -> Option<RateLimitKeyV0<T::AccountId>> {
+    let mut slice = bytes;
+    let decoded = RateLimitKeyV0::<T::AccountId>::decode(&mut slice).ok()?;
+    if slice.is_empty() {
+        Some(decoded)
+    } else {
+        None
+    }
+}
+
+fn decode_value(bytes: &[u8]) -> Option<u64> {
+    let mut slice = bytes;
+    u64::decode(&mut slice).ok()
+}
+
+fn legacy_to_modern<AccountId: Ord + Clone>(
+    legacy: RateLimitKeyV0<AccountId>,
+    child_accounts: &BTreeSet<AccountId>,
+    delegate_accounts: &BTreeSet<AccountId>,
+) -> Option<RateLimitKey<AccountId>> {
+    match legacy {
+        RateLimitKeyV0::SetSNOwnerHotkey(_) => None,
+        RateLimitKeyV0::NetworkLastRegistered => Some(RateLimitKey::NetworkLastRegistered),
+        RateLimitKeyV0::LastTxBlock(account) => Some(RateLimitKey::LastTxBlock(account)),
+        RateLimitKeyV0::LastTxBlockChildKeyTake(account) => {
+            if child_accounts.contains(&account) {
+                Some(RateLimitKey::LastTxBlockChildKeyTake(account))
+            } else {
+                None
+            }
+        }
+        RateLimitKeyV0::LastTxBlockDelegateTake(account) => {
+            if delegate_accounts.contains(&account) {
+                Some(RateLimitKey::LastTxBlockDelegateTake(account))
+            } else {
+                None
+            }
+        }
+    }
+}
+
+fn collect_child_related_accounts<T: Config>() -> (BTreeSet<T::AccountId>, Weight)
+where
+    T::AccountId: Ord + Clone,
+{
+    let mut accounts = BTreeSet::new();
+    let mut reads = 0u64;
+
+    for (parent, _, children) in ChildKeys::<T>::iter() {
+        accounts.insert(parent.clone());
+        for (_, child) in children {
+            accounts.insert(child.clone());
+        }
+        reads = reads.saturating_add(1);
+    }
+
+    for (_, parent, (children, _)) in PendingChildKeys::<T>::iter() {
+        accounts.insert(parent.clone());
+        for (_, child) in children {
+            accounts.insert(child.clone());
+        }
+        reads = reads.saturating_add(1);
+    }
+
+    for (child, _, parents) in ParentKeys::<T>::iter() {
+        accounts.insert(child.clone());
+        for (_, parent) in parents {
+            accounts.insert(parent.clone());
+        }
+        reads = reads.saturating_add(1);
+    }
+
+    (accounts, T::DbWeight::get().reads(reads))
+}
+
+fn collect_delegate_accounts<T: Config>() -> (BTreeSet<T::AccountId>, Weight)
+where
+    T::AccountId: Ord + Clone,
+{
+    let mut accounts = BTreeSet::new();
+    let mut reads = 0u64;
+
+    for (account, _) in Delegates::<T>::iter() {
+        accounts.insert(account.clone());
+        reads = reads.saturating_add(1);
+    }
+
+    (accounts, T::DbWeight::get().reads(reads))
+}
diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs
@@ -27,6 +27,7 @@ pub mod migrate_network_lock_reduction_interval;
 pub mod migrate_orphaned_storage_items;
 pub mod migrate_populate_owned_hotkeys;
 pub mod migrate_rao;
+pub mod migrate_rate_limit_keys;
 pub mod migrate_rate_limiting_last_blocks;
 pub mod migrate_remove_commitments_rate_limit;
 pub mod migrate_remove_network_modality;
diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs
@@ -1027,6 +1027,117 @@ fn test_migrate_last_tx_block_delegate_take() {
     });
 }
 
+#[test]
+fn test_migrate_rate_limit_keys() {
+    new_test_ext(1).execute_with(|| {
+        const MIGRATION_NAME: &[u8] = b"migrate_rate_limit_keys";
+        let prefix = {
+            let pallet_prefix = twox_128("SubtensorModule".as_bytes());
+            let storage_prefix = twox_128("LastRateLimitedBlock".as_bytes());
+            [pallet_prefix, storage_prefix].concat()
+        };
+
+        // Seed new-format entries that must survive the migration untouched.
+        let new_last_account = U256::from(10);
+        SubtensorModule::set_last_tx_block(&new_last_account, 555);
+        let new_child_account = U256::from(11);
+        SubtensorModule::set_last_tx_block_childkey(&new_child_account, 777);
+        let new_delegate_account = U256::from(12);
+        SubtensorModule::set_last_tx_block_delegate_take(&new_delegate_account, 888);
+
+        // Legacy NetworkLastRegistered entry (index 1)
+        let mut legacy_network_key = prefix.clone();
+        legacy_network_key.push(1u8);
+        sp_io::storage::set(&legacy_network_key, &111u64.encode());
+
+        // Legacy LastTxBlock entry (index 2) for an account that already has a new-format value.
+        let mut legacy_last_key = prefix.clone();
+        legacy_last_key.push(2u8);
+        legacy_last_key.extend_from_slice(&new_last_account.encode());
+        sp_io::storage::set(&legacy_last_key, &666u64.encode());
+
+        // Legacy LastTxBlockChildKeyTake entry (index 3)
+        let legacy_child_account = U256::from(3);
+        ChildKeys::<Test>::insert(
+            legacy_child_account,
+            NetUid::from(0),
+            vec![(0u64, U256::from(99))],
+        );
+        let mut legacy_child_key = prefix.clone();
+        legacy_child_key.push(3u8);
+        legacy_child_key.extend_from_slice(&legacy_child_account.encode());
+        sp_io::storage::set(&legacy_child_key, &333u64.encode());
+
+        // Legacy LastTxBlockDelegateTake entry (index 4)
+        let legacy_delegate_account = U256::from(4);
+        Delegates::<Test>::insert(legacy_delegate_account, 500u16);
+        let mut legacy_delegate_key = prefix.clone();
+        legacy_delegate_key.push(4u8);
+        legacy_delegate_key.extend_from_slice(&legacy_delegate_account.encode());
+        sp_io::storage::set(&legacy_delegate_key, &444u64.encode());
+
+        let weight = crate::migrations::migrate_rate_limit_keys::migrate_rate_limit_keys::<Test>();
+        assert!(
+            HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should be marked as executed"
+        );
+        assert!(!weight.is_zero(), "Migration weight should be non-zero");
+
+        // Legacy entries were migrated and cleared.
+        assert_eq!(
+            SubtensorModule::get_network_last_lock_block(),
+            111u64,
+            "Network last lock block should match migrated value"
+        );
+        assert!(
+            sp_io::storage::get(&legacy_network_key).is_none(),
+            "Legacy network entry should be cleared"
+        );
+
+        assert_eq!(
+            SubtensorModule::get_last_tx_block(&new_last_account),
+            666u64,
+            "LastTxBlock should reflect the merged legacy value"
+        );
+        assert!(
+            sp_io::storage::get(&legacy_last_key).is_none(),
+            "Legacy LastTxBlock entry should be cleared"
+        );
+
+        assert_eq!(
+            SubtensorModule::get_last_tx_block_childkey_take(&legacy_child_account),
+            333u64,
+            "Child key take block should be migrated"
+        );
+        assert!(
+            sp_io::storage::get(&legacy_child_key).is_none(),
+            "Legacy child take entry should be cleared"
+        );
+
+        assert_eq!(
+            SubtensorModule::get_last_tx_block_delegate_take(&legacy_delegate_account),
+            444u64,
+            "Delegate take block should be migrated"
+        );
+        assert!(
+            sp_io::storage::get(&legacy_delegate_key).is_none(),
+            "Legacy delegate take entry should be cleared"
+        );
+
+        // New-format entries remain untouched.
+        assert_eq!(
+            SubtensorModule::get_last_tx_block_childkey_take(&new_child_account),
+            777u64,
+            "Existing child take entry should be preserved"
+        );
+        assert_eq!(
+            SubtensorModule::get_last_tx_block_delegate_take(&new_delegate_account),
+            888u64,
+            "Existing delegate take entry should be preserved"
+        );
+    });
+}
+
 #[test]
 fn test_migrate_fix_root_subnet_tao() {
     new_test_ext(1).execute_with(|| {
