diff --git a/pallets/subtensor/rpc/src/lib.rs b/pallets/subtensor/rpc/src/lib.rs index 2df65242d..b5749988e 100644 --- a/pallets/subtensor/rpc/src/lib.rs +++ b/pallets/subtensor/rpc/src/lib.rs @@ -92,6 +92,13 @@ pub trait SubtensorCustomApi { metagraph_index: Vec, at: Option, ) -> RpcResult>; + #[method(name = "subnetInfo_getColdkeyAutoStakeHotkey")] + fn get_coldkey_auto_stake_hotkey( + &self, + coldkey: AccountId32, + netuid: NetUid, + at: Option, + ) -> RpcResult>; #[method(name = "subnetInfo_getSelectiveMechagraph")] fn get_selective_mechagraph( &self, @@ -474,6 +481,24 @@ where } } + fn get_coldkey_auto_stake_hotkey( + &self, + coldkey: AccountId32, + netuid: NetUid, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + match api.get_coldkey_auto_stake_hotkey(at, coldkey, netuid) { + Ok(result) => Ok(result.encode()), + Err(e) => Err(Error::RuntimeError(format!( + "Unable to get coldkey auto stake hotkey: {e:?}" + )) + .into()), + } + } + fn get_selective_mechagraph( &self, netuid: NetUid, diff --git a/pallets/subtensor/runtime-api/src/lib.rs b/pallets/subtensor/runtime-api/src/lib.rs index a2df8a351..f1107bad0 100644 --- a/pallets/subtensor/runtime-api/src/lib.rs +++ b/pallets/subtensor/runtime-api/src/lib.rs @@ -45,6 +45,7 @@ sp_api::decl_runtime_apis! { fn get_dynamic_info(netuid: NetUid) -> Option>; fn get_subnet_state(netuid: NetUid) -> Option>; fn get_selective_metagraph(netuid: NetUid, metagraph_indexes: Vec) -> Option>; + fn get_coldkey_auto_stake_hotkey(coldkey: AccountId32, netuid: NetUid) -> Option; fn get_selective_mechagraph(netuid: NetUid, subid: MechId, metagraph_indexes: Vec) -> Option>; fn get_subnet_to_prune() -> Option; } diff --git a/pallets/subtensor/src/benchmarks.rs b/pallets/subtensor/src/benchmarks.rs index 98307f521..5b73ac3f4 100644 --- a/pallets/subtensor/src/benchmarks.rs +++ b/pallets/subtensor/src/benchmarks.rs @@ -1560,9 +1560,21 @@ mod pallet_benchmarks { #[benchmark] fn set_coldkey_auto_stake_hotkey() { let coldkey: T::AccountId = whitelisted_caller(); - let hot: T::AccountId = account("A", 0, 1); + let netuid = NetUid::from(1); + let hotkey: T::AccountId = account("A", 0, 1); + SubtokenEnabled::::insert(netuid, true); + Subtensor::::init_new_network(netuid, 1); + let amount = 900_000_000_000; + + Subtensor::::add_balance_to_coldkey_account(&coldkey.clone(), amount); + + assert_ok!(Subtensor::::burned_register( + RawOrigin::Signed(coldkey.clone()).into(), + netuid, + hotkey.clone() + )); #[extrinsic_call] - _(RawOrigin::Signed(coldkey.clone()), hot.clone()); + _(RawOrigin::Signed(coldkey.clone()), netuid, hotkey.clone()); } } diff --git a/pallets/subtensor/src/coinbase/run_coinbase.rs b/pallets/subtensor/src/coinbase/run_coinbase.rs index 3f2f715df..8748fe8ee 100644 --- a/pallets/subtensor/src/coinbase/run_coinbase.rs +++ b/pallets/subtensor/src/coinbase/run_coinbase.rs @@ -514,7 +514,7 @@ impl Pallet { } let owner: T::AccountId = Owner::::get(&hotkey); - let maybe_dest = AutoStakeDestination::::get(&owner); + let maybe_dest = AutoStakeDestination::::get(&owner, netuid); // Always stake but only emit event if autostake is set. let destination = maybe_dest.clone().unwrap_or(hotkey.clone()); diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index ab0809050..e33598472 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -1140,9 +1140,16 @@ pub mod pallet { #[pallet::storage] // --- MAP ( cold ) --> Vec | Returns the vector of hotkeys controlled by this coldkey. pub type OwnedHotkeys = StorageMap<_, Blake2_128Concat, T::AccountId, Vec, ValueQuery>; - #[pallet::storage] // --- MAP ( cold ) --> hot | Returns the hotkey a coldkey will autostake to with mining rewards. - pub type AutoStakeDestination = - StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, OptionQuery>; + #[pallet::storage] // --- DMAP ( cold, netuid )--> hot | Returns the hotkey a coldkey will autostake to with mining rewards. + pub type AutoStakeDestination = StorageDoubleMap< + _, + Blake2_128Concat, + T::AccountId, + Identity, + NetUid, + T::AccountId, + OptionQuery, + >; #[pallet::storage] // --- DMAP ( cold ) --> (block_expected, new_coldkey) | Maps coldkey to the block to swap at and new coldkey. pub type ColdkeySwapScheduled = StorageMap< diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index 485dce649..8883fefa3 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -1086,7 +1086,7 @@ mod dispatches { /// Weight is calculated based on the number of database reads and writes. #[pallet::call_index(71)] #[pallet::weight((Weight::from_parts(161_700_000, 0) - .saturating_add(T::DbWeight::get().reads(15_u64)) + .saturating_add(T::DbWeight::get().reads(16_u64)) .saturating_add(T::DbWeight::get().writes(9)), DispatchClass::Operational, Pays::No))] pub fn swap_coldkey( origin: OriginFor, @@ -2292,16 +2292,36 @@ mod dispatches { /// * `hotkey` (T::AccountId): /// - The hotkey account to designate as the autostake destination. #[pallet::call_index(114)] - #[pallet::weight((Weight::from_parts(5_170_000, 0) - .saturating_add(T::DbWeight::get().reads(0_u64)) + #[pallet::weight((Weight::from_parts(29_930_000, 0) + .saturating_add(T::DbWeight::get().reads(3_u64)) .saturating_add(T::DbWeight::get().writes(1)), DispatchClass::Normal, Pays::No))] pub fn set_coldkey_auto_stake_hotkey( origin: T::RuntimeOrigin, + netuid: NetUid, hotkey: T::AccountId, ) -> DispatchResult { let coldkey = ensure_signed(origin)?; + ensure!(Self::if_subnet_exist(netuid), Error::::SubnetNotExists); + ensure!( + Uids::::contains_key(netuid, &hotkey), + Error::::HotKeyNotRegisteredInSubNet + ); + + let current_hotkey = AutoStakeDestination::::get(coldkey.clone(), netuid); + if let Some(current_hotkey) = current_hotkey { + ensure!( + current_hotkey != hotkey, + Error::::SameAutoStakeHotkeyAlreadySet + ); + } + + AutoStakeDestination::::insert(coldkey.clone(), netuid, hotkey.clone()); - AutoStakeDestination::::insert(coldkey, hotkey.clone()); + Self::deposit_event(Event::AutoStakeDestinationSet { + coldkey, + netuid, + hotkey, + }); Ok(()) } diff --git a/pallets/subtensor/src/macros/errors.rs b/pallets/subtensor/src/macros/errors.rs index fbae1e107..e241ff068 100644 --- a/pallets/subtensor/src/macros/errors.rs +++ b/pallets/subtensor/src/macros/errors.rs @@ -256,6 +256,8 @@ mod errors { CannotAffordLockCost, /// exceeded the rate limit for associating an EVM key. EvmKeyAssociateRateLimitExceeded, + /// Same auto stake hotkey already set + SameAutoStakeHotkeyAlreadySet, /// The UID map for the subnet could not be cleared UidMapCouldNotBeCleared, /// Trimming would exceed the max immune neurons percentage diff --git a/pallets/subtensor/src/macros/events.rs b/pallets/subtensor/src/macros/events.rs index f2d134c18..c34219d53 100644 --- a/pallets/subtensor/src/macros/events.rs +++ b/pallets/subtensor/src/macros/events.rs @@ -442,5 +442,19 @@ mod events { /// The minimum allowed UIDs for a subnet have been set. MinAllowedUidsSet(NetUid, u16), + + /// The auto stake destination has been set. + /// + /// - **coldkey**: The account ID of the coldkey. + /// - **netuid**: The network identifier. + /// - **hotkey**: The account ID of the hotkey. + AutoStakeDestinationSet { + /// The account ID of the coldkey. + coldkey: T::AccountId, + /// The network identifier. + netuid: NetUid, + /// The account ID of the hotkey. + hotkey: T::AccountId, + }, } } diff --git a/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs b/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs new file mode 100644 index 000000000..13f165db4 --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs @@ -0,0 +1,82 @@ +use super::*; +use crate::AccountIdOf; +use frame_support::{ + IterableStorageMap, + pallet_prelude::{Blake2_128Concat, OptionQuery}, + storage_alias, + traits::Get, + weights::Weight, +}; +use scale_info::prelude::string::String; + +/// Module containing deprecated storage format for AutoStakeDestination +pub mod deprecated_auto_stake_destination_format { + use super::*; + + #[storage_alias] + pub(super) type AutoStakeDestination = + StorageMap, Blake2_128Concat, AccountIdOf, AccountIdOf, OptionQuery>; +} + +/// Migrate the AutoStakeDestination map from single map to double map format +pub fn migrate_auto_stake_destination() -> Weight { + use deprecated_auto_stake_destination_format as old; + + let migration_name = b"migrate_auto_stake_destination".to_vec(); + let mut weight = T::DbWeight::get().reads(1); + + if HasMigrationRun::::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) + ); + + // ------------------------------ + // Step 1: Migrate AutoStakeDestination entries + // ------------------------------ + + let curr_keys: Vec> = old::AutoStakeDestination::::iter_keys().collect(); + let root_netuid = NetUid::ROOT; + let netuids: Vec = as IterableStorageMap>::iter() + .map(|(netuid, _)| netuid) + .collect(); + + for coldkey in &curr_keys { + weight.saturating_accrue(T::DbWeight::get().reads(1)); + + if let Some(hotkey) = old::AutoStakeDestination::::get(coldkey) { + for netuid in netuids.iter() { + if *netuid == root_netuid { + continue; + } + AutoStakeDestination::::insert(coldkey, netuid, hotkey.clone()); + } + + old::AutoStakeDestination::::remove(coldkey); + + weight.saturating_accrue(T::DbWeight::get().writes(netuids.len() as u64)); + } + } + + // ------------------------------ + // Step 2: Mark Migration as Completed + // ------------------------------ + + HasMigrationRun::::insert(&migration_name, true); + weight = weight.saturating_add(T::DbWeight::get().writes(1)); + + log::info!( + "Migration '{:?}' completed successfully. {} entries migrated.", + String::from_utf8_lossy(&migration_name), + curr_keys.len() + ); + + weight +} diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs index e7c50c008..ef2df8bde 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -4,6 +4,7 @@ use frame_support::pallet_prelude::Weight; use sp_io::KillStorageResult; use sp_io::hashing::twox_128; use sp_io::storage::clear_prefix; +pub mod migrate_auto_stake_destination; pub mod migrate_chain_identity; pub mod migrate_coldkey_swap_scheduled; pub mod migrate_commit_reveal_settings; diff --git a/pallets/subtensor/src/rpc_info/subnet_info.rs b/pallets/subtensor/src/rpc_info/subnet_info.rs index 11d60b2bc..6a7966b4f 100644 --- a/pallets/subtensor/src/rpc_info/subnet_info.rs +++ b/pallets/subtensor/src/rpc_info/subnet_info.rs @@ -403,4 +403,11 @@ impl Pallet { user_liquidity_enabled, }) } + + pub fn get_coldkey_auto_stake_hotkey( + coldkey: T::AccountId, + netuid: NetUid, + ) -> Option { + AutoStakeDestination::::get(coldkey, netuid) + } } diff --git a/pallets/subtensor/src/swap/swap_coldkey.rs b/pallets/subtensor/src/swap/swap_coldkey.rs index 8180650bf..afee144e5 100644 --- a/pallets/subtensor/src/swap/swap_coldkey.rs +++ b/pallets/subtensor/src/swap/swap_coldkey.rs @@ -155,6 +155,12 @@ impl Pallet { SubnetOwner::::insert(netuid, new_coldkey.clone()); } weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + + if let Some(old_auto_stake_hotkey) = AutoStakeDestination::::get(old_coldkey, netuid) + { + AutoStakeDestination::::remove(old_coldkey, netuid); + AutoStakeDestination::::insert(new_coldkey, netuid, old_auto_stake_hotkey); + } } // 3. Swap Stake. @@ -178,11 +184,6 @@ impl Pallet { weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2)); } - if let Some(old_auto_stake_hotkey) = AutoStakeDestination::::get(old_coldkey) { - AutoStakeDestination::::remove(old_coldkey); - AutoStakeDestination::::insert(new_coldkey, old_auto_stake_hotkey); - } - // 4. Swap TotalColdkeyAlpha (DEPRECATED) // for netuid in Self::get_all_subnet_netuids() { // let old_alpha_stake: u64 = TotalColdkeyAlpha::::get(old_coldkey, netuid); diff --git a/pallets/subtensor/src/tests/auto_stake_hotkey.rs b/pallets/subtensor/src/tests/auto_stake_hotkey.rs new file mode 100644 index 000000000..242d469c2 --- /dev/null +++ b/pallets/subtensor/src/tests/auto_stake_hotkey.rs @@ -0,0 +1,112 @@ +use super::mock::*; +use crate::*; +use frame_support::{assert_noop, assert_ok}; +use sp_core::U256; +use subtensor_runtime_common::NetUid; + +#[test] +fn test_set_coldkey_auto_stake_hotkey_subnet_not_exists() { + new_test_ext(1).execute_with(|| { + let coldkey = U256::from(1); + let hotkey = U256::from(2); + let netuid = NetUid::from(999); // Non-existent subnet + + assert_noop!( + SubtensorModule::set_coldkey_auto_stake_hotkey( + RuntimeOrigin::signed(coldkey), + netuid, + hotkey, + ), + Error::::SubnetNotExists + ); + }); +} + +#[test] +fn test_set_coldkey_auto_stake_hotkey_hotkey_not_registered() { + new_test_ext(1).execute_with(|| { + let subnet_owner_ck = U256::from(0); + let subnet_owner_hk = U256::from(1); + + let coldkey = U256::from(10); + let hotkey = U256::from(11); // Hotkey not registered in subnet + + let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck); + + assert_noop!( + SubtensorModule::set_coldkey_auto_stake_hotkey( + RuntimeOrigin::signed(coldkey), + netuid, + hotkey, + ), + Error::::HotKeyNotRegisteredInSubNet + ); + }); +} + +#[test] +fn test_set_coldkey_auto_stake_hotkey_success() { + new_test_ext(1).execute_with(|| { + let subnet_owner_ck = U256::from(0); + let subnet_owner_hk = U256::from(1); + + let coldkey = U256::from(10); + let hotkey = U256::from(11); + + Owner::::insert(hotkey, coldkey); + OwnedHotkeys::::insert(coldkey, vec![hotkey]); + + let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck); + Uids::::insert(netuid, hotkey, 1); + + // Verify no destination is set initially + assert_eq!(AutoStakeDestination::::get(coldkey, netuid), None); + + // Call should succeed + assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey( + RuntimeOrigin::signed(coldkey), + netuid, + hotkey, + )); + + // Verify destination is now set + assert_eq!( + AutoStakeDestination::::get(coldkey, netuid), + Some(hotkey) + ); + }); +} + +#[test] +fn test_set_coldkey_auto_stake_hotkey_same_hotkey_again() { + new_test_ext(1).execute_with(|| { + let subnet_owner_ck = U256::from(0); + let subnet_owner_hk = U256::from(1); + + let coldkey = U256::from(10); + let hotkey = U256::from(11); + + Owner::::insert(hotkey, coldkey); + OwnedHotkeys::::insert(coldkey, vec![hotkey]); + + let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck); + Uids::::insert(netuid, hotkey, 1); + + // First call should succeed + assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey( + RuntimeOrigin::signed(coldkey), + netuid, + hotkey, + )); + + // Second call with same hotkey should fail + assert_noop!( + SubtensorModule::set_coldkey_auto_stake_hotkey( + RuntimeOrigin::signed(coldkey), + netuid, + hotkey, + ), + Error::::SameAutoStakeHotkeyAlreadySet + ); + }); +} diff --git a/pallets/subtensor/src/tests/coinbase.rs b/pallets/subtensor/src/tests/coinbase.rs index a4ea3988b..b53308ae9 100644 --- a/pallets/subtensor/src/tests/coinbase.rs +++ b/pallets/subtensor/src/tests/coinbase.rs @@ -2797,6 +2797,7 @@ fn test_incentive_is_autostaked_to_owner_destination() { // Set autostake destination for the miner's coldkey assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey( RuntimeOrigin::signed(miner_ck), + netuid, dest_hk, )); diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index 816c87837..b37c25f04 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -1502,6 +1502,110 @@ fn test_migrate_commit_reveal_settings_values_access() { }); } +#[test] +fn test_migrate_auto_stake_destination() { + new_test_ext(1).execute_with(|| { + // ------------------------------ + // Step 1: Simulate Old Storage Entries + // ------------------------------ + const MIGRATION_NAME: &[u8] = b"migrate_auto_stake_destination"; + let netuids = [NetUid::ROOT, NetUid::from(1), NetUid::from(2), NetUid::from(42)]; + for netuid in &netuids { + NetworksAdded::::insert(*netuid, true); + } + + let pallet_prefix = twox_128("SubtensorModule".as_bytes()); + let storage_prefix = twox_128("AutoStakeDestination".as_bytes()); + + // Create test accounts + let coldkey1: U256 = U256::from(1); + let coldkey2: U256 = U256::from(2); + let hotkey1: U256 = U256::from(100); + let hotkey2: U256 = U256::from(200); + + // Construct storage keys for old format (StorageMap) + let mut key1 = Vec::new(); + key1.extend_from_slice(&pallet_prefix); + key1.extend_from_slice(&storage_prefix); + key1.extend_from_slice(&Blake2_128Concat::hash(&coldkey1.encode())); + + let mut key2 = Vec::new(); + key2.extend_from_slice(&pallet_prefix); + key2.extend_from_slice(&storage_prefix); + key2.extend_from_slice(&Blake2_128Concat::hash(&coldkey2.encode())); + + // Store old format entries + put_raw(&key1, &hotkey1.encode()); + put_raw(&key2, &hotkey2.encode()); + + // Verify old entries are stored + assert_eq!(get_raw(&key1), Some(hotkey1.encode())); + assert_eq!(get_raw(&key2), Some(hotkey2.encode())); + + assert!( + !HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should not have run yet" + ); + + // ------------------------------ + // Step 2: Run the Migration + // ------------------------------ + let weight = crate::migrations::migrate_auto_stake_destination::migrate_auto_stake_destination::(); + + assert!( + HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should be marked as run" + ); + + // ------------------------------ + // Step 3: Verify Migration Effects + // ------------------------------ + + // Verify new format entries exist + for netuid in &netuids { + if *netuid == NetUid::ROOT { + assert_eq!( + AutoStakeDestination::::get(coldkey1, NetUid::ROOT), + None + ); + assert_eq!( + AutoStakeDestination::::get(coldkey2, NetUid::ROOT), + None + ); + } else { + assert_eq!( + AutoStakeDestination::::get(coldkey1, *netuid), + Some(hotkey1) + ); + assert_eq!( + AutoStakeDestination::::get(coldkey2, *netuid), + Some(hotkey2) + ); + } + + } + + // Verify old format entries are cleared + assert_eq!(get_raw(&key1), None, "Old storage entry 1 should be cleared"); + assert_eq!(get_raw(&key2), None, "Old storage entry 2 should be cleared"); + + // Verify weight calculation + assert!(!weight.is_zero(), "Migration weight should be non-zero"); + + // ------------------------------ + // Step 4: Test Migration Idempotency + // ------------------------------ + let weight_second_run = crate::migrations::migrate_auto_stake_destination::migrate_auto_stake_destination::(); + + // Second run should only read the migration flag + assert_eq!( + weight_second_run, + ::DbWeight::get().reads(1), + "Second run should only read the migration flag" + ); + }); +} + #[test] fn test_migrate_crv3_v2_to_timelocked() { new_test_ext(1).execute_with(|| { diff --git a/pallets/subtensor/src/tests/mod.rs b/pallets/subtensor/src/tests/mod.rs index b9f4ff536..a0105a6ff 100644 --- a/pallets/subtensor/src/tests/mod.rs +++ b/pallets/subtensor/src/tests/mod.rs @@ -1,3 +1,4 @@ +mod auto_stake_hotkey; mod batch_tx; mod children; mod coinbase; diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index e597728a6..931653bfc 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -223,7 +223,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion { // `spec_version`, and `authoring_version` are the same between Wasm and native. // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use // the compatible custom types. - spec_version: 323, + spec_version: 324, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1, @@ -2403,6 +2403,10 @@ impl_runtime_apis! { pallet_subtensor::Pallet::::get_network_to_prune() } + fn get_coldkey_auto_stake_hotkey(coldkey: AccountId32, netuid: NetUid) -> Option { + SubtensorModule::get_coldkey_auto_stake_hotkey(coldkey, netuid) + } + fn get_selective_mechagraph(netuid: NetUid, mecid: MechId, metagraph_indexes: Vec) -> Option> { SubtensorModule::get_selective_mechagraph(netuid, mecid, metagraph_indexes) }