refactor rate limit check

Author: camfairchild

pallets/subtensor/src/staking/set_children.rs

@@ -61,7 +61,7 @@ impl<T: Config> Pallet<T> {
 
         // Set last transaction block
         let current_block = Self::get_current_block_as_u64();
-        Self::set_last_transaction_block(
+        Self::set_last_transaction_block_on_subnet(
             &hotkey,
             netuid,
             &TransactionType::SetChildren,
@@ -325,7 +325,7 @@ impl<T: Config> Pallet<T> {
 
         // Set last transaction block
         let current_block = Self::get_current_block_as_u64();
-        Self::set_last_transaction_block(
+        Self::set_last_transaction_block_on_subnet(
             &hotkey,
             netuid,
             &TransactionType::SetChildkeyTake,
@@ -336,7 +336,7 @@ impl<T: Config> Pallet<T> {
         ChildkeyTake::<T>::insert(hotkey.clone(), netuid, take);
 
         // Update the last transaction block
-        Self::set_last_transaction_block(
+        Self::set_last_transaction_block_on_subnet(
             &hotkey,
             netuid,
             &TransactionType::SetChildkeyTake,

pallets/subtensor/src/tests/children.rs

@@ -937,7 +937,7 @@ fn test_childkey_take_rate_limiting() {
         // Helper function to log rate limit information
         let log_rate_limit_info = || {
             let current_block = SubtensorModule::get_current_block_as_u64();
-            let last_block = SubtensorModule::get_last_transaction_block(
+            let last_block = SubtensorModule::get_last_transaction_block_on_subnet(
                 &hotkey,
                 netuid,
                 &TransactionType::SetChildkeyTake,
@@ -947,7 +947,7 @@ fn test_childkey_take_rate_limiting() {
                 &hotkey,
                 netuid,
             );
-            let limit = SubtensorModule::get_rate_limit(&TransactionType::SetChildkeyTake, 0);
+            let limit = SubtensorModule::get_rate_limit_on_subnet(&TransactionType::SetChildkeyTake, netuid);
             log::info!(
                 "Rate limit info: current_block: {}, last_block: {}, limit: {}, passes: {}, diff: {}",
                 current_block,
@@ -3609,7 +3609,8 @@ fn test_set_children_rate_limit_fail_then_succeed() {
 
         // Try again after rate limit period has passed
         // Check rate limit
-        let limit = SubtensorModule::get_rate_limit(&TransactionType::SetChildren, netuid);
+        let limit =
+            SubtensorModule::get_rate_limit_on_subnet(&TransactionType::SetChildren, netuid);
 
         // Step that many blocks
         step_block(limit as u16);

pallets/subtensor/src/tests/mock.rs

@@ -709,7 +709,7 @@ pub fn mock_set_children(coldkey: &U256, parent: &U256, netuid: u16, child_vec:
 #[allow(dead_code)]
 pub fn step_rate_limit(transaction_type: &TransactionType, netuid: u16) {
     // Check rate limit
-    let limit = SubtensorModule::get_rate_limit(transaction_type, netuid);
+    let limit = SubtensorModule::get_rate_limit_on_subnet(transaction_type, netuid);
 
     // Step that many blocks
     step_block(limit as u16);

pallets/subtensor/src/utils/rate_limiting.rs

@@ -34,12 +34,33 @@ impl<T: Config> Pallet<T> {
     // ==== Rate Limiting =====
     // ========================
     /// Get the rate limit for a specific transaction type
-    pub fn get_rate_limit(tx_type: &TransactionType, _netuid: u16) -> u64 {
+    pub fn get_rate_limit(tx_type: &TransactionType) -> u64 {
         match tx_type {
             TransactionType::SetChildren => 7200, // Cannot set children twice within a day
             TransactionType::SetChildkeyTake => TxChildkeyTakeRateLimit::<T>::get(),
             TransactionType::Unknown => 0, // Default to no limit for unknown types (no limit)
         }
+        }
+    }
+
+    pub fn get_rate_limit_on_subnet(tx_type: &TransactionType, _netuid: u16) -> u64 {
+        #[allow(clippy::match_single_binding)]
+        match tx_type {
+            _ => Self::get_rate_limit(tx_type),
+        }
+    }
+
+    pub fn check_passes_rate_limit(limit: u64, block: u64, last_block: u64) -> bool {
+        // Allow the first transaction (when last_block is 0) or if the rate limit has passed
+        last_block == 0 || block.saturating_sub(last_block) >= limit
+    }
+
+    pub fn passes_rate_limit(tx_type: &TransactionType, key: &T::AccountId) -> bool {
+        let block: u64 = Self::get_current_block_as_u64();
+        let limit: u64 = Self::get_rate_limit(tx_type);
+        let last_block: u64 = Self::get_last_transaction_block(key, tx_type);
+
+        Self::check_passes_rate_limit(limit, block, last_block)
     }
 
     /// Check if a transaction should be rate limited on a specific subnet
@@ -49,15 +70,19 @@ impl<T: Config> Pallet<T> {
         netuid: u16,
     ) -> bool {
         let block: u64 = Self::get_current_block_as_u64();
-        let limit: u64 = Self::get_rate_limit(tx_type, netuid);
-        let last_block: u64 = Self::get_last_transaction_block(hotkey, netuid, tx_type);
+        let limit: u64 = Self::get_rate_limit_on_subnet(tx_type, netuid);
+        let last_block: u64 = Self::get_last_transaction_block_on_subnet(hotkey, netuid, tx_type);
 
-        // Allow the first transaction (when last_block is 0) or if the rate limit has passed
-        last_block == 0 || block.saturating_sub(last_block) >= limit
+        Self::check_passes_rate_limit(limit, block, last_block)
+    }
+
+    /// Get the block number of the last transaction for a specific key, and transaction type
+    pub fn get_last_transaction_block(key: &T::AccountId, tx_type: &TransactionType) -> u64 {
+        Self::get_last_transaction_block_on_subnet(key, 0, tx_type)
     }
 
     /// Get the block number of the last transaction for a specific hotkey, network, and transaction type
-    pub fn get_last_transaction_block(
+    pub fn get_last_transaction_block_on_subnet(
         hotkey: &T::AccountId,
         netuid: u16,
         tx_type: &TransactionType,
@@ -66,15 +91,20 @@ impl<T: Config> Pallet<T> {
         TransactionKeyLastBlock::<T>::get((hotkey, netuid, tx_as_u16))
     }
 
+    /// Set the block number of the last transaction for a specific key, and transaction type
+    pub fn set_last_transaction_block(key: &T::AccountId, tx_type: &TransactionType, block: u64) {
+        Self::set_last_transaction_block_on_subnet(key, 0, tx_type, block)
+    }
+
     /// Set the block number of the last transaction for a specific hotkey, network, and transaction type
-    pub fn set_last_transaction_block(
-        hotkey: &T::AccountId,
+    pub fn set_last_transaction_block_on_subnet(
+        key: &T::AccountId,
         netuid: u16,
         tx_type: &TransactionType,
         block: u64,
     ) {
-        let tx_as_u16: u16 = (*tx_type).into();
-        TransactionKeyLastBlock::<T>::insert((hotkey, netuid, tx_as_u16), block);
+		let tx_as_u16: u16 = (*tx_type).into();
+        TransactionKeyLastBlock::<T>::insert((key, netuid, tx_as_u16), block);
     }
 
     pub fn set_last_tx_block(key: &T::AccountId, block: u64) {