Introduce stake jobs batching and delayed execution

pallets/subtensor/src/benchmarks.rs

@@ -135,48 +135,138 @@ benchmarks! {
     assert_ok!(Subtensor::<T>::do_burned_registration(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey.clone()));
   }: add_stake(RawOrigin::Signed( coldkey.clone() ), hotkey, netuid, amount)
 
-  benchmark_remove_stake{
-    let caller: T::AccountId = whitelisted_caller::<T::AccountId>();
+  benchmark_add_stake_aggregate {
+    let caller: T::AccountId = whitelisted_caller::<AccountIdOf<T>>();
+    let caller_origin = <T as frame_system::Config>::RuntimeOrigin::from(RawOrigin::Signed(caller.clone()));
+    let netuid: u16 = 1;
+    let version_key: u64 = 1;
+    let tempo: u16 = 1;
+    let modality: u16 = 0;
+    let seed : u32 = 1;
+
+    Subtensor::<T>::init_new_network(netuid, tempo);
+
+    Subtensor::<T>::set_burn(netuid, 1);
+    Subtensor::<T>::set_network_registration_allowed( netuid, true );
+
+    Subtensor::<T>::set_max_allowed_uids( netuid, 4096 );
+    assert_eq!(Subtensor::<T>::get_max_allowed_uids(netuid), 4096);
+
+    let coldkey: T::AccountId = account("Test", 0, seed);
+    let hotkey: T::AccountId = account("Alice", 0, seed);
+
+    let amount: u64 = 600000;
+    let amount_to_be_staked = 1000000000u64;
+    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), amount_to_be_staked);
+
+    assert_ok!(Subtensor::<T>::do_burned_registration(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey.clone()));
+  }: add_stake_aggregate(RawOrigin::Signed( coldkey.clone() ), hotkey, netuid, amount)
+
+  benchmark_remove_stake_limit_aggregate{
+    let caller: T::AccountId = whitelisted_caller::<AccountIdOf<T>>();
+    let caller_origin = <T as frame_system::Config>::RuntimeOrigin::from(RawOrigin::Signed(caller.clone()));
     let netuid: u16 = 1;
     let tempo: u16 = 1;
+    let modality: u16 = 0;
     let seed : u32 = 1;
 
+    // Set our total stake to 1000 TAO
     Subtensor::<T>::increase_total_stake(1_000_000_000_000);
+
     Subtensor::<T>::init_new_network(netuid, tempo);
-    SubtokenEnabled::<T>::insert(netuid, true);
-    Subtensor::<T>::set_network_registration_allowed(netuid, true);
-    Subtensor::<T>::set_max_allowed_uids(netuid, 4096);
+    Subtensor::<T>::set_network_registration_allowed( netuid, true );
+
+    Subtensor::<T>::set_max_allowed_uids( netuid, 4096 );
+    assert_eq!(Subtensor::<T>::get_max_allowed_uids(netuid), 4096);
+
+    let coldkey: T::AccountId = account("Test", 0, seed);
+    let hotkey: T::AccountId = account("Alice", 0, seed);
+    Subtensor::<T>::set_burn(netuid, 1);
+
+    let limit: u64 = 1_000_000_000;
+    let tao_reserve = 150_000_000_000_u64;
+    let alpha_in = 100_000_000_000_u64;
+    SubnetTAO::<T>::insert(netuid, tao_reserve);
+    SubnetAlphaIn::<T>::insert(netuid, alpha_in);
+
+    let wallet_bal = 1000000u32.into();
+    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), wallet_bal);
+
+    assert_ok!(Subtensor::<T>::do_burned_registration(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey.clone()));
+
+    let u64_staked_amt = 100_000_000_000;
+    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), u64_staked_amt);
+
+    assert_ok!(Subtensor::<T>::add_stake(RawOrigin::Signed( coldkey.clone() ).into() , hotkey.clone(), netuid, u64_staked_amt));
+
+    let amount_unstaked: u64 = 30_000_000_000;
+  }: remove_stake_limit_aggregate(RawOrigin::Signed( coldkey.clone() ), hotkey.clone(), netuid, amount_unstaked, limit, false)
+
+  benchmark_remove_stake_aggregate{
+    let caller: T::AccountId = whitelisted_caller::<AccountIdOf<T>>();
+    let caller_origin = <T as frame_system::Config>::RuntimeOrigin::from(RawOrigin::Signed(caller.clone()));
+    let netuid: u16 = 1;
+    let version_key: u64 = 1;
+    let tempo: u16 = 1;
+    let modality: u16 = 0;
+    let seed : u32 = 1;
+
+    // Set our total stake to 1000 TAO
+    Subtensor::<T>::increase_total_stake(1_000_000_000_000);
+
+    Subtensor::<T>::init_new_network(netuid, tempo);
+    Subtensor::<T>::set_network_registration_allowed( netuid, true );
+
+    Subtensor::<T>::set_max_allowed_uids( netuid, 4096 );
     assert_eq!(Subtensor::<T>::get_max_allowed_uids(netuid), 4096);
 
     let coldkey: T::AccountId = account("Test", 0, seed);
     let hotkey: T::AccountId = account("Alice", 0, seed);
+      Subtensor::<T>::set_burn(netuid, 1);
+
+    let wallet_bal = 1000000u32.into();
+    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), wallet_bal);
+
+    assert_ok!(Subtensor::<T>::do_burned_registration(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey.clone()));
 
+    // Stake 10% of our current total staked TAO
+    let u64_staked_amt = 100_000_000_000;
+    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), u64_staked_amt);
+
+    assert_ok!( Subtensor::<T>::add_stake(RawOrigin::Signed( coldkey.clone() ).into() , hotkey.clone(), netuid, u64_staked_amt));
+
+    let amount_unstaked: u64 = 600000;
+  }: remove_stake_aggregate(RawOrigin::Signed( coldkey.clone() ), hotkey.clone(), netuid, amount_unstaked)
+
+  benchmark_add_stake_limit_aggregate {
+    let caller: T::AccountId = whitelisted_caller::<AccountIdOf<T>>();
+    let caller_origin = <T as frame_system::Config>::RuntimeOrigin::from(RawOrigin::Signed(caller.clone()));
+    let netuid: u16 = 1;
+    let tempo: u16 = 1;
+    let modality: u16 = 0;
+    let seed : u32 = 1;
+
+    Subtensor::<T>::init_new_network(netuid, tempo);
     Subtensor::<T>::set_burn(netuid, 1);
-    let wallet_bal = 9_999_999_999_999u64.into();
-    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey, wallet_bal);
+    Subtensor::<T>::set_network_registration_allowed( netuid, true );
+    Subtensor::<T>::set_max_allowed_uids( netuid, 4096 );
 
-    assert_ok!(Subtensor::<T>::do_burned_registration(
-        RawOrigin::Signed(coldkey.clone()).into(),
-        netuid,
-        hotkey.clone()
-    ));
+    let coldkey: T::AccountId = account("Test", 0, seed);
+    let hotkey: T::AccountId = account("Alice", 0, seed);
 
-    let tao_to_stake = 100_000_000_000u64;
-    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey, tao_to_stake);
-    assert_ok!( Subtensor::<T>::add_stake(
-        RawOrigin::Signed(coldkey.clone()).into(),
-        hotkey.clone(),
-        netuid,
-        tao_to_stake
-    ));
+    let amount = 900_000_000_000;
+    let limit: u64 = 6_000_000_000;
+    let amount_to_be_staked = 440_000_000_000;
+    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), amount);
 
-    let actual_alpha_minted: u64 = Subtensor::<T>::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid);
-    assert!(actual_alpha_minted > 0, "No alpha minted after add_stake");
+    let tao_reserve = 150_000_000_000_u64;
+    let alpha_in = 100_000_000_000_u64;
+    SubnetTAO::<T>::insert(netuid, tao_reserve);
+    SubnetAlphaIn::<T>::insert(netuid, alpha_in);
 
-    SubnetAlphaOut::<T>::insert(netuid, actual_alpha_minted * 2);
+    assert_ok!(Subtensor::<T>::do_burned_registration(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey.clone()));
+  }: add_stake_limit_aggregate(RawOrigin::Signed( coldkey.clone() ), hotkey, netuid, amount_to_be_staked, limit, false)
 
-    let amount_unstaked: u64 = actual_alpha_minted / 2;
-  }: remove_stake(RawOrigin::Signed( coldkey.clone() ), hotkey.clone(), netuid, amount_unstaked)
 
   benchmark_serve_axon{
     let caller: T::AccountId = whitelisted_caller::<AccountIdOf<T>>();

pallets/subtensor/src/lib.rs

@@ -268,6 +268,80 @@ pub mod pallet {
         /// Additional information about the subnet
         pub additional: Vec<u8>,
     }
+
+    /// Data structure for stake related jobs.
+    #[derive(Encode, Decode, TypeInfo, Clone, PartialEq, Eq, Debug)]
+    pub enum StakeJob<AccountId> {
+        /// Represents a job for "add_stake" operation
+        AddStake {
+            /// Hotkey account
+            hotkey: AccountId,
+            /// Coldkey account
+            coldkey: AccountId,
+            /// Subnet ID
+            netuid: u16,
+            /// The amount of stake to be added to the hotkey staking account.
+            stake_to_be_added: u64,
+        },
+        /// Represents a job for "remove_stake" operation
+        RemoveStake {
+            /// Hotkey account
+            hotkey: AccountId,
+            /// Coldkey account
+            coldkey: AccountId,
+            /// Subnet ID
+            netuid: u16,
+            /// Alpha value
+            alpha_unstaked: u64,
+        },
+        /// Represents a job for "add_stake_limit" operation
+        AddStakeLimit {
+            /// Coldkey account
+            coldkey: AccountId,
+            /// Hotkey account
+            hotkey: AccountId,
+            /// Subnet ID
+            netuid: u16,
+            /// The amount of stake to be added to the hotkey staking account.
+            stake_to_be_added: u64,
+            /// The limit price expressed in units of RAO per one Alpha.
+            limit_price: u64,
+            /// Allows partial execution of the amount. If set to false, this becomes
+            /// fill or kill type or order.
+            allow_partial: bool,
+        },
+        /// Represents a job for "remove_stake_limit" operation
+        RemoveStakeLimit {
+            /// Coldkey account
+            coldkey: AccountId,
+            /// Hotkey account
+            hotkey: AccountId,
+            /// Subnet ID
+            netuid: u16,
+            /// The amount of stake to be added to the hotkey staking account.
+            alpha_unstaked: u64,
+            /// The limit price
+            limit_price: u64,
+            /// Allows partial execution of the amount. If set to false, this becomes
+            /// fill or kill type or order.
+            allow_partial: bool,
+        },
+        /// Represents a job for "unstake_all" operation
+        UnstakeAll {
+            /// Coldkey account
+            coldkey: AccountId,
+            /// Hotkey account
+            hotkey: AccountId,
+        },
+        /// Represents a job for "unstake_all_alpha" operation
+        UnstakeAllAlpha {
+            /// Coldkey account
+            coldkey: AccountId,
+            /// Hotkey account
+            hotkey: AccountId,
+        },
+    }
+
     /// ============================
     /// ==== Staking + Accounts ====
     /// ============================
@@ -816,6 +890,21 @@ pub mod pallet {
     pub type SenateRequiredStakePercentage<T> =
         StorageValue<_, u64, ValueQuery, DefaultSenateRequiredStakePercentage<T>>;
 
+    #[pallet::storage]
+    pub type StakeJobs<T: Config> = StorageDoubleMap<
+        _,
+        Blake2_128Concat,
+        BlockNumberFor<T>, // first key: current block number
+        Twox64Concat,
+        u64, // second key: unique job ID
+        StakeJob<T::AccountId>,
+        OptionQuery,
+    >;
+
+    #[pallet::storage]
+    /// Ensures unique IDs for StakeJobs storage map
+    pub type NextStakeJobId<T> = StorageValue<_, u64, ValueQuery, DefaultZeroU64<T>>;
+
     /// ============================
     /// ==== Staking Variables ====
     /// ============================
@@ -2057,6 +2146,101 @@ where
                     Self::get_priority_staking(who, hotkey, *amount_unstaked),
                 )
             }
+            Some(Call::add_stake_aggregate {
+                hotkey,
+                netuid,
+                amount_staked,
+            }) => {
+                if ColdkeySwapScheduled::<T>::contains_key(who) {
+                    return InvalidTransaction::Custom(
+                        CustomTransactionError::ColdkeyInSwapSchedule.into(),
+                    )
+                    .into();
+                }
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_add_stake(
+                        who,
+                        hotkey,
+                        *netuid,
+                        *amount_staked,
+                        *amount_staked,
+                        false,
+                    ),
+                    Self::get_priority_staking(who, hotkey, *amount_staked),
+                )
+            }
+            Some(Call::add_stake_limit_aggregate {
+                hotkey,
+                netuid,
+                amount_staked,
+                limit_price,
+                allow_partial,
+            }) => {
+                if ColdkeySwapScheduled::<T>::contains_key(who) {
+                    return InvalidTransaction::Custom(
+                        CustomTransactionError::ColdkeyInSwapSchedule.into(),
+                    )
+                    .into();
+                }
+
+                //Calculate the maximum amount that can be executed with price limit
+                let max_amount = Pallet::<T>::get_max_amount_add(*netuid, *limit_price);
+
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_add_stake(
+                        who,
+                        hotkey,
+                        *netuid,
+                        *amount_staked,
+                        max_amount,
+                        *allow_partial,
+                    ),
+                    Self::get_priority_staking(who, hotkey, *amount_staked),
+                )
+            }
+            Some(Call::remove_stake_aggregate {
+                hotkey,
+                netuid,
+                amount_unstaked,
+            }) => {
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_remove_stake(
+                        who,
+                        hotkey,
+                        *netuid,
+                        *amount_unstaked,
+                        *amount_unstaked,
+                        false,
+                    ),
+                    Self::get_priority_staking(who, hotkey, *amount_unstaked),
+                )
+            }
+            Some(Call::remove_stake_limit_aggregate {
+                hotkey,
+                netuid,
+                amount_unstaked,
+                limit_price,
+                allow_partial,
+            }) => {
+                // Calculate the maximum amount that can be executed with price limit
+                let max_amount = Pallet::<T>::get_max_amount_remove(*netuid, *limit_price);
+
+                // Fully validate the user input
+                Self::result_to_validity(
+                    Pallet::<T>::validate_remove_stake(
+                        who,
+                        hotkey,
+                        *netuid,
+                        *amount_unstaked,
+                        max_amount,
+                        *allow_partial,
+                    ),
+                    Self::get_priority_staking(who, hotkey, *amount_unstaked),
+                )
+            }
             Some(Call::move_stake {
                 origin_hotkey,
                 destination_hotkey,