Schedule Coldkey Swap

pallets/subtensor/src/benchmarks.rs

@@ -428,4 +428,31 @@ 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;
+    let block_number: u64 = Subtensor::<T>::get_current_block_as_u64();
+    let nonce = 0;
+
+    // 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.clone(), vec![], block_number, nonce)
 }

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/events.rs

@@ -150,5 +150,14 @@ mod events {
                 <T as frame_system::Config>::AccountId,
             >>::Balance,
         },
+        /// A coldkey swap has been scheduled
+        ColdkeySwapScheduled {
+            /// The account ID of the old coldkey
+            old_coldkey: T::AccountId,
+            /// The account ID of the new coldkey
+            new_coldkey: T::AccountId,
+            /// The arbitration block for the coldkey swap
+            arbitration_block: u64,
+        },
     }
 }

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},
@@ -83,6 +84,9 @@ pub mod pallet {
     /// order of migrations. (i.e. always increasing)
     const STORAGE_VERSION: StorageVersion = StorageVersion::new(6);
 
+    /// Minimum balance required to perform a coldkey swap
+    pub const MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP: u64 = 1_000_000_000; // 1 TAO in RAO
+
     #[pallet::pallet]
     #[pallet::without_storage_info]
     #[pallet::storage_version(STORAGE_VERSION)]
@@ -366,6 +370,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,9 +387,35 @@ pub mod pallet {
         ValueQuery,
         DefaultAccountTake<T>,
     >;
-    #[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::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 {
@@ -1301,21 +1334,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)
                 }
             }
         }
@@ -2037,15 +2078,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 ------------------------------------------------------------
@@ -2278,10 +2323,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";
 
@@ -2301,6 +2348,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) {
@@ -2377,6 +2437,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/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!(