-
Notifications
You must be signed in to change notification settings - Fork 167
/
VersionedLDBAVLStorage.scala
250 lines (214 loc) · 10.6 KB
/
VersionedLDBAVLStorage.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
package scorex.crypto.authds.avltree.batch
import com.google.common.primitives.Ints
import scorex.crypto.authds.avltree.batch.Constants.{DigestType, HashFnType, hashFn}
import scorex.crypto.authds.avltree.batch.VersionedLDBAVLStorage.{topNodeHashKey, topNodeHeightKey}
import scorex.crypto.authds.avltree.batch.serialization.{BatchAVLProverManifest, BatchAVLProverSubtree, ProxyInternalNode}
import scorex.crypto.authds.{ADDigest, ADKey}
import scorex.util.encode.Base58
import scorex.crypto.hash
import scorex.crypto.hash.Digest32
import scorex.db.{LDBKVStore, LDBVersionedStore}
import scorex.util.ScorexLogging
import scala.collection.mutable
import scala.util.{Failure, Try}
/**
* Persistent versioned authenticated AVL+ tree implementation on top of versioned LevelDB storage
*
* @param store - level db storage to save the tree in
*/
class VersionedLDBAVLStorage(store: LDBVersionedStore)
extends VersionedAVLStorage[DigestType] with ScorexLogging {
import VersionedLDBAVLStorage.nodeLabel
private def restorePrunedRootNode(): Try[(ProverNodes[DigestType], Int)] = Try {
val rootNode = VersionedLDBAVLStorage.fetch(ADKey @@ store.get(topNodeHashKey).get)(store)
val rootHeight = Ints.fromByteArray(store.get(topNodeHeightKey).get)
rootNode -> rootHeight
}
/**
* Restore prover's tree from database with just pruned top node (so only one node is fetched)
*/
def restorePrunedProver(): Try[BatchAVLProver[DigestType, HashFnType]] = {
restorePrunedRootNode().map { recoveredTop =>
new BatchAVLProver(StateTreeParameters.keySize, StateTreeParameters.valueSize, Some(recoveredTop))(hashFn)
}
}
override def rollback(version: ADDigest): Try[(ProverNodes[DigestType], Int)] = Try {
if (!this.version.contains(version)) { // do not rollback to self
store.rollbackTo(version)
}
}.flatMap(_ => restorePrunedRootNode())
.recoverWith { case e =>
log.warn(s"Failed to recover tree for digest ${Base58.encode(version)}:", e)
Failure(e)
}
override def version: Option[ADDigest] = store.lastVersionID.map(d => ADDigest @@ d)
def rollbackVersions: Iterable[ADDigest] = store.rollbackVersions().map(d => ADDigest @@ d)
override def update[K <: Array[Byte], V <: Array[Byte]](prover: BatchAVLProver[DigestType, _],
additionalData: Seq[(K, V)]): Try[Unit] = {
val digestWrapper = prover.digest
val indexes = Seq(topNodeHashKey -> nodeLabel(prover.topNode),
topNodeHeightKey -> Ints.toByteArray(prover.rootNodeHeight))
val toInsert = serializedVisitedNodes(prover.topNode, isTop = true)
val toRemove = prover.removedNodes().map(rn => rn.label)
val toUpdate = indexes ++ toInsert
val toUpdateWithWrapped = toUpdate ++ additionalData
store.update(digestWrapper, toRemove, toUpdateWithWrapped)
}
//todo: optimize the method
private def serializedVisitedNodes(node: ProverNodes[DigestType],
isTop: Boolean): Array[(Array[Byte], Array[Byte])] = {
// Should always serialize top node. It may not be new if it is the creation of the tree
if (node.isNew || isTop) {
val pair: (Array[Byte], Array[Byte]) = (nodeLabel(node), VersionedLDBAVLStorage.noStoreSerializer.toBytes(node))
node match {
case n: InternalProverNode[DigestType] =>
val leftSubtree = serializedVisitedNodes(n.left, isTop = false)
val rightSubtree = serializedVisitedNodes(n.right, isTop = false)
pair +: (leftSubtree ++ rightSubtree)
case _: ProverLeaf[DigestType] => Array(pair)
}
} else {
Array.empty
}
}
//todo: this method is not used, should be removed on next scrypto update?
override def update(prover: BatchAVLProver[DigestType, _]): Try[Unit] = update(prover, Nil)
/**
* Dump current state of AVL+ tree into another (non-versioned) storage. Tree is dumped in
* (manifest, subtrees) format
*
* @param dumpStorage - non-versioned storage to dump tree to
* @param manifestDepth - depth of manifest tree
* @param expectedRootHash - expected UTXO set authenticating tree root hash
* @return - hash of root node of tree, or failure if an error (e.g. in database) happened
*/
def dumpSnapshot(dumpStorage: LDBKVStore, manifestDepth: Byte, expectedRootHash: Array[Byte]): Try[Array[Byte]] = {
store.processSnapshot { dbReader =>
def subtreeLoop(label: DigestType, builder: mutable.ArrayBuilder[Byte]): Unit = {
val nodeBytes = dbReader.get(label)
builder ++= nodeBytes
val node = VersionedLDBAVLStorage.noStoreSerializer.parseBytes(nodeBytes)
node match {
case in: ProxyInternalNode[DigestType] =>
subtreeLoop(Digest32 @@@ in.leftLabel, builder)
subtreeLoop(Digest32 @@@ in.rightLabel, builder)
case _ =>
}
}
def dumpSubtree(sid: DigestType): Try[Unit] = {
val builder = mutable.ArrayBuilder.make[Byte]()
builder.sizeHint(200000)
subtreeLoop(sid, builder)
dumpStorage.insert(sid, builder.result())
}
def manifestLoop(nodeDbKey: Array[Byte], level: Int, manifestBuilder: mutable.ArrayBuilder[Byte]): Unit = {
val nodeBytes = dbReader.get(nodeDbKey)
manifestBuilder ++= nodeBytes
val node = VersionedLDBAVLStorage.noStoreSerializer.parseBytes(nodeBytes)
node match {
case in: ProxyInternalNode[DigestType] if level == manifestDepth =>
dumpSubtree(Digest32 @@@ in.leftLabel)
dumpSubtree(Digest32 @@@ in.rightLabel)
case in: ProxyInternalNode[DigestType] =>
manifestLoop(in.leftLabel, level + 1, manifestBuilder)
manifestLoop(in.rightLabel, level + 1, manifestBuilder)
case _ =>
// do nothing for a leaf
}
}
val rootNodeLabel = dbReader.get(topNodeHashKey)
val rootNodeHeight = Ints.fromByteArray(dbReader.get(topNodeHeightKey)).toByte
require(rootNodeLabel.sameElements(expectedRootHash), "Root node hash changed")
val manifestBuilder = mutable.ArrayBuilder.make[Byte]()
manifestBuilder.sizeHint(200000)
manifestBuilder += rootNodeHeight
manifestBuilder += manifestDepth
manifestLoop(rootNodeLabel, level = 1, manifestBuilder)
val manifestBytes = manifestBuilder.result()
dumpStorage.insert(rootNodeLabel, manifestBytes)
rootNodeLabel
}
}
}
object VersionedLDBAVLStorage {
private[batch] val topNodeHashKey: Array[Byte] = Array.fill(StateTreeParameters.labelSize)(123: Byte)
private[batch] val topNodeHeightKey: Array[Byte] = Array.fill(StateTreeParameters.labelSize)(124: Byte)
// prefixes used to encode node type (internal or leaf) in database
private[batch] val InternalNodePrefix: Byte = 0: Byte
private[batch] val LeafPrefix: Byte = 1: Byte
/**
* Node serializer which is restoring nodes with only children labels stored in, without possibility to restore
* children automatically when requested
*/
val noStoreSerializer = new ProverNodeSerializer(store = null)
/**
* Fetch tree node from database by its database id (hash of node contents)
*
* @param dbKey - database key node of interest is stored under (hash of the node)
* @param store - database
* @return node read from the database (or throws exception if there is no such node), in case of internal node it
* returns pruned internal node, so a node where left and right children not stored, only their hashes
*/
def fetch(dbKey: ADKey)(store: LDBVersionedStore): ProverNodes[DigestType] = {
val bytes = store(dbKey)
val node = new ProverNodeSerializer(store).parseBytes(bytes)
node.isNew = false
node
}
/**
* Node label (hash). Used as database key for node contents
*/
private[batch] def nodeLabel(node: ProverNodes[DigestType]): Array[Byte] = node.label
/**
* Reconstruct versioned AVL+ tree storage from dump
* @param manifest - tree started from root node cut at fixed depth
* @param chunks - iterator (lazy sequence) of subtrees under manifest tree
* @param additionalData - additional fields (metadata) to write into database
* @param store - db to store restored tree to
* @return db with reconstructed tree written to, or error
*/
def recreate(manifest: BatchAVLProverManifest[DigestType],
chunks: Iterator[BatchAVLProverSubtree[DigestType]],
additionalData: Iterator[(Array[Byte], Array[Byte])],
store: LDBVersionedStore): Try[VersionedLDBAVLStorage] = {
//todo: the function below copy-pasted from BatchAVLProver, eliminate boilerplate?
def idCollector(node: ProverNodes[DigestType]): Iterator[(Array[Byte], Array[Byte])] = {
val pair: (Array[Byte], Array[Byte]) = (nodeLabel(node), noStoreSerializer.toBytes(node))
node match {
case n: ProxyInternalNode[DigestType] if n.isEmpty =>
Iterator(pair)
case i: InternalProverNode[DigestType] =>
Iterator(pair) ++ idCollector(i.left) ++ idCollector(i.right)
case _: ProverLeaf[DigestType] =>
Iterator(pair)
}
}
val rootNode = manifest.root
val rootNodeHeight = manifest.rootHeight
val digestWrapper = VersionedLDBAVLStorage.digest(rootNode.label, rootNodeHeight)
val indices = Iterator(topNodeHashKey -> nodeLabel(rootNode), topNodeHeightKey -> Ints.toByteArray(rootNodeHeight))
val nodesIterator = idCollector(manifest.root) ++
chunks.flatMap(subtree => idCollector(subtree.subtreeTop))
store.update(digestWrapper, toRemove = Nil, toUpdate = indices ++ nodesIterator ++ additionalData).map { _ =>
new VersionedLDBAVLStorage(store)
}
}
/**
* Calculate tree digest, given root node label(hash) and root node height, by appending height to the hash
*/
def digest[D <: hash.Digest](rootNodeLabel: D, rootNodeHeight: Int): ADDigest = {
assert(rootNodeHeight >= 0 && rootNodeHeight < 256)
// rootNodeHeight should never be more than 255, so the toByte conversion is safe (though it may cause an incorrect
// sign on the signed byte if rootHeight>127, but we handle that case correctly on decoding the byte back to int in the
// verifier, by adding 256 if it's negative).
// The reason rootNodeHeight should never be more than 255 is that if height is more than 255,
// then the AVL tree has at least 2^{255/1.4405} = 2^177 leaves, which is more than the number of atoms on planet Earth.
ADDigest @@ (rootNodeLabel :+ rootNodeHeight.toByte)
}
/**
* splits 33-byte digest into 32-bytes hash and 1-byte tree height
*/
def splitDigest(digest: ADDigest): (Array[Byte], Byte) = {
digest.dropRight(1) -> digest.last
}
}