Added ScalaDoc, added basic test for iterateAVLTree

Author: jellymlg

avldb/src/main/scala/scorex/crypto/authds/avltree/batch/VersionedLDBAVLStorage.scala

@@ -1,7 +1,6 @@
 package scorex.crypto.authds.avltree.batch
 
 import com.google.common.primitives.Ints
-import org.ergoplatform.serialization.ManifestSerializer.MainnetManifestDepth
 import scorex.crypto.authds.avltree.batch.Constants.{DigestType, HashFnType, hashFn}
 import scorex.crypto.authds.avltree.batch.VersionedLDBAVLStorage.{noStoreSerializer, topNodeHashKey, topNodeHeightKey}
 import scorex.crypto.authds.avltree.batch.serialization.{BatchAVLProverManifest, BatchAVLProverSubtree, ProxyInternalNode}
@@ -155,7 +154,13 @@ class VersionedLDBAVLStorage(store: LDBVersionedStore)
     }
   }
 
-  def iterateAVLTree(fromIndex: Int)(handleSubtree: BatchAVLProverSubtree[DigestType] => Unit): Unit =
+  /**
+   * Split the AVL+ tree to 2^depth number of subtrees and process them
+   * @param fromIndex - only start processing subtrees from this index
+   * @param depth - depth at whitch to split AVL+ tree to subtrees
+   * @param handleSubtree - function to process subtree
+   */
+  def iterateAVLTree(fromIndex: Int, depth: Int)(handleSubtree: BatchAVLProverSubtree[DigestType] => Unit): Unit =
     store.processSnapshot { dbReader =>
 
       var current: Int = 0
@@ -172,7 +177,7 @@ class VersionedLDBAVLStorage(store: LDBVersionedStore)
 
       def proxyLoop(label: Array[Byte], level: Int): Unit =
         noStoreSerializer.parseBytes(dbReader.get(label)) match {
-          case in: ProxyInternalNode[DigestType] if level == MainnetManifestDepth =>
+          case in: ProxyInternalNode[DigestType] if level == depth =>
             if(current >= fromIndex) handleSubtree(subtree(in.leftLabel))
             current += 1
             if(current >= fromIndex) handleSubtree(subtree(in.rightLabel))

avldb/src/test/scala/scorex/crypto/authds/avltree/batch/VersionedLDBAVLStorageSpecification.scala

@@ -12,10 +12,11 @@ import scorex.crypto.authds.{ADDigest, ADKey, ADValue, SerializedAdProof}
 import scorex.util.encode.Base16
 import scorex.crypto.hash.{Blake2b256, Digest32}
 import scorex.db.{LDBFactory, LDBVersionedStore}
-import scorex.util.ByteArrayBuilder
+import scorex.util.{ByteArrayBuilder, ModifierId, bytesToId}
 import scorex.util.serialization.VLQByteBufferWriter
 import scorex.utils.{Random => RandomBytes}
 
+import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.duration._
 import scala.concurrent.{Await, Future}
@@ -370,4 +371,22 @@ class VersionedLDBAVLStorageSpecification
     }
   }
 
+  property("iterate AVL tree") {
+    val prover = createPersistentProver()
+    val current = 11
+    val depth = 6
+    blockchainWorkflowTest(prover)
+
+    val chunkBuffer: ArrayBuffer[(ModifierId,Array[Array[Byte]])] = ArrayBuffer.empty[(ModifierId,Array[Array[Byte]])]
+
+    prover.storage.asInstanceOf[VersionedLDBAVLStorage].iterateAVLTree(current, depth) { subtree =>
+      chunkBuffer += ((
+        bytesToId(subtree.id),
+        subtree.leafValues.toArray
+        ))
+    }
+
+    chunkBuffer.size shouldBe math.pow(2, depth) - current
+  }
+
 }

src/main/scala/org/ergoplatform/nodeView/history/ErgoHistory.scala

@@ -8,7 +8,7 @@ import org.ergoplatform.mining.AutolykosPowScheme
 import org.ergoplatform.modifiers.history._
 import org.ergoplatform.modifiers.history.header.{Header, PreGenesisHeader}
 import org.ergoplatform.modifiers.{BlockSection, ErgoFullBlock, NonHeaderBlockSection}
-import org.ergoplatform.nodeView.history.UtxoSetScanner.StartUtxoSetScanWithHistory
+import org.ergoplatform.nodeView.history.UtxoSetScanner.InitializeUtxoSetScannerWithHistory
 import org.ergoplatform.nodeView.history.extra.ExtraIndexer.ReceivableMessages.StartExtraIndexer
 import org.ergoplatform.nodeView.history.extra.ExtraIndexer.{IndexedHeightKey, NewestVersion, NewestVersionBytes, SchemaVersionKey, getIndex}
 import org.ergoplatform.nodeView.history.storage.HistoryStorage
@@ -304,7 +304,7 @@ object ErgoHistory extends ScorexLogging {
 
     // set history for utxo set scanner, if bootstrapping by snapshot
     if(ergoSettings.nodeSettings.utxoSettings.utxoBootstrap) {
-      context.system.eventStream.publish(StartUtxoSetScanWithHistory(history))
+      context.system.eventStream.publish(InitializeUtxoSetScannerWithHistory(history))
     }
 
     history

src/main/scala/org/ergoplatform/nodeView/history/UtxoSetScanner.scala

@@ -12,6 +12,7 @@ import org.ergoplatform.nodeView.state.UtxoState
 import org.ergoplatform.nodeView.wallet.ErgoWallet
 import org.ergoplatform.nodeView.wallet.ErgoWalletActorMessages.ScanBoxesFromUtxoSnapshot
 import org.ergoplatform.serialization.ManifestSerializer
+import org.ergoplatform.serialization.ManifestSerializer.MainnetManifestDepth
 import org.ergoplatform.wallet.boxes.ErgoBoxSerializer
 import scorex.crypto.authds.avltree.batch.VersionedLDBAVLStorage
 import scorex.crypto.hash.Blake2b256
@@ -25,6 +26,12 @@ import scala.concurrent.Await
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.duration.Duration
 
+/**
+ * This class is used to provide the current UTXO set for wallet scans when bootstrapping
+ * by UTXO set snapshot. This is done by creating a snapshot of the UTXO set, deserializing
+ * the raw bytes to ErgoBoxes and sending them to the wallet actor in chunks.
+ * @param nodeView - NodeView actor to get wallet and UTXOs from
+ */
 class UtxoSetScanner(nodeView: ActorRef) extends Actor with ScorexLogging {
 
   private var history: ErgoHistory = _
@@ -33,22 +40,39 @@ class UtxoSetScanner(nodeView: ActorRef) extends Actor with ScorexLogging {
   private implicit val timeout: Timeout = Timeout(10, TimeUnit.SECONDS)
   private implicit val duration: Duration = Duration.create(10, TimeUnit.SECONDS)
 
+  /**
+   * Internal buffer that holds deserialized AVL subtrees until they are sent to wallet
+   */
   private val chunkBuffer: ArrayBuffer[(ModifierId,Array[ErgoBox])] = ArrayBuffer.empty[(ModifierId,Array[ErgoBox])]
 
+  /**
+   * Reads the current progress of the scanner.
+   * @return (current segment, total segments)
+   */
   private def readProgress(): (Int, Int) =
     historyStorage.getIndex(utxoSetScanProgressKey).map(ByteBuffer.wrap).map { buffer =>
       val current = buffer.getInt
       val total = buffer.getInt
       (current, total)
     }.getOrElse((0, 0))
 
+  /**
+   * Writes progress to db.
+   * @param current - current retrieved segment
+   * @param total - total segment count
+   */
   private def writeProgress(current: Int, total: Int): Unit = {
     val buffer: ByteBuffer = ByteBuffer.allocate(8)
     buffer.putInt(current)
     buffer.putInt(total)
     historyStorage.insert(Array((utxoSetScanProgressKey, buffer.array)), Array.empty[BlockSection])
   }
 
+  /**
+   * Send deserialized AVL subtrees to wallet for scanning.
+   * @param wallet - wallet to send to
+   * @param current - current retrieved segment
+   */
   private def sendBufferToWallet(wallet: ErgoWallet, current: Int): Unit = {
     wallet.scanUtxoSnapshot(ScanBoxesFromUtxoSnapshot(chunkBuffer, current, MainnetTotal))
     writeProgress(current, MainnetTotal)
@@ -69,11 +93,12 @@ class UtxoSetScanner(nodeView: ActorRef) extends Actor with ScorexLogging {
     )
 
     val initialized: Boolean = Await.result(wallet.getWalletStatus.map(_.initialized), duration)
-    if(!initialized) return
+    if(!initialized) // wallet is not initialized
+      return
 
     log.info(s"Starting UTXO set snapshot scan for $total chunks")
 
-    state.persistentProver.storage.asInstanceOf[VersionedLDBAVLStorage].iterateAVLTree(current) { subtree =>
+    state.persistentProver.storage.asInstanceOf[VersionedLDBAVLStorage].iterateAVLTree(current, MainnetManifestDepth) { subtree =>
       current += 1
 
       chunkBuffer += ((
@@ -93,6 +118,7 @@ class UtxoSetScanner(nodeView: ActorRef) extends Actor with ScorexLogging {
 
     if(current == total) {
       log.info(s"Successfully scanned $total Utxo set subtrees")
+      // send newest block to wallet, if blocks were applied since scan began it will go back to scan them
       wallet.scanPersistent(history.bestFullBlockOpt.get)
     }else {
       log.error(s"Inconsistent Utxo set scan state: $current scanned subtrees out of $total")
@@ -101,27 +127,40 @@ class UtxoSetScanner(nodeView: ActorRef) extends Actor with ScorexLogging {
   }
 
   override def receive: Receive = {
-    case StartUtxoSetScanWithHistory(history: ErgoHistory) =>
+    case InitializeUtxoSetScannerWithHistory(history: ErgoHistory) =>
       this.history = history
       run()
     case StartUtxoSetScan(rescan: Boolean) =>
-      if(readProgress()._1 == 0 || rescan) writeProgress(0, MainnetTotal)
+      if(readProgress()._1 == 0 || //
+        rescan) // start over UTXO set scan
+        writeProgress(0, MainnetTotal)
       run()
   }
 
   override def preStart(): Unit = {
-    context.system.eventStream.subscribe(self, classOf[StartUtxoSetScanWithHistory])
+    context.system.eventStream.subscribe(self, classOf[InitializeUtxoSetScannerWithHistory])
     context.system.eventStream.subscribe(self, classOf[StartUtxoSetScan])
   }
 
 }
 
 object UtxoSetScanner {
 
-  case class StartUtxoSetScanWithHistory(history: ErgoHistory)
+  /**
+   * Initialize UTXO set scanner with database and try continuing scan if possible
+   * @param history - database handle
+   */
+  case class InitializeUtxoSetScannerWithHistory(history: ErgoHistory)
 
+  /**
+   * Start scanning UTXO set, or continue if the scan was interrupted, or start over if rescan = true
+   * @param rescan - whether to start over or just continue scan
+   */
   case class StartUtxoSetScan(rescan: Boolean)
 
+  /**
+   * Number of subtrees to divide AVL tree to
+   */
   private val MainnetTotal: Int = math.pow(2, ManifestSerializer.MainnetManifestDepth).toInt
 
   private val utxoSetScanProgressKey: ByteArrayWrapper =

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

@@ -198,6 +198,9 @@ object WalletScanLogic extends ScorexLogging {
 
     val scanRes = ScanResults(myOutputs, Seq.empty, Seq.empty)
 
+    /** Pass subtreeId as blockId; set height to 0 so when UTXO set scan is finished normal wallet scan
+     *  will start with the first non-pruned block (see [[ErgoWalletState.expectedNextBlockHeight]])
+     */
     updateRegistryAndOffchain(registry, offChainRegistry, outputsFilter, scanRes, subtreeId, 0)
   }