wallet: extract ProTx UTXO locking into LockProTxCoins

UdjinM6 · Jan 29, 2025 · 0efc9bb · 0efc9bb
1 parent 4e81799
commit 0efc9bb
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 15 deletions.
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
@@ -914,19 +914,15 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const CWalletTx::Confirmatio
     CWalletTx& wtx = (*ret.first).second;
     bool fInsertedNew = ret.second;
     bool fUpdated = update_wtx && update_wtx(wtx, fInsertedNew);
+    std::set<COutPoint> candidates;
     if (fInsertedNew) {
         wtx.m_confirm = confirm;
         wtx.nTimeReceived = GetTime();
         wtx.nOrderPos = IncOrderPosNext(&batch);
         wtx.m_it_wtxOrdered = wtxOrdered.insert(std::make_pair(wtx.nOrderPos, &wtx));
         wtx.nTimeSmart = ComputeTimeSmart(wtx);
         AddToSpends(hash, &batch);
-
-        // TODO: refactor duplicated code between CWallet::AddToWallet and CWallet::AutoLockMasternodeCollaterals
-        auto candidates{AddWalletUTXOs(wtx.tx, /*ret_dups=*/true)};
-        for (const auto& utxo : ListProTxCoins(candidates)) {
-            LockCoin(utxo, &batch);
-        }
+        candidates = AddWalletUTXOs(wtx.tx, /*ret_dups=*/true);
     }
 
     if (!fInsertedNew)
@@ -942,15 +938,12 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const CWalletTx::Confirmatio
             assert(wtx.m_confirm.hashBlock == confirm.hashBlock);
             assert(wtx.m_confirm.block_height == confirm.block_height);
         }
-
-        // TODO: refactor duplicated code with case fInstertedNew
-        auto candidates{AddWalletUTXOs(wtx.tx, /*ret_dups=*/false)};
+        candidates = AddWalletUTXOs(wtx.tx, /*ret_dups=*/false);
         if (!candidates.empty()) fUpdated = true;
-        for (const auto& utxo : ListProTxCoins(candidates)) {
-            LockCoin(utxo, &batch);
-        }
     }
 
+    LockProTxCoins(candidates, &batch);
+
     //// debug print
     WalletLogPrintf("AddToWallet %s  %s%s\n", hash.ToString(), (fInsertedNew ? "new" : ""), (fUpdated ? "update" : ""));
 
@@ -4058,9 +4051,7 @@ void CWallet::AutoLockMasternodeCollaterals()
         candidates.insert(tx_utxos.begin(), tx_utxos.end());
     }
     WalletBatch batch(GetDatabase());
-    for (const auto& utxo : ListProTxCoins(candidates)) {
-        LockCoin(utxo, &batch);
-    }
+    LockProTxCoins(candidates, &batch);
 }
 
 DBErrors CWallet::ZapSelectTx(std::vector<uint256>& vHashIn, std::vector<uint256>& vHashOut)
@@ -4516,6 +4507,14 @@ std::vector<COutPoint> CWallet::ListProTxCoins(const std::set<COutPoint>& utxos)
     return m_chain->listMNCollaterials(candidates);
 }
 
+void CWallet::LockProTxCoins(const std::set<COutPoint>& utxos, WalletBatch* batch)
+{
+    AssertLockHeld(cs_wallet);
+    for (const auto& utxo : ListProTxCoins(utxos)) {
+        LockCoin(utxo, batch);
+    }
+}
+
 /** @} */ // end of Actions
 
 void CWallet::GetKeyBirthTimes(std::map<CKeyID, int64_t>& mapKeyBirth) const {

diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -1047,6 +1047,7 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
     std::vector<COutPoint> ListLockedCoins() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     std::vector<COutPoint> ListProTxCoins() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     std::vector<COutPoint> ListProTxCoins(const std::set<COutPoint>& utxos) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    void LockProTxCoins(const std::set<COutPoint>& utxos, WalletBatch* batch = nullptr) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     /*
      * Rescan abort properties