Schedule Coldkey Swap

pallets/subtensor/src/benchmarks.rs

@@ -428,4 +428,29 @@ reveal_weights {
     let _ = Subtensor::<T>::commit_weights(<T as frame_system::Config>::RuntimeOrigin::from(RawOrigin::Signed(hotkey.clone())), netuid, commit_hash);
 
   }: reveal_weights(RawOrigin::Signed(hotkey.clone()), netuid, uids, weight_values, salt, version_key)
+
+  schedule_coldkey_swap {
+    let seed: u32 = 1;
+    let old_coldkey: T::AccountId = account("OldColdkey", 0, seed);
+    let new_coldkey: T::AccountId = account("NewColdkey", 0, seed + 1);
+    let hotkey: T::AccountId = account("Hotkey", 0, seed);
+
+    let netuid = 1u16;
+    let tempo = 1u16;
+
+    // Initialize the network
+    Subtensor::<T>::init_new_network(netuid, tempo);
+    Subtensor::<T>::set_network_registration_allowed(netuid, true);
+
+    // Add balance to the old coldkey account
+    let amount_to_be_staked: u64 = 1000000u32.into();
+    Subtensor::<T>::add_balance_to_coldkey_account(&old_coldkey.clone(), amount_to_be_staked+1000000000);
+       // Burned register the hotkey with the old coldkey
+       assert_ok!(Subtensor::<T>::burned_register(
+        RawOrigin::Signed(old_coldkey.clone()).into(),
+        netuid,
+        hotkey.clone()
+    ));
+
+  }: schedule_coldkey_swap(RawOrigin::Signed(old_coldkey.clone()), new_coldkey)
 }

pallets/subtensor/src/errors.rs

@@ -146,5 +146,13 @@ mod errors {
         NoBalanceToTransfer,
         /// Same coldkey
         SameColdkey,
+        /// The coldkey is in arbitration
+        ColdkeyIsInArbitration,
+        /// The new coldkey is already registered for the drain
+        DuplicateColdkey,
+        /// Error thrown on a coldkey swap.
+        SwapError,
+        /// Insufficient Balance to Schedule coldkey swap
+        InsufficientBalanceToPerformColdkeySwap,
     }
 }

pallets/subtensor/src/lib.rs

@@ -18,6 +18,7 @@ use frame_support::{
 use codec::{Decode, Encode};
 use frame_support::sp_runtime::transaction_validity::InvalidTransaction;
 use frame_support::sp_runtime::transaction_validity::ValidTransaction;
+use pallet_balances::Call as BalancesCall;
 use scale_info::TypeInfo;
 use sp_runtime::{
     traits::{DispatchInfoOf, Dispatchable, PostDispatchInfoOf, SignedExtension},
@@ -366,6 +367,9 @@ pub mod pallet {
     #[pallet::storage] // --- MAP ( cold ) --> Vec<hot> | Returns the vector of hotkeys controlled by this coldkey.
     pub type OwnedHotkeys<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
+    #[pallet::storage] // --- DMAP ( cold ) --> Vec<hot> | Maps coldkey to hotkeys that stake to it
+    pub type StakingHotkeys<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
     #[pallet::storage] // --- MAP ( hot ) --> take | Returns the hotkey delegation take. And signals that this key is open for delegation.
     pub type Delegates<T: Config> =
         StorageMap<_, Blake2_128Concat, T::AccountId, u16, ValueQuery, DefaultDefaultTake<T>>;
@@ -380,6 +384,36 @@ pub mod pallet {
         ValueQuery,
         DefaultAccountTake<T>,
     >;
+
+    #[pallet::type_value]
+    /// Default value for hotkeys.
+    pub fn EmptyAccounts<T: Config>() -> Vec<T::AccountId> {
+        vec![]
+    }
+    #[pallet::type_value]
+    /// Default stake interval.
+    pub fn DefaultArbitrationPeriod<T: Config>() -> u64 {
+        7200 * 3
+    }
+    #[pallet::storage] // ---- StorageItem Global Used Work.
+    pub type ArbitrationPeriod<T: Config> =
+        StorageValue<_, u64, ValueQuery, DefaultArbitrationPeriod<T>>;
+    #[pallet::storage] // --- MAP ( cold ) --> Vec<wallet_to_drain_to> | Returns a list of keys to drain to, if there are two, we extend the period.
+    pub type ColdkeySwapDestinations<T: Config> = StorageMap<
+        _,
+        Blake2_128Concat,
+        T::AccountId,
+        Vec<T::AccountId>,
+        ValueQuery,
+        EmptyAccounts<T>,
+    >;
+    #[pallet::storage] // --- MAP ( cold ) --> u64 | Block when the coldkey will be arbitrated.
+    pub type ColdkeyArbitrationBlock<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, u64, ValueQuery>;
+    #[pallet::storage] // --- MAP ( u64 ) --> Vec<coldkeys_to_drain>  | Coldkeys to drain on the specific block.
+    pub type ColdkeysToSwapAtBlock<T: Config> =
+        StorageMap<_, Identity, u64, Vec<T::AccountId>, ValueQuery, EmptyAccounts<T>>;
+
     /// -- ITEM (switches liquid alpha on)
     #[pallet::type_value]
     pub fn DefaultLiquidAlpha<T: Config>() -> bool {
@@ -1225,6 +1259,13 @@ pub mod pallet {
 
                     Stake::<T>::insert(hotkey.clone(), coldkey.clone(), stake);
 
+                    // Update StakingHotkeys map
+                    let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
+                    if !staking_hotkeys.contains(hotkey) {
+                        staking_hotkeys.push(hotkey.clone());
+                        StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
+                    }
+
                     next_uid = next_uid.checked_add(1).expect(
                         "should not have total number of hotkey accounts larger than u16::MAX",
                     );
@@ -1291,21 +1332,29 @@ pub mod pallet {
         // 	* 'n': (BlockNumberFor<T>):
         // 		- The number of the block we are initializing.
         fn on_initialize(_block_number: BlockNumberFor<T>) -> Weight {
+            // Unstake all and transfer pending coldkeys
+            let swap_weight = match Self::swap_coldkeys_this_block() {
+                Ok(weight) => weight,
+                Err(_) => Weight::from_parts(0, 0),
+            };
+
             let block_step_result = Self::block_step();
             match block_step_result {
                 Ok(_) => {
                     // --- If the block step was successful, return the weight.
-                    log::info!("Successfully ran block step.");
+                    log::debug!("Successfully ran block step.");
                     Weight::from_parts(110_634_229_000_u64, 0)
                         .saturating_add(T::DbWeight::get().reads(8304_u64))
                         .saturating_add(T::DbWeight::get().writes(110_u64))
+                        .saturating_add(swap_weight)
                 }
                 Err(e) => {
                     // --- If the block step was unsuccessful, return the weight anyway.
                     log::error!("Error while stepping block: {:?}", e);
                     Weight::from_parts(110_634_229_000_u64, 0)
                         .saturating_add(T::DbWeight::get().reads(8304_u64))
                         .saturating_add(T::DbWeight::get().writes(110_u64))
+                        .saturating_add(swap_weight)
                 }
             }
         }
@@ -1337,7 +1386,9 @@ pub mod pallet {
                 // Doesn't check storage version. TODO: Remove after upgrade
                 .saturating_add(migration::migration5_total_issuance::<T>(false))
                 // Populate OwnedHotkeys map for coldkey swap. Doesn't update storage vesion.
-                .saturating_add(migration::migrate_populate_owned::<T>());
+                .saturating_add(migration::migrate_populate_owned::<T>())
+                // Populate StakingHotkeys map for coldkey swap. Doesn't update storage vesion.
+                .saturating_add(migration::migrate_populate_staking_hotkeys::<T>());
 
             weight
         }
@@ -2025,15 +2076,19 @@ pub mod pallet {
         ///
         /// Weight is calculated based on the number of database reads and writes.
         #[pallet::call_index(72)]
-        #[pallet::weight((Weight::from_parts(1_940_000_000, 0)
-		.saturating_add(T::DbWeight::get().reads(272))
-		.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
-        pub fn unstake_all_and_transfer_to_new_coldkey(
+        #[pallet::weight((Weight::from_parts(21_000_000, 0)
+		.saturating_add(T::DbWeight::get().reads(3))
+		.saturating_add(T::DbWeight::get().writes(3)), DispatchClass::Operational, Pays::No))]
+        pub fn schedule_coldkey_swap(
             origin: OriginFor<T>,
             new_coldkey: T::AccountId,
+            work: Vec<u8>,
+            block_number: u64,
+            nonce: u64,
         ) -> DispatchResult {
-            let current_coldkey = ensure_signed(origin)?;
-            Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, new_coldkey)
+            // Attain the calling coldkey from the origin.
+            let old_coldkey: T::AccountId = ensure_signed(origin)?;
+            Self::do_schedule_coldkey_swap(&old_coldkey, &new_coldkey, work, block_number, nonce)
         }
 
         // ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------
@@ -2266,10 +2321,12 @@ impl<T: Config + Send + Sync + TypeInfo> sp_std::fmt::Debug for SubtensorSignedE
     }
 }
 
-impl<T: Config + Send + Sync + TypeInfo> SignedExtension for SubtensorSignedExtension<T>
+impl<T: Config + Send + Sync + TypeInfo + pallet_balances::Config> SignedExtension
+    for SubtensorSignedExtension<T>
 where
     T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
     <T as frame_system::Config>::RuntimeCall: IsSubType<Call<T>>,
+    <T as frame_system::Config>::RuntimeCall: IsSubType<BalancesCall<T>>,
 {
     const IDENTIFIER: &'static str = "SubtensorSignedExtension";
 
@@ -2289,6 +2346,19 @@ where
         _info: &DispatchInfoOf<Self::Call>,
         _len: usize,
     ) -> TransactionValidity {
+        // Check if the call is one of the balance transfer types we want to reject
+        if let Some(balances_call) = call.is_sub_type() {
+            match balances_call {
+                BalancesCall::transfer_allow_death { .. }
+                | BalancesCall::transfer_keep_alive { .. }
+                | BalancesCall::transfer_all { .. } => {
+                    if Pallet::<T>::coldkey_in_arbitration(who) {
+                        return Err(TransactionValidityError::Invalid(InvalidTransaction::Call));
+                    }
+                }
+                _ => {} // Other Balances calls are allowed
+            }
+        }
         match call.is_sub_type() {
             Some(Call::commit_weights { netuid, .. }) => {
                 if Self::check_weights_min_stake(who) {
@@ -2365,6 +2435,16 @@ where
                 priority: Self::get_priority_vanilla(),
                 ..Default::default()
             }),
+            Some(Call::dissolve_network { .. }) => {
+                if Pallet::<T>::coldkey_in_arbitration(who) {
+                    Err(TransactionValidityError::Invalid(InvalidTransaction::Call))
+                } else {
+                    Ok(ValidTransaction {
+                        priority: Self::get_priority_vanilla(),
+                        ..Default::default()
+                    })
+                }
+            }
             _ => Ok(ValidTransaction {
                 priority: Self::get_priority_vanilla(),
                 ..Default::default()

pallets/subtensor/src/migration.rs

@@ -539,3 +539,67 @@ pub fn migrate_populate_owned<T: Config>() -> Weight {
         Weight::zero()
     }
 }
+
+/// Populate the StakingHotkeys map from Stake map
+pub fn migrate_populate_staking_hotkeys<T: Config>() -> Weight {
+    // Setup migration weight
+    let mut weight = T::DbWeight::get().reads(1);
+    let migration_name = "Populate StakingHotkeys map";
+
+    // Check if this migration is needed (if StakingHotkeys map is empty)
+    let migrate = StakingHotkeys::<T>::iter().next().is_none();
+
+    // Only runs if the migration is needed
+    if migrate {
+        info!(target: LOG_TARGET_1, ">>> Starting Migration: {}", migration_name);
+
+        let mut longest_hotkey_vector: usize = 0;
+        let mut longest_coldkey: Option<T::AccountId> = None;
+        let mut keys_touched: u64 = 0;
+        let mut storage_reads: u64 = 0;
+        let mut storage_writes: u64 = 0;
+
+        // Iterate through all Owner entries
+        Stake::<T>::iter().for_each(|(hotkey, coldkey, stake)| {
+            storage_reads = storage_reads.saturating_add(1); // Read from Owner storage
+            if stake > 0 {
+                let mut hotkeys = StakingHotkeys::<T>::get(&coldkey);
+                storage_reads = storage_reads.saturating_add(1); // Read from StakingHotkeys storage
+
+                // Add the hotkey if it's not already in the vector
+                if !hotkeys.contains(&hotkey) {
+                    hotkeys.push(hotkey);
+                    keys_touched = keys_touched.saturating_add(1);
+
+                    // Update longest hotkey vector info
+                    if longest_hotkey_vector < hotkeys.len() {
+                        longest_hotkey_vector = hotkeys.len();
+                        longest_coldkey = Some(coldkey.clone());
+                    }
+
+                    // Update the StakingHotkeys storage
+                    StakingHotkeys::<T>::insert(&coldkey, hotkeys);
+                    storage_writes = storage_writes.saturating_add(1); // Write to StakingHotkeys storage
+                }
+
+                // Accrue weight for reads and writes
+                weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
+            }
+        });
+
+        // Log migration results
+        info!(
+            target: LOG_TARGET_1,
+            "Migration {} finished. Keys touched: {}, Longest hotkey vector: {}, Storage reads: {}, Storage writes: {}",
+            migration_name, keys_touched, longest_hotkey_vector, storage_reads, storage_writes
+        );
+        if let Some(c) = longest_coldkey {
+            info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {:?}", c);
+        }
+
+        weight
+    } else {
+        info!(target: LOG_TARGET_1, "Migration {} already done!", migration_name);
+        Weight::zero()
+    }
+}

pallets/subtensor/src/registration.rs

@@ -41,6 +41,10 @@ impl<T: Config> Pallet<T> {
     ) -> DispatchResult {
         // --- 1. Check that the caller has signed the transaction. (the coldkey of the pairing)
         let coldkey = ensure_signed(origin)?;
+        ensure!(
+            !Self::coldkey_in_arbitration(&coldkey),
+            Error::<T>::ColdkeyIsInArbitration
+        );
         log::info!(
             "do_registration( coldkey:{:?} netuid:{:?} hotkey:{:?} )",
             coldkey,

pallets/subtensor/src/root.rs

@@ -483,6 +483,10 @@ impl<T: Config> Pallet<T> {
 
         // --- 1. Ensure that the call originates from a signed source and retrieve the caller's account ID (coldkey).
         let coldkey = ensure_signed(origin)?;
+        ensure!(
+            !Self::coldkey_in_arbitration(&coldkey),
+            Error::<T>::ColdkeyIsInArbitration
+        );
         log::info!(
             "do_root_register( coldkey: {:?}, hotkey: {:?} )",
             coldkey,
@@ -623,6 +627,10 @@ impl<T: Config> Pallet<T> {
     ) -> dispatch::DispatchResult {
         // Check the caller's signature. This is the coldkey of a registered account.
         let coldkey = ensure_signed(origin)?;
+        ensure!(
+            !Self::coldkey_in_arbitration(&coldkey),
+            Error::<T>::ColdkeyIsInArbitration
+        );
         log::info!(
             "do_set_root_weights( origin:{:?} netuid:{:?}, uids:{:?}, values:{:?})",
             coldkey,
@@ -744,6 +752,10 @@ impl<T: Config> Pallet<T> {
     ) -> DispatchResultWithPostInfo {
         // --- 1. Ensure that the caller has signed with their coldkey.
         let coldkey = ensure_signed(origin.clone())?;
+        ensure!(
+            !Self::coldkey_in_arbitration(&coldkey),
+            Error::<T>::ColdkeyIsInArbitration
+        );
 
         // --- 2. Ensure that the calling coldkey owns the associated hotkey.
         ensure!(
@@ -797,6 +809,10 @@ impl<T: Config> Pallet<T> {
     pub fn user_add_network(origin: T::RuntimeOrigin) -> dispatch::DispatchResult {
         // --- 0. Ensure the caller is a signed user.
         let coldkey = ensure_signed(origin)?;
+        ensure!(
+            !Self::coldkey_in_arbitration(&coldkey),
+            Error::<T>::ColdkeyIsInArbitration
+        );
 
         // --- 1. Rate limit for network registrations.
         let current_block = Self::get_current_block_as_u64();
@@ -885,6 +901,10 @@ impl<T: Config> Pallet<T> {
     pub fn user_remove_network(origin: T::RuntimeOrigin, netuid: u16) -> dispatch::DispatchResult {
         // --- 1. Ensure the function caller is a signed user.
         let coldkey = ensure_signed(origin)?;
+        ensure!(
+            !Self::coldkey_in_arbitration(&coldkey),
+            Error::<T>::ColdkeyIsInArbitration
+        );
 
         // --- 2. Ensure this subnet exists.
         ensure!(