Segregate updating received utxos and spent utxos (#4093)

* Fix UTXOHandlingTest cases

* Reduce log level

* Add comments, fix bugs so that we can transition from TxoState.ConfirmedReceived -> TxoState.PendingConfirmationsReceived

* Fix BroadcastSpent/BroadcastReceive bugs

* Add scaladoc to updateUtxoStates()

* Break receive and spent tests into two test cases

* Remove unneeded database write

* Fix rebase

wallet-test/src/test/scala/org/bitcoins/wallet/UTXOHandlingTest.scala

@@ -3,6 +3,7 @@ package org.bitcoins.wallet
 import org.bitcoins.core.protocol.script.EmptyScriptPubKey
 import org.bitcoins.core.wallet.utxo.TxoState
 import org.bitcoins.core.wallet.utxo.TxoState._
+import org.bitcoins.crypto.DoubleSha256DigestBE
 import org.bitcoins.testkit.wallet.BitcoinSWalletTest
 import org.bitcoins.testkit.wallet.WalletTestUtil._
 import org.scalatest.FutureOutcome
@@ -17,46 +18,141 @@ class UTXOHandlingTest extends BitcoinSWalletTest {
     withNewWallet(test, getBIP39PasswordOpt())(getFreshWalletAppConfig)
   }
 
-  it must "correctly update txo state based on confirmations" in { wallet =>
-    val utxo = sampleSegwitUTXO(EmptyScriptPubKey,
-                                state = TxoState.Reserved
-    ) //state doesn't matter here
-    val requiredConfs = 6
-    assert(wallet.walletConfig.requiredConfirmations == requiredConfs)
+  val utxo = sampleSegwitUTXO(EmptyScriptPubKey,
+                              state = TxoState.Reserved
+  ) //state doesn't matter here
+  val requiredConfs = 6
+  val reserved = utxo.copyWithState(Reserved)
 
-    val immatureCoinbase = utxo.copyWithState(ImmatureCoinbase)
-    val pendingConfReceived = utxo.copyWithState(PendingConfirmationsReceived)
-    val spendingTxId = randomTXID
-    val pendingConfSpent = utxo
-      .copyWithSpendingTxId(spendingTxId)
-      .copyWithState(PendingConfirmationsSpent)
-    val confReceived = utxo
-      .copyWithState(ConfirmedReceived)
-    val confSpent = utxo
-      .copyWithSpendingTxId(spendingTxId)
-      .copyWithState(ConfirmedSpent)
-    val reserved = utxo.copyWithState(Reserved)
+  val immatureCoinbase = utxo.copyWithState(ImmatureCoinbase)
 
-    assert(wallet.updateTxoWithConfs(reserved, 1) == reserved)
+  val confReceived = utxo
+    .copyWithState(ConfirmedReceived)
 
-    assert(wallet.updateTxoWithConfs(immatureCoinbase, 10) == immatureCoinbase)
-    assert(wallet.updateTxoWithConfs(immatureCoinbase, 101) == confReceived)
+  val pendingConfReceived = utxo.copyWithState(PendingConfirmationsReceived)
 
-    assert(
-      wallet.updateTxoWithConfs(pendingConfReceived, 1) == pendingConfReceived)
-    assert(
-      wallet.updateTxoWithConfs(pendingConfReceived,
-                                requiredConfs) == confReceived)
+  it must "correct update receive txo state based on confirmations" in {
+    wallet =>
+      //it must stay reserved if we are receiving confirmations
+      assert(wallet.updateReceivedTxoWithConfs(reserved, 1) == reserved)
 
-    assert(wallet.updateTxoWithConfs(pendingConfSpent, 1) == pendingConfSpent)
-    assert(
-      wallet.updateTxoWithConfs(pendingConfSpent, requiredConfs) == confSpent)
+      //it must stay immature coinbase if we don't have > 101 confirmations
+      assert(
+        wallet.updateReceivedTxoWithConfs(immatureCoinbase,
+                                          10) == immatureCoinbase)
+      assert(
+        wallet.updateReceivedTxoWithConfs(immatureCoinbase,
+                                          101) == confReceived)
 
-    assert(wallet.updateTxoWithConfs(confSpent, 1) == confSpent)
-    assert(wallet.updateTxoWithConfs(confSpent, requiredConfs) == confSpent)
+      //we must stay pending confirmations received with only 1 confirmation
+      assert(
+        wallet.updateReceivedTxoWithConfs(pendingConfReceived,
+                                          1) == pendingConfReceived)
 
-    assert(wallet.updateTxoWithConfs(confReceived, 1) == confReceived)
-    assert(
-      wallet.updateTxoWithConfs(confReceived, requiredConfs) == confReceived)
+      //must be able to go back to broadcast received if we don't have confirmations (re-org scenario)
+      assert(
+        wallet
+          .updateReceivedTxoWithConfs(pendingConfReceived, 0)
+          .state == TxoState.BroadcastReceived)
+
+      //transition from TxoState.ConfirmedReceived -> TxoState.PendingConfirmationsReceived (reorg scenario)
+      assert(
+        wallet.updateReceivedTxoWithConfs(confReceived, 1) == confReceived
+          .copyWithState(PendingConfirmationsReceived))
+
+      //it must stay TxoState.ConfirmedReceived if we keep receiving confirmations
+      assert(
+        wallet.updateReceivedTxoWithConfs(confReceived,
+                                          requiredConfs) == confReceived)
+
+      //it must stay broadcast received with 0 confirmations
+      val broadcastReceived = utxo
+        .copyWithState(TxoState.BroadcastReceived)
+      assert(
+        wallet.updateReceivedTxoWithConfs(broadcastReceived,
+                                          0) == broadcastReceived)
+  }
+
+  it must "correctly update spent txo state based on confirmations" in {
+    wallet =>
+      assert(wallet.walletConfig.requiredConfirmations == requiredConfs)
+
+      val spendingTxId = randomTXID
+      val pendingConfSpent = utxo
+        .copyWithSpendingTxId(spendingTxId)
+        .copyWithState(PendingConfirmationsSpent)
+
+      val confSpent = utxo
+        .copyWithSpendingTxId(spendingTxId)
+        .copyWithState(ConfirmedSpent)
+
+      val withSpendingTxId =
+        reserved.copyWithSpendingTxId(DoubleSha256DigestBE.empty)
+
+      //it must transition from reserved to broacast spent
+      assert(
+        wallet
+          .updateSpentTxoWithConfs(withSpendingTxId, 0)
+          .state == TxoState.BroadcastSpent)
+
+      //it must transition from reserved to spent
+      assert(
+        wallet
+          .updateSpentTxoWithConfs(withSpendingTxId, 1)
+          .state == TxoState.PendingConfirmationsSpent)
+
+      //cannot spend an immature coinbase output
+      assertThrows[RuntimeException] {
+        //cannot have utxo spent from coinbase before 101 blocks
+        wallet.updateSpentTxoWithConfs(immatureCoinbase, 100)
+      }
+
+      val pendingConfReceivedWithTxId = pendingConfReceived
+        .copyWithSpendingTxId(DoubleSha256DigestBE.empty)
+
+      val expectedConfSpent = pendingConfReceived
+        .copyWithSpendingTxId(DoubleSha256DigestBE.empty)
+        .copyWithState(TxoState.ConfirmedSpent)
+
+      val updated =
+        wallet.updateSpentTxoWithConfs(pendingConfReceivedWithTxId,
+                                       requiredConfs)
+      //it must transition from Pending Confirmation Received -> Pending Confirmation Spent
+      assert(updated == expectedConfSpent)
+
+      //must stay at pending confirmations spent with only 1 confirmation
+      assert(
+        wallet.updateSpentTxoWithConfs(pendingConfSpent, 1) == pendingConfSpent)
+
+      //must transition from PendingConfirmationsSpent -> ConfirmedSpent
+      assert(
+        wallet.updateSpentTxoWithConfs(pendingConfSpent,
+                                       requiredConfs) == confSpent)
+
+      //transition form TxoState.ConfirmedSpent -> TxoState.PendingConfirmationSpent (reorg scenario)
+      assert(
+        wallet.updateSpentTxoWithConfs(confSpent, 1) == confSpent.copyWithState(
+          PendingConfirmationsSpent))
+
+      //stay confirmed if we are already confirmed
+      assert(
+        wallet.updateSpentTxoWithConfs(confSpent, requiredConfs) == confSpent)
+
+      val expectedConfReceivedWithTxid =
+        confReceived.copyWithSpendingTxId(DoubleSha256DigestBE.empty)
+
+      //TxoState.ConfirmedReceived -> TxoState.ConfirmedSpent
+      assert(
+        wallet.updateSpentTxoWithConfs(
+          expectedConfReceivedWithTxid,
+          requiredConfs) == expectedConfReceivedWithTxid.copyWithState(
+          TxoState.ConfirmedSpent))
+
+      //it must stay BroadcastSpent with 0 confirmations
+      val broadcastSpent = utxo
+        .copyWithSpendingTxId(DoubleSha256DigestBE.empty)
+        .copyWithState(TxoState.BroadcastSpent)
+      assert(
+        wallet.updateSpentTxoWithConfs(broadcastSpent, 0) == broadcastSpent)
   }
 }

wallet/src/main/scala/org/bitcoins/wallet/internal/TransactionProcessing.scala

@@ -323,8 +323,7 @@ private[bitcoins] trait TransactionProcessing extends WalletLogger {
       toBeUpdated = outputsBeingSpent
         .map(markAsSpent(_, transaction.txIdBE))
         .flatten
-      processed <- spendingInfoDAO.updateAllSpendingInfoDb(toBeUpdated)
-      _ <- updateUtxoConfirmedStates(processed)
+      processed <- updateUtxoSpentConfirmedStates(toBeUpdated)
     } yield {
       processed
     }
@@ -503,7 +502,7 @@ private[bitcoins] trait TransactionProcessing extends WalletLogger {
 
           // Update Txo State
           updateTxDbF.flatMap(_ =>
-            updateUtxoConfirmedState(foundTxo).flatMap {
+            updateUtxoReceiveConfirmedStates(foundTxo).flatMap {
               case Some(txo) =>
                 logger.debug(
                   s"Updated block_hash of txo=${txo.txid.hex} new block hash=${blockHash.hex}")

wallet/src/main/scala/org/bitcoins/wallet/internal/UtxoHandling.scala

@@ -82,9 +82,25 @@ private[wallet] trait UtxoHandling extends WalletLogger {
     }
   }
 
-  private[wallet] def updateUtxoConfirmedState(
+  private[wallet] def updateUtxoSpentConfirmedStates(
       txo: SpendingInfoDb): Future[Option[SpendingInfoDb]] = {
-    updateUtxoConfirmedStates(Vector(txo)).map(_.headOption)
+    updateUtxoSpentConfirmedStates(Vector(txo)).map(_.headOption)
+  }
+
+  private[wallet] def updateUtxoSpentConfirmedStates(
+      txos: Vector[SpendingInfoDb]): Future[Vector[SpendingInfoDb]] = {
+    updateUtxoStates(txos, updateSpentTxoWithConfs)
+  }
+
+  private[wallet] def updateUtxoReceiveConfirmedStates(
+      txo: SpendingInfoDb): Future[Option[SpendingInfoDb]] = {
+    updateUtxoReceiveConfirmedStates(Vector(txo))
+      .map(_.headOption)
+  }
+
+  private[wallet] def updateUtxoReceiveConfirmedStates(
+      txos: Vector[SpendingInfoDb]): Future[Vector[SpendingInfoDb]] = {
+    updateUtxoStates(txos, updateReceivedTxoWithConfs)
   }
 
   /** Returns a map of the SpendingInfoDbs with their relevant block.
@@ -124,10 +140,31 @@ private[wallet] trait UtxoHandling extends WalletLogger {
     }
   }
 
+  private[wallet] def updateSpentTxoWithConfs(
+      txo: SpendingInfoDb,
+      confs: Int): SpendingInfoDb = {
+    txo.state match {
+      case TxoState.ImmatureCoinbase =>
+        sys.error(
+          s"Cannot update txo with received state=${TxoState.ImmatureCoinbase}")
+      case TxoState.Reserved | TxoState.PendingConfirmationsSpent |
+          TxoState.ConfirmedSpent | TxoState.BroadcastSpent |
+          TxoState.PendingConfirmationsReceived | TxoState.BroadcastReceived |
+          TxoState.ConfirmedReceived =>
+        if (confs >= walletConfig.requiredConfirmations) {
+          txo.copyWithState(TxoState.ConfirmedSpent)
+        } else if (confs == 0) {
+          txo.copyWithState(TxoState.BroadcastSpent)
+        } else {
+          txo.copyWithState(TxoState.PendingConfirmationsSpent)
+        }
+    }
+  }
+
   /** Updates the SpendingInfoDb to the correct state based
     * on the number of confirmations it has received
     */
-  private[wallet] def updateTxoWithConfs(
+  private[wallet] def updateReceivedTxoWithConfs(
       txo: SpendingInfoDb,
       confs: Int): SpendingInfoDb = {
     txo.state match {
@@ -138,28 +175,30 @@ private[wallet] trait UtxoHandling extends WalletLogger {
           else
             txo.copyWithState(TxoState.PendingConfirmationsReceived)
         } else txo
-      case TxoState.PendingConfirmationsReceived | BroadcastReceived =>
-        if (confs >= walletConfig.requiredConfirmations)
+      case TxoState.PendingConfirmationsReceived | BroadcastReceived |
+          TxoState.ConfirmedReceived =>
+        if (confs >= walletConfig.requiredConfirmations) {
           txo.copyWithState(TxoState.ConfirmedReceived)
-        else txo.copyWithState(PendingConfirmationsReceived)
-      case TxoState.PendingConfirmationsSpent | BroadcastSpent =>
-        if (confs >= walletConfig.requiredConfirmations)
-          txo.copyWithState(TxoState.ConfirmedSpent)
-        else txo.copyWithState(PendingConfirmationsSpent)
+        } else if (confs == 0) {
+          txo.copyWithState(TxoState.BroadcastReceived)
+        } else txo.copyWithState(PendingConfirmationsReceived)
       case TxoState.Reserved =>
-        // We should keep the utxo as reserved so it is not used in
-        // a future transaction that it should not be in
-        txo
-      case TxoState.ConfirmedReceived | TxoState.ConfirmedSpent =>
+        //do nothing if we have reserved the utxo
         txo
+      case state: SpentState =>
+        sys.error(s"Cannot update spendingInfoDb in spent state=$state")
+
     }
   }
 
   /** Updates all the given SpendingInfoDbs to the correct state
     * based on how many confirmations they have received
+    * @param spendingInfoDbs the utxos we need to update
+    * @param fn the function used to transition the [[TxoState]] given a utxo and number of confirmations
     */
-  private[wallet] def updateUtxoConfirmedStates(
-      spendingInfoDbs: Vector[SpendingInfoDb]): Future[
+  private def updateUtxoStates(
+      spendingInfoDbs: Vector[SpendingInfoDb],
+      fn: (SpendingInfoDb, Int) => SpendingInfoDb): Future[
     Vector[SpendingInfoDb]] = {
     val relevantBlocksF: Future[
       Map[Option[DoubleSha256DigestBE], Vector[SpendingInfoDb]]] = {
@@ -185,12 +224,12 @@ private[wallet] trait UtxoHandling extends WalletLogger {
                 s"Given txos exist in block (${blockHashWithConfs.blockHash.hex}) that we do not have or that has been reorged! $txos")
               Vector.empty
             case Some(confs) =>
-              txos.map(updateTxoWithConfs(_, confs))
+              txos.map(fn(_, confs))
           }
         case (None, txos) =>
           logger.debug(
             s"Currently have ${txos.size} transactions in the mempool")
-          Vector.empty
+          txos
       }.toVector
 
       toUpdateFs
@@ -351,7 +390,7 @@ private[wallet] trait UtxoHandling extends WalletLogger {
     for {
       updatedUtxos <- updatedUtxosF
       // update the confirmed utxos
-      updatedConfirmed <- updateUtxoConfirmedStates(updatedUtxos)
+      updatedConfirmed <- updateUtxoReceiveConfirmedStates(updatedUtxos)
 
       // update the utxos that are in blocks but not considered confirmed yet
       pendingConf = updatedUtxos.filterNot(utxo =>
@@ -378,7 +417,10 @@ private[wallet] trait UtxoHandling extends WalletLogger {
     for {
       infos <- spendingInfoDAO.findAllPendingConfirmation
       _ = logger.debug(s"Updating states of ${infos.size} pending utxos...")
-      updatedInfos <- updateUtxoConfirmedStates(infos)
-    } yield updatedInfos
+      receivedUtxos = infos.filter(_.state.isInstanceOf[ReceivedState])
+      spentUtxos = infos.filter(_.state.isInstanceOf[SpentState])
+      updatedReceivedInfos <- updateUtxoReceiveConfirmedStates(receivedUtxos)
+      updatedSpentInfos <- updateUtxoSpentConfirmedStates(spentUtxos)
+    } yield (updatedReceivedInfos ++ updatedSpentInfos).toVector
   }
 }