diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index c52c589ad6..c73ef19ca0 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -2295,16 +2295,22 @@ impl> #[derive(Encode, Decode, Clone, PartialEq, Eq, Debug, TypeInfo)] pub enum RateLimitKey { // 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 index 20651f68fc..ddde143cf9 100644 --- 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::()) // Migrate last block rate limiting storage items .saturating_add(migrations::migrate_rate_limiting_last_blocks::migrate_obsolete_rate_limiting_last_blocks_storage::()) + // Re-encode rate limit keys after introducing OwnerHyperparamUpdate variant + .saturating_add(migrations::migrate_rate_limit_keys::migrate_rate_limit_keys::()) // Migrate remove network modality .saturating_add(migrations::migrate_remove_network_modality::migrate_remove_network_modality::()) // 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 new file mode 100644 index 0000000000..e6e331fb63 --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_rate_limit_keys.rs @@ -0,0 +1,241 @@ +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 { + SetSNOwnerHotkey(NetUid), + NetworkLastRegistered, + LastTxBlock(AccountId), + LastTxBlockChildKeyTake(AccountId), + LastTxBlockDelegateTake(AccountId), +} + +pub fn migrate_rate_limit_keys() -> Weight +where + T::AccountId: Ord + Clone, +{ + let mut weight = T::DbWeight::get().reads(1); + + if HasMigrationRun::::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::(); + let (delegate_accounts, delegate_weight) = collect_delegate_accounts::(); + 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 Some(encoded_key) = old_storage_key.get(prefix.len()..) else { + continue; + }; + if encoded_key.is_empty() { + continue; + } + + let Some(decoded_legacy) = decode_legacy::(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::::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::::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 { + let pallet_hash = twox_128(pallet.as_bytes()); + let storage_hash = twox_128(storage.as_bytes()); + [pallet_hash, storage_hash].concat() +} + +fn decode_legacy(bytes: &[u8]) -> Option> { + let mut slice = bytes; + let decoded = RateLimitKeyV0::::decode(&mut slice).ok()?; + if slice.is_empty() { + Some(decoded) + } else { + None + } +} + +fn decode_value(bytes: &[u8]) -> Option { + let mut slice = bytes; + u64::decode(&mut slice).ok() +} + +fn legacy_to_modern( + legacy: RateLimitKeyV0, + child_accounts: &BTreeSet, + delegate_accounts: &BTreeSet, +) -> Option> { + 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() -> (BTreeSet, Weight) +where + T::AccountId: Ord + Clone, +{ + let mut accounts = BTreeSet::new(); + let mut reads = 0u64; + + for (parent, _, children) in ChildKeys::::iter() { + accounts.insert(parent.clone()); + for (_, child) in children { + accounts.insert(child.clone()); + } + reads = reads.saturating_add(1); + } + + for (_, parent, (children, _)) in PendingChildKeys::::iter() { + accounts.insert(parent.clone()); + for (_, child) in children { + accounts.insert(child.clone()); + } + reads = reads.saturating_add(1); + } + + for (child, _, parents) in ParentKeys::::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() -> (BTreeSet, Weight) +where + T::AccountId: Ord + Clone, +{ + let mut accounts = BTreeSet::new(); + let mut reads = 0u64; + + for (account, _) in Delegates::::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 index e92b86aa29..57a993f2ff 100644 --- 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 index ede32aa06c..53ac3409d7 100644 --- 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::::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::::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::(); + assert!( + HasMigrationRun::::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(|| {