feat: rework contract lock (#332)

substrate-node/pallets/pallet-smart-contract/src/lib.rs

@@ -5,13 +5,15 @@ use frame_support::{
     dispatch::DispatchResultWithPostInfo,
     ensure,
     traits::{
-        Currency, ExistenceRequirement::KeepAlive, Get, LockableCurrency, Vec, WithdrawReasons,
+        Currency, ExistenceRequirement::KeepAlive, Get, LockIdentifier, LockableCurrency, Vec, WithdrawReasons,
     },
     weights::Pays,
 };
 use frame_system::{self as system, ensure_signed};
-use sp_runtime::{traits::SaturatedConversion, DispatchError, DispatchResult, Perbill, Percent};
-
+use sp_runtime::{
+    traits::{SaturatedConversion, CheckedSub}, 
+    DispatchError, DispatchResult, Perbill, Percent
+};
 use pallet_balances;
 use pallet_tfgrid;
 use pallet_tfgrid::types as pallet_tfgrid_types;
@@ -36,6 +38,8 @@ pub mod weights;
 
 pub use weights::WeightInfo;
 
+pub const GRID_LOCK_ID: LockIdentifier = *b"gridlock";
+
 pub trait Config:
     system::Config + pallet_tfgrid::Config + pallet_timestamp::Config + pallet_balances::Config
 {
@@ -107,7 +111,8 @@ decl_error! {
         NodeHasRentContract,
         NodeIsNotDedicated,
         NodeNotAvailableToDeploy,
-        CannotUpdateContractInGraceState
+        CannotUpdateContractInGraceState,
+        NumOverflow
     }
 }
 
@@ -694,7 +699,7 @@ impl<T: Config> Module<T> {
         let contract = Self::handle_grace(contract, usable_balance, amount_due)?;
 
         // Handle contract lock operations
-        Self::handle_lock(contract, amount_due);
+        Self::handle_lock(contract, amount_due)?;
 
         // Always emit a contract billed event
         let contract_bill = types::ContractBill {
@@ -771,30 +776,30 @@ impl<T: Config> Module<T> {
         Ok(contract)
     }
 
-    fn handle_lock(contract: &mut types::Contract, amount_due: BalanceOf<T>) {
+    fn handle_lock(contract: &mut types::Contract, amount_due: BalanceOf<T>) -> DispatchResult {
         let now = <timestamp::Module<T>>::get().saturated_into::<u64>() / 1000;
-        let mut contract_lock = ContractLock::<T>::get(contract.contract_id);
-        let new_amount_locked = contract_lock.amount_locked + amount_due;
 
         // increment cycles billed and update the internal lock struct
+        let mut contract_lock = ContractLock::<T>::get(contract.contract_id);
         contract_lock.lock_updated = now;
         contract_lock.cycles += 1;
-        contract_lock.amount_locked = new_amount_locked;
+        contract_lock.amount_locked = contract_lock.amount_locked + amount_due;
         ContractLock::<T>::insert(contract.contract_id, &contract_lock);
 
         // Contract is in grace state, don't actually lock tokens or distribute rewards
         if matches!(contract.state, types::ContractState::GracePeriod(_)) {
-            return;
+            return Ok(())
         }
 
-        let twin = pallet_tfgrid::Twins::<T>::get(contract.twin_id);
         // Only lock an amount from the user's balance if the contract is in create state
-        // Update lock for contract and ContractLock in storage
+        // The lock is specified on the user's account, since a user can have multiple contracts
+        // Just extend the lock with the amount due for this contract billing period (lock will be created if not exists)
+        let twin = pallet_tfgrid::Twins::<T>::get(contract.twin_id);
         <T as Config>::Currency::extend_lock(
-            contract.contract_id.to_be_bytes(),
+            GRID_LOCK_ID,
             &twin.account_id,
-            new_amount_locked.into(),
-            WithdrawReasons::RESERVE,
+            amount_due,
+            WithdrawReasons::all(),
         );
 
         let is_canceled = matches!(contract.state, types::ContractState::Deleted(_));
@@ -803,11 +808,28 @@ impl<T: Config> Module<T> {
         // When the cultivation rewards are ready to be distributed or it's in delete state
         // Unlock all reserved balance and distribute
         if contract_lock.cycles >= T::DistributionFrequency::get() || canceled_and_not_zero {
-            // Remove lock
+            // Deprecated locking system
+            // If there is a lock with ID being the contract ID, remove it
+            // Code can be removed in a later phase
             <T as Config>::Currency::remove_lock(
                 contract.contract_id.to_be_bytes(),
                 &twin.account_id,
             );
+
+            // First remove the lock, calculate how much locked balance needs to be unlocked and re-lock the remaining locked balance 
+            let locked_balance = Self::get_locked_balance(&twin.account_id);            
+            let new_locked_balance= match locked_balance.checked_sub(&contract_lock.amount_locked) {
+                Some(b) => b,
+                None => BalanceOf::<T>::saturated_from(0 as u128)
+            };
+            <T as Config>::Currency::remove_lock(GRID_LOCK_ID, &twin.account_id);
+            <T as Config>::Currency::set_lock(
+                GRID_LOCK_ID,
+                &twin.account_id,
+                new_locked_balance,
+                WithdrawReasons::all(),
+            );
+
             // Fetch the default pricing policy
             let pricing_policy = pallet_tfgrid::PricingPolicies::<T>::get(1);
             // Distribute cultivation rewards
@@ -825,6 +847,8 @@ impl<T: Config> Module<T> {
             contract_lock.cycles = 0;
             ContractLock::<T>::insert(contract.contract_id, &contract_lock);
         }
+
+        Ok(())
     }
 
     fn calculate_contract_cost_tft(
@@ -1323,6 +1347,17 @@ impl<T: Config> Module<T> {
         BalanceOf::<T>::saturated_from(b)
     }
 
+    fn get_locked_balance(account_id: &T::AccountId) -> BalanceOf<T> {
+        let usable_balance = Self::get_usable_balance(account_id);
+        let free_balance = <T as Config>::Currency::free_balance(account_id);
+
+        let locked_balance = free_balance.checked_sub(&usable_balance);
+        match locked_balance {
+            Some(balance) => balance,
+            None => BalanceOf::<T>::saturated_from(0 as u128)
+        }
+    }
+
     // cu1 = MAX(cru/2, mru/4)
     // cu2 = MAX(cru, mru/8)
     // cu3 = MAX(cru/4, mru/2)

substrate-node/pallets/pallet-smart-contract/src/tests.rs

@@ -1,7 +1,7 @@
 use crate::{mock::*, Error, RawEvent};
 use frame_support::{
     assert_noop, assert_ok,
-    traits::{OnFinalize, OnInitialize},
+    traits::{OnFinalize, OnInitialize, LockableCurrency, WithdrawReasons},
 };
 use frame_system::RawOrigin;
 use sp_runtime::{Perbill, Percent};
@@ -1741,6 +1741,25 @@ fn test_cu_calculation() {
     })
 }
 
+#[test]
+fn test_lock() {
+    new_test_ext().execute_with(|| {
+        let id: u64 = 1;
+        Balances::set_lock(
+            id.to_be_bytes(),
+            &bob(),
+            100,
+            WithdrawReasons::RESERVE,
+        );
+
+        let usable_balance = Balances::usable_balance(&bob());
+        let free_balance = Balances::free_balance(&bob());
+
+        let locked_balance = free_balance - usable_balance;
+        assert_eq!(locked_balance, 100);
+    })
+}
+
 #[test]
 fn test_percent() {
     let cost: u64 = 1000;