Auto merge of #3106 - braddmiller:3046-mergetoaddress-concurrent, r=str4d

Add Note Locking to z_mergetoaddress

Adds note locking to `z_mergetoaddress` allowing it to be invoked multiple times before previous `z_mergetoaddress` operations have finished.

Reference issue [#3046](https://github.com/zcash/zcash/issues/3046)

Co-authored-by: Eirik Ogilvie-Wigley <eirik@z.cash>

qa/rpc-tests/wallet_mergetoaddress.py

@@ -324,13 +324,27 @@ class WalletMergeToAddressTest (BitcoinTestFramework):
         self.sync_all()
 
         # Verify maximum number of notes which node 0 can shield can be set by the limit parameter
-        result = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0, 50, 2)
+        # Also check that we can set off a second merge before the first one is complete
-        assert_equal(result["mergingUTXOs"], Decimal('0'))
+
+        # myzaddr has 5 notes at this point
+        result1 = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0.0001, 50, 2)
+        result2 = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0.0001, 50, 2)
+
+        # First merge should select from all notes
+        assert_equal(result1["mergingUTXOs"], Decimal('0'))
         # Remaining UTXOs are only counted if we are trying to merge any UTXOs
-        assert_equal(result["remainingUTXOs"], Decimal('0'))
+        assert_equal(result1["remainingUTXOs"], Decimal('0'))
-        assert_equal(result["mergingNotes"], Decimal('2'))
+        assert_equal(result1["mergingNotes"], Decimal('2'))
-        assert_equal(result["remainingNotes"], Decimal('3'))
+        assert_equal(result1["remainingNotes"], Decimal('3'))
-        wait_and_assert_operationid_status(self.nodes[0], result['opid'])
+
+        # Second merge should ignore locked notes
+        assert_equal(result2["mergingUTXOs"], Decimal('0'))
+        assert_equal(result2["remainingUTXOs"], Decimal('0'))
+        assert_equal(result2["mergingNotes"], Decimal('2'))
+        assert_equal(result2["remainingNotes"], Decimal('1'))
+        wait_and_assert_operationid_status(self.nodes[0], result1['opid'])
+        wait_and_assert_operationid_status(self.nodes[0], result2['opid'])
+
         self.sync_all()
         self.nodes[1].generate(1)
         self.sync_all()
@@ -340,7 +354,7 @@ class WalletMergeToAddressTest (BitcoinTestFramework):
         assert_equal(result["mergingUTXOs"], Decimal('10'))
         assert_equal(result["remainingUTXOs"], Decimal('7'))
         assert_equal(result["mergingNotes"], Decimal('2'))
-        assert_equal(result["remainingNotes"], Decimal('2'))
+        assert_equal(result["remainingNotes"], Decimal('1'))
         wait_and_assert_operationid_status(self.nodes[0], result['opid'])
         # Don't sync node 2 which rejects the tx due to its mempooltxinputlimit
         sync_blocks(self.nodes[:2])

src/wallet/asyncrpcoperation_mergetoaddress.cpp

@@ -98,6 +98,7 @@ AsyncRPCOperation_mergetoaddress::AsyncRPCOperation_mergetoaddress(
 
     // Lock UTXOs
     lock_utxos();
+    lock_notes();
 
     // Enable payment disclosure if requested
     paymentDisclosureMode = fExperimentalMode && GetBoolArg("-paymentdisclosure", false);
@@ -111,6 +112,7 @@ void AsyncRPCOperation_mergetoaddress::main()
 {
     if (isCancelled()) {
         unlock_utxos(); // clean up
+        unlock_notes();
         return;
     }
 
@@ -173,6 +175,7 @@ void AsyncRPCOperation_mergetoaddress::main()
     LogPrintf("%s", s);
 
     unlock_utxos(); // clean up
+    unlock_notes(); // clean up
 
     // !!! Payment disclosure START
     if (success && paymentDisclosureMode && paymentDisclosureData_.size() > 0) {
@@ -927,3 +930,24 @@ void AsyncRPCOperation_mergetoaddress::unlock_utxos() {
         pwalletMain->UnlockCoin(std::get<0>(utxo));
     }
 }
+
+
+/**
+ * Lock input notes
+ */
+ void AsyncRPCOperation_mergetoaddress::lock_notes() {
+    LOCK2(cs_main, pwalletMain->cs_wallet);
+    for (auto note : noteInputs_) {
+        pwalletMain->LockNote(std::get<0>(note));
+    }
+}
+
+/**
+ * Unlock input notes
+ */
+void AsyncRPCOperation_mergetoaddress::unlock_notes() {
+    LOCK2(cs_main, pwalletMain->cs_wallet);
+    for (auto note : noteInputs_) {
+        pwalletMain->UnlockNote(std::get<0>(note));
+    }
+}

src/wallet/asyncrpcoperation_mergetoaddress.h

@@ -121,6 +121,10 @@ private:
 
     void unlock_utxos();
 
+    void lock_notes();
+
+    void unlock_notes();
+
     // payment disclosure!
     std::vector<PaymentDisclosureKeyInfo> paymentDisclosureData_;
 };

src/wallet/gtest/test_wallet.cpp

@@ -1046,3 +1046,36 @@ TEST(wallet_tests, MarkAffectedTransactionsDirty) {
     wallet.MarkAffectedTransactionsDirty(wtx2);
     EXPECT_FALSE(wallet.mapWallet[hash].fDebitCached);
 }
+
+TEST(wallet_tests, NoteLocking) {
+    TestWallet wallet;
+
+    auto sk = libzcash::SpendingKey::random();
+    wallet.AddSpendingKey(sk);
+
+    auto wtx = GetValidReceive(sk, 10, true);
+    auto wtx2 = GetValidReceive(sk, 10, true);
+
+    JSOutPoint jsoutpt {wtx.GetHash(), 0, 0};
+    JSOutPoint jsoutpt2 {wtx2.GetHash(),0, 0};
+
+    // Test selective locking
+    wallet.LockNote(jsoutpt);
+    EXPECT_TRUE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n));
+    EXPECT_FALSE(wallet.IsLockedNote(jsoutpt2.hash, jsoutpt2.js, jsoutpt2.n));
+
+    // Test selective unlocking
+    wallet.UnlockNote(jsoutpt);
+    EXPECT_FALSE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n));
+
+    // Test multiple locking
+    wallet.LockNote(jsoutpt);
+    wallet.LockNote(jsoutpt2);
+    EXPECT_TRUE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n));
+    EXPECT_TRUE(wallet.IsLockedNote(jsoutpt2.hash, jsoutpt2.js, jsoutpt2.n));
+
+    // Test unlock all
+    wallet.UnlockAllNotes();
+    EXPECT_FALSE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n));
+    EXPECT_FALSE(wallet.IsLockedNote(jsoutpt2.hash, jsoutpt2.js, jsoutpt2.n));
+}

src/wallet/wallet.cpp

@@ -3443,6 +3443,42 @@ void CWallet::ListLockedCoins(std::vector<COutPoint>& vOutpts)
     }
 }
 
+
+// Note Locking Operations
+
+void CWallet::LockNote(JSOutPoint& output)
+{
+    AssertLockHeld(cs_wallet); // setLockedNotes
+    setLockedNotes.insert(output);
+}
+
+void CWallet::UnlockNote(JSOutPoint& output)
+{
+    AssertLockHeld(cs_wallet); // setLockedNotes
+    setLockedNotes.erase(output);
+}
+
+void CWallet::UnlockAllNotes()
+{
+    AssertLockHeld(cs_wallet); // setLockedNotes
+    setLockedNotes.clear();
+}
+
+bool CWallet::IsLockedNote(uint256 hash, size_t js, uint8_t n) const
+{
+    AssertLockHeld(cs_wallet); // setLockedNotes
+    JSOutPoint outpt(hash, js, n);
+
+    return (setLockedNotes.count(outpt) > 0);
+}
+
+std::vector<JSOutPoint> CWallet::ListLockedNotes()
+{
+    AssertLockHeld(cs_wallet); // setLockedNotes
+    std::vector<JSOutPoint> vOutpts(setLockedNotes.begin(), setLockedNotes.end());
+    return vOutpts;
+}
+
 /** @} */ // end of Actions
 
 class CAffectedKeysVisitor : public boost::static_visitor<void> {
@@ -3732,6 +3768,11 @@ void CWallet::GetFilteredNotes(
                 continue;
             }
             
+            // skip locked notes
+            if (IsLockedNote(jsop.hash, jsop.js, jsop.n)) {
+                continue;
+            }
+
             int i = jsop.js; // Index into CTransaction.vjoinsplit
             int j = jsop.n; // Index into JSDescription.ciphertexts
 

src/wallet/wallet.h

@@ -883,6 +883,7 @@ public:
     CPubKey vchDefaultKey;
 
     std::set<COutPoint> setLockedCoins;
+    std::set<JSOutPoint> setLockedNotes;
 
     int64_t nTimeFirstKey;
 
@@ -903,6 +904,14 @@ public:
     void UnlockAllCoins();
     void ListLockedCoins(std::vector<COutPoint>& vOutpts);
 
+
+    bool IsLockedNote(uint256 hash, size_t js, uint8_t n) const;
+    void LockNote(JSOutPoint& output);
+    void UnlockNote(JSOutPoint& output);
+    void UnlockAllNotes();
+    std::vector<JSOutPoint> ListLockedNotes();
+
+
     /**
      * keystore implementation
      * Generate a new key