Persist unconfirmed wallet transactions and restore them at start-up

src/main/scala/org/ergoplatform/nodeView/ErgoNodeViewHolder.scala

@@ -29,6 +29,7 @@ import org.ergoplatform.modifiers.history.extension.Extension
 
 import scala.annotation.tailrec
 import scala.collection.mutable
+import scala.concurrent.ExecutionContext
 import scala.util.{Failure, Success, Try}
 
 /**
@@ -87,12 +88,35 @@ abstract class ErgoNodeViewHolder[State <: ErgoState[State]](settings: ErgoSetti
         Escalate
     }
 
+  override def preStart(): Unit = {
+    super.preStart()
+    restoreUnconfirmedWalletTransactions()
+  }
+
   override def postStop(): Unit = {
     log.warn("Stopping ErgoNodeViewHolder")
     history().closeStorage()
     minimalState().closeStorage()
   }
 
+  /**
+    * The memory pool is not persisted, so a restart drops every transaction which was waiting in it.
+    * The wallet does keep its own unconfirmed transactions, so ask it for them and put them back,
+    * instead of waiting for peers to gossip them again - which they may never do.
+    */
+  private def restoreUnconfirmedWalletTransactions(): Unit = {
+    implicit val ec: ExecutionContext = context.dispatcher
+    vault().unconfirmedTransactionsToRestore.onComplete {
+      case Success(txs) if txs.nonEmpty =>
+        log.info(s"Putting ${txs.size} unconfirmed wallet transaction(s) back into the memory pool")
+        txs.foreach(tx => self ! RestoredTransaction(UnconfirmedTransaction(tx, None)))
+      case Success(_) =>
+        log.debug("No unconfirmed wallet transactions to restore")
+      case Failure(t) =>
+        log.warn("Could not read unconfirmed wallet transactions to restore: ", t)
+    }
+  }
+
   /**
     * Update NodeView with new components and notify subscribers of changed components
     *
@@ -658,6 +682,17 @@ abstract class ErgoNodeViewHolder[State <: ErgoState[State]](settings: ErgoSetti
       txModify(unconfirmedTx)
     case LocallyGeneratedTransaction(unconfirmedTx) =>
       sender() ! txModify(unconfirmedTx)
+    case RestoredTransaction(unconfirmedTx) =>
+      txModify(unconfirmedTx) match {
+        case _: ProcessingOutcome.Accepted =>
+          log.info(s"Unconfirmed wallet transaction ${unconfirmedTx.id} is back in the memory pool")
+        case outcome =>
+          // the transaction can not be brought back, e.g. it got on the blockchain while the node
+          // was down, or a conflicting one did. There is no point in keeping it for the next restart
+          log.info(s"Unconfirmed wallet transaction ${unconfirmedTx.id} was not accepted back " +
+            s"into the memory pool ($outcome), forgetting it")
+          vault().forgetUnconfirmedTransactions(Seq(unconfirmedTx.id))
+      }
     case RecheckedTransactions(unconfirmedTxs) =>
       val updatedPool = memoryPool().put(unconfirmedTxs)
       updateNodeView(updatedMempool = Some(updatedPool))
@@ -737,6 +772,12 @@ object ErgoNodeViewHolder {
       */
     case class TransactionFromRemote(unconfirmedTx: UnconfirmedTransaction)
 
+    /**
+      * Wrapper for a wallet transaction which was unconfirmed when the node was stopped and is being
+      * put back into the memory pool now
+      */
+    case class RestoredTransaction(unconfirmedTx: UnconfirmedTransaction)
+
     /**
       * Wrapper for transactions which sit in mempool for long enough time, so `CleanWorker` is re-checking their
       * validity and then sending via this message to update the mempool

src/main/scala/org/ergoplatform/nodeView/wallet/ErgoWallet.scala

@@ -9,7 +9,7 @@ import org.ergoplatform.nodeView.wallet.ErgoWalletActorMessages._
 import org.ergoplatform.settings.{ErgoSettings, Parameters}
 import org.ergoplatform.wallet.boxes.{ReemissionData, ReplaceCompactCollectBoxSelector}
 import org.ergoplatform.core.VersionTag
-import scorex.util.ScorexLogging
+import scorex.util.{ModifierId, ScorexLogging}
 
 import scala.util.{Failure, Success, Try}
 
@@ -46,6 +46,14 @@ class ErgoWallet(historyReader: ErgoHistoryReader, settings: ErgoSettings, param
     this
   }
 
+  /**
+    * Tell the wallet to stop keeping the given unconfirmed transactions across restarts, e.g.
+    * because the memory pool refused them when they were re-submitted.
+    */
+  def forgetUnconfirmedTransactions(ids: Seq[ModifierId]): Unit = {
+    walletActor ! ForgetUnconfirmedTransactions(ids)
+  }
+
   def scanPersistent(modifier: BlockSection): ErgoWallet = {
     modifier match {
       case fb: ErgoFullBlock =>

src/main/scala/org/ergoplatform/nodeView/wallet/ErgoWalletActor.scala

@@ -4,11 +4,14 @@ import akka.actor.SupervisorStrategy.{Restart, Stop}
 import akka.actor._
 import akka.pattern.StatusReply
 import org.ergoplatform.ErgoBox._
+import org.ergoplatform.modifiers.ErgoFullBlock
+import org.ergoplatform.modifiers.mempool.ErgoTransaction
 import org.ergoplatform.network.ErgoNodeViewSynchronizerMessages.{ChangedMempool, ChangedState}
 import org.ergoplatform.nodeView.history.ErgoHistoryReader
 import org.ergoplatform.nodeView.mempool.ErgoMemPoolReader
 import org.ergoplatform.nodeView.state.ErgoStateReader
 import org.ergoplatform.nodeView.wallet.ErgoWalletServiceUtils.DeriveNextKeyResult
+import org.ergoplatform.nodeView.wallet.persistence.WalletStorage
 import org.ergoplatform.sdk.wallet.secrets.DerivationPath
 import org.ergoplatform.settings._
 import org.ergoplatform.wallet.Constants.ScanId
@@ -18,8 +21,9 @@ import org.ergoplatform._
 import org.ergoplatform.core.VersionTag
 import org.ergoplatform.sdk.SecretString
 import org.ergoplatform.utils.ScorexEncoding
-import scorex.util.ScorexLogging
+import scorex.util.{ModifierId, ScorexLogging, bytesToId}
 
+import scala.annotation.tailrec
 import scala.concurrent.duration._
 import scala.util.{Failure, Success}
 
@@ -76,7 +80,7 @@ class ErgoWalletActor(settings: ErgoSettings,
       val ws = settings.walletSettings
       // Try to read wallet from json file or test mnemonic provided in a config file
       val newState = ergoWalletService.readWallet(state, ws.testMnemonic.map(SecretString.create(_)), ws.testKeysQty, ws.secretStorage)
-      context.become(loadedWallet(newState))
+      context.become(loadedWallet(ergoWalletService.restoreOffChainState(newState)))
       unstashAll()
     case _ => // stashing all messages until wallet is setup
       stash()
@@ -223,14 +227,23 @@ class ErgoWalletActor(settings: ErgoSettings,
     /* SCAN COMMANDS */
     //scan mempool transaction
     case ScanOffChain(tx) =>
-      val dustLimit = settings.walletSettings.dustLimit
-      val newWalletBoxes = WalletScanLogic.extractWalletOutputs(tx, None, state.walletVars, dustLimit)
-      val inputs = WalletScanLogic.extractInputBoxes(tx)
-      val newState = state.copy(offChainRegistry =
-        state.offChainRegistry.updateOnTransaction(newWalletBoxes, inputs, state.walletVars.externalScans)
-      )
+      val (newState, walletAffected) = ergoWalletService.scanOffChainUpdate(state, tx)
+      if (walletAffected) {
+        // the transaction is kept until it gets on the blockchain, so that a restart in the meantime
+        // does not make the wallet consider its inputs spendable again
+        state.storage.addUnconfirmedTransaction(tx, state.fullHeight) match {
+          case Success(_) =>
+          case Failure(t) => log.error(s"Could not store unconfirmed transaction ${tx.id}: ", t)
+        }
+      }
       context.become(loadedWallet(newState))
 
+    case ReadUnconfirmedTransactions =>
+      sender() ! unconfirmedTransactionsToRestore(state)
+
+    case ForgetUnconfirmedTransactions(ids) =>
+      forgetUnconfirmedTransactions(state, ids)
+
     // rescan=true means we serve a user request for rescan from arbitrary height
     case ScanInThePast(blockHeight, rescan) =>
       val nextBlockHeight = state.expectedNextBlockHeight(blockHeight, settings.nodeSettings.isFullBlocksPruned)
@@ -275,6 +288,7 @@ class ErgoWalletActor(settings: ErgoSettings,
               case Success(updatedState) =>
                 updatedState
             }
+          forgetConfirmedAndExpired(newState, newBlock)
           context.become(loadedWallet(newState))
         } else if (nextBlockHeight < newBlock.height) {
           log.warn(s"Wallet: skipped blocks found starting from $nextBlockHeight, going back to scan them")
@@ -489,6 +503,48 @@ class ErgoWalletActor(settings: ErgoSettings,
     sender() ! txsToSend
   }
 
+  /**
+    * Stored unconfirmed transactions worth putting back into the memory pool. Transactions which did
+    * not make it onto the blockchain for too long are dropped instead of being re-submitted forever.
+    */
+  private def unconfirmedTransactionsToRestore(state: ErgoWalletState): Seq[ErgoTransaction] = {
+    val (fresh, expired) = state.storage.readUnconfirmedTransactions().partition { case (_, seenAt) =>
+      state.fullHeight - seenAt <= WalletStorage.UnconfirmedTxLifetimeInBlocks
+    }
+    if (expired.nonEmpty) {
+      forgetUnconfirmedTransactions(state, expired.map(_._1.id))
+    }
+    ErgoWalletActor.orderByDependency(fresh.map(_._1))
+  }
+
+  /**
+    * Drop the transactions of a just applied block from the store, and give up on the ones which
+    * stayed unconfirmed for longer than [[WalletStorage.UnconfirmedTxLifetimeInBlocks]] blocks.
+    */
+  private def forgetConfirmedAndExpired(state: ErgoWalletState, block: ErgoFullBlock): Unit = {
+    val seenAtHeights = state.storage.unconfirmedTransactionHeights
+    if (seenAtHeights.nonEmpty) {
+      val confirmed = block.transactions.map(_.id).filter(seenAtHeights.contains)
+      val expired = seenAtHeights.collect {
+        case (id, seenAt) if block.height - seenAt > WalletStorage.UnconfirmedTxLifetimeInBlocks => id
+      }.toSeq
+      if (expired.nonEmpty) {
+        log.warn(s"Wallet gave up on ${expired.size} transaction(s) still unconfirmed at height ${block.height}")
+      }
+      val toForget = (confirmed ++ expired).distinct
+      if (toForget.nonEmpty) {
+        forgetUnconfirmedTransactions(state, toForget)
+      }
+    }
+  }
+
+  private def forgetUnconfirmedTransactions(state: ErgoWalletState, ids: Seq[ModifierId]): Unit = {
+    state.storage.removeUnconfirmedTransactions(ids) match {
+      case Success(_) =>
+      case Failure(t) => log.error("Could not forget unconfirmed transactions: ", t)
+    }
+  }
+
   override def receive: Receive = emptyWallet
 
   private def wrapLegalExc[T](e: Throwable): Failure[T] =
@@ -503,6 +559,42 @@ class ErgoWalletActor(settings: ErgoSettings,
 
 object ErgoWalletActor extends ScorexLogging {
 
+  /**
+    * Order transactions so that a transaction spending an output of another one comes after it.
+    *
+    * Unconfirmed transactions do form such chains, and both consumers of this ordering need it: the
+    * off-chain registry only nets a spending out if it has seen the box being spent already, and the
+    * memory pool refuses a transaction whose inputs it does not know yet.
+    *
+    * Transactions with no producer among `txs` keep their relative order. A cycle is impossible
+    * between valid transactions, but should one be given, its members are appended unordered rather
+    * than dropped.
+    */
+  def orderByDependency(txs: Seq[ErgoTransaction]): Seq[ErgoTransaction] = {
+    val producerOf: Map[ModifierId, ModifierId] =
+      txs.flatMap(tx => tx.outputs.map(out => bytesToId(out.id) -> tx.id)).toMap
+
+    @tailrec
+    def loop(remaining: Seq[ErgoTransaction],
+             ordered: Set[ModifierId],
+             acc: Seq[ErgoTransaction]): Seq[ErgoTransaction] = {
+      if (remaining.isEmpty) {
+        acc
+      } else {
+        val (ready, blocked) = remaining.partition { tx =>
+          tx.inputs.forall(in => producerOf.get(bytesToId(in.boxId)).forall(ordered.contains))
+        }
+        if (ready.isEmpty) {
+          acc ++ blocked
+        } else {
+          loop(blocked, ordered ++ ready.map(_.id), acc ++ ready)
+        }
+      }
+    }
+
+    loop(txs, Set.empty, Seq.empty)
+  }
+
   /** Start actor and register its proper closing into coordinated shutdown */
   def apply(settings: ErgoSettings,
             parameters: Parameters,

src/main/scala/org/ergoplatform/nodeView/wallet/ErgoWalletActorMessages.scala

@@ -51,6 +51,21 @@ object ErgoWalletActorMessages {
    */
   final case class ScanOnChain(block: ErgoFullBlock)
 
+  /**
+   * Read wallet-related transactions which were not on the blockchain yet when the node was stopped,
+   * so that they can be put back into the memory pool. Answered with a `Seq[ErgoTransaction]`,
+   * ordered so that a transaction spending an output of another one comes after it.
+   */
+  final case object ReadUnconfirmedTransactions
+
+  /**
+   * Stop keeping the given unconfirmed transactions across restarts, e.g. because the memory pool
+   * refused them
+   *
+   * @param ids - identifiers of the transactions to forget
+   */
+  final case class ForgetUnconfirmedTransactions(ids: Seq[ModifierId])
+
   /**
    * Rollback to previous version of the wallet, by throwing away effects of blocks after the version
    *

src/main/scala/org/ergoplatform/nodeView/wallet/ErgoWalletReader.scala

@@ -97,6 +97,14 @@ trait ErgoWalletReader extends NodeViewComponent {
   def transactionById(id: ModifierId): Future[Option[AugWalletTransaction]] =
     (walletActor ? GetTransaction(id)).mapTo[Option[AugWalletTransaction]]
 
+  /**
+    * Wallet-related transactions which were not on the blockchain yet when the node was stopped, so
+    * that they can be put back into the memory pool. Ordered so that a transaction spending an
+    * output of another one comes after it.
+    */
+  def unconfirmedTransactionsToRestore: Future[Seq[ErgoTransaction]] =
+    (walletActor ? ReadUnconfirmedTransactions).mapTo[Seq[ErgoTransaction]]
+
   def generateTransaction(requests: Seq[TransactionGenerationRequest],
                           inputsRaw: Seq[String] = Seq.empty,
                           dataInputsRaw: Seq[String] = Seq.empty): Future[Try[ErgoTransaction]] =

src/main/scala/org/ergoplatform/nodeView/wallet/ErgoWalletService.scala

@@ -220,6 +220,21 @@ trait ErgoWalletService {
     */
   def scanBlockUpdate(state: ErgoWalletState, block: ErgoFullBlock, dustLimit: Option[Long]): Try[ErgoWalletState]
 
+  /**
+    * Update the off-chain part of the wallet with a transaction which is not on the blockchain yet.
+    *
+    * @return the updated state, and whether the transaction is of any interest to the wallet, that
+    *         is whether it pays the wallet or spends one of its boxes
+    */
+  def scanOffChainUpdate(state: ErgoWalletState, tx: ErgoTransaction): (ErgoWalletState, Boolean)
+
+  /**
+    * Rebuild the off-chain part of the wallet from the transactions which were unconfirmed when the
+    * node was stopped. Without it the wallet forgets that it already spent their inputs and can
+    * build a conflicting transaction spending the same boxes a second time.
+    */
+  def restoreOffChainState(state: ErgoWalletState): ErgoWalletState
+
   /**
     * Sign a transaction
     */
@@ -596,6 +611,33 @@ class ErgoWalletServiceImpl(override val ergoSettings: ErgoSettings) extends Erg
         state.copy(registry = reg, offChainRegistry = offReg, outputsFilter = Some(updatedOutputsFilter))
       }
 
+  override def scanOffChainUpdate(state: ErgoWalletState, tx: ErgoTransaction): (ErgoWalletState, Boolean) = {
+    val dustLimit = ergoSettings.walletSettings.dustLimit
+    val newWalletBoxes = WalletScanLogic.extractWalletOutputs(tx, None, state.walletVars, dustLimit)
+    val inputs = WalletScanLogic.extractInputBoxes(tx)
+
+    def spendsWalletBox: Boolean =
+      state.offChainRegistry.offChainBoxes.exists(box => inputs.contains(box.boxId)) ||
+        tx.inputs.exists(input => state.registry.getBox(input.boxId).isDefined)
+
+    val newState = state.copy(offChainRegistry =
+      state.offChainRegistry.updateOnTransaction(newWalletBoxes, inputs, state.walletVars.externalScans)
+    )
+    (newState, newWalletBoxes.nonEmpty || spendsWalletBox)
+  }
+
+  override def restoreOffChainState(state: ErgoWalletState): ErgoWalletState = {
+    val stored = state.storage.readUnconfirmedTransactions().map(_._1)
+    if (stored.isEmpty) {
+      state
+    } else {
+      log.info(s"Wallet is restoring ${stored.size} unconfirmed transaction(s) kept across restart")
+      ErgoWalletActor.orderByDependency(stored).foldLeft(state) { case (acc, tx) =>
+        scanOffChainUpdate(acc, tx)._1
+      }
+    }
+  }
+
   override def updateUtxoState(state: ErgoWalletState): ErgoWalletState = {
     (state.mempoolReaderOpt, state.stateReaderOpt) match {
       case (Some(mr), Some(sr)) =>