Browse Source

Implement addBatchInDisk for big trees

Implement addBatchInDisk for big trees, which does not puts the tree in
memory, and works directly over the db data, parallelizing for each CPU.
master
arnaucube 2 years ago
parent
commit
6cf1e58d9f
10 changed files with 659 additions and 104 deletions
  1. +119
    -58
      addbatch_test.go
  2. +2
    -1
      circomproofs_test.go
  3. +2
    -2
      go.mod
  4. +288
    -5
      go.sum
  5. +3
    -2
      helpers_test.go
  6. +2
    -1
      testvectors/circom/go-data-generator/generator_test.go
  7. +202
    -5
      tree.go
  8. +35
    -25
      tree_test.go
  9. +3
    -3
      vt.go
  10. +3
    -2
      vt_test.go

+ 119
- 58
addbatch_test.go

@ -11,7 +11,9 @@ import (
"time"
qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
"go.vocdoni.io/dvote/db/pebbledb"
)
var debug = true
@ -37,12 +39,12 @@ func debugTime(descr string, time1, time2 time.Duration) {
}
func testInit(c *qt.C, n int) (*Tree, *Tree) {
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -68,7 +70,7 @@ func TestAddBatchTreeEmpty(t *testing.T) {
nLeafs := 1024
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -91,7 +93,7 @@ func TestAddBatchTreeEmpty(t *testing.T) {
}
time1 := time.Since(start)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -99,7 +101,7 @@ func TestAddBatchTreeEmpty(t *testing.T) {
tree2.dbgInit()
start = time.Now()
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
time2 := time.Since(start)
if debug {
@ -107,7 +109,7 @@ func TestAddBatchTreeEmpty(t *testing.T) {
printTestContext(" ", nLeafs, "Poseidon", "memory")
tree2.dbg.print(" ")
}
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree, tree2)
@ -118,7 +120,7 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
nLeafs := 1027
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -133,7 +135,7 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
}
}
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -146,9 +148,9 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
keys = append(keys, k)
values = append(values, v)
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree, tree2)
@ -165,13 +167,13 @@ func randomBytes(n int) []byte {
func TestAddBatchTestVector1(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -198,20 +200,20 @@ func TestAddBatchTestVector1(t *testing.T) {
}
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
// 2nd test vectors
database1, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err = badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err = NewTree(database1, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
database2, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err = badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err = NewTree(database2, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -242,9 +244,9 @@ func TestAddBatchTestVector1(t *testing.T) {
}
}
indexes, err = tree2.AddBatch(keys, values)
invalids, err = tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
}
@ -253,13 +255,13 @@ func TestAddBatchTestVector2(t *testing.T) {
// test vector with unbalanced tree
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -286,9 +288,9 @@ func TestAddBatchTestVector2(t *testing.T) {
}
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
@ -298,13 +300,13 @@ func TestAddBatchTestVector3(t *testing.T) {
// test vector with unbalanced tree
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -331,9 +333,9 @@ func TestAddBatchTestVector3(t *testing.T) {
}
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
@ -347,13 +349,13 @@ func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
nLeafs := 8
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -371,9 +373,9 @@ func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
}
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
}
@ -407,7 +409,7 @@ func TestAddBatchTreeNotEmptyFewLeafs(t *testing.T) {
values = append(values, v)
}
start = time.Now()
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
time2 := time.Since(start)
if debug {
@ -415,7 +417,7 @@ func TestAddBatchTreeNotEmptyFewLeafs(t *testing.T) {
printTestContext(" ", nLeafs, "Poseidon", "memory")
tree2.dbg.print(" ")
}
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
@ -450,7 +452,7 @@ func TestAddBatchTreeNotEmptyEnoughLeafs(t *testing.T) {
values = append(values, v)
}
start = time.Now()
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
time2 := time.Since(start)
if debug {
@ -458,7 +460,7 @@ func TestAddBatchTreeNotEmptyEnoughLeafs(t *testing.T) {
printTestContext(" ", nLeafs, "Poseidon", "memory")
tree2.dbg.print(" ")
}
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
}
@ -495,9 +497,9 @@ func TestAddBatchTreeEmptyRepeatedLeafs(t *testing.T) {
}
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, nRepeatedKeys)
c.Check(len(invalids), qt.Equals, nRepeatedKeys)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
}
@ -527,9 +529,9 @@ func TestAddBatchTreeNotEmptyFewLeafsRepeatedLeafs(t *testing.T) {
}
}
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, initialNLeafs)
c.Assert(len(invalids), qt.Equals, initialNLeafs)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
}
@ -664,7 +666,7 @@ func TestAddBatchTreeNotEmpty(t *testing.T) {
values = append(values, v)
}
start = time.Now()
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
time2 := time.Since(start)
if debug {
@ -672,7 +674,7 @@ func TestAddBatchTreeNotEmpty(t *testing.T) {
printTestContext(" ", nLeafs, "Poseidon", "memory")
tree2.dbg.print(" ")
}
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
@ -697,7 +699,7 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
}
time1 := time.Since(start)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -729,7 +731,7 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
values = append(values, v)
}
start = time.Now()
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
time2 := time.Since(start)
if debug {
@ -737,7 +739,7 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
printTestContext(" ", nLeafs, "Poseidon", "memory")
tree2.dbg.print(" ")
}
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
@ -774,7 +776,7 @@ func TestAddBatchBench(t *testing.T) {
func benchAdd(t *testing.T, ks, vs [][]byte) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -794,7 +796,7 @@ func benchAdd(t *testing.T, ks, vs [][]byte) {
func benchAddBatch(t *testing.T, ks, vs [][]byte) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -827,7 +829,7 @@ func TestDbgStats(t *testing.T) {
}
// 1
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -841,7 +843,7 @@ func TestDbgStats(t *testing.T) {
}
// 2
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -854,7 +856,7 @@ func TestDbgStats(t *testing.T) {
c.Assert(len(invalids), qt.Equals, 0)
// 3
database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database3, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree3, err := NewTree(database3, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -889,7 +891,7 @@ func TestLoadVT(t *testing.T) {
nLeafs := 1024
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -902,9 +904,9 @@ func TestLoadVT(t *testing.T) {
keys = append(keys, k)
values = append(values, v)
}
indexes, err := tree.AddBatch(keys, values)
invalids, err := tree.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
rTx := tree.db.ReadTx()
defer rTx.Discard()
@ -925,7 +927,7 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
nLeafs := 1024
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -946,29 +948,29 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
}
}
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree2.db.Close() //nolint:errcheck
tree2.dbgInit()
indexes, err := tree2.AddBatch(keys, values)
invalids, err := tree2.AddBatch(keys, values)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree, tree2)
// use tree3 to add nil value array
database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database3, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree3, err := NewTree(database3, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
defer tree3.db.Close() //nolint:errcheck
indexes, err = tree3.AddBatch(keys, nil)
invalids, err = tree3.AddBatch(keys, nil)
c.Assert(err, qt.IsNil)
c.Check(len(indexes), qt.Equals, 0)
c.Check(len(invalids), qt.Equals, 0)
checkRoots(c, tree, tree3)
kAux, proofV, siblings, existence, err := tree2.GenProof(keys[9])
@ -1003,6 +1005,65 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
c.Check(verif, qt.IsFalse)
}
func TestAddBatchThresholdInDisk(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
defer tree1.db.Close() //nolint:errcheck
database2, err := pebbledb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
defer tree2.db.Close() //nolint:errcheck
database3, err := pebbledb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree3, err := NewTree(database3, 256, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
defer tree3.db.Close() //nolint:errcheck
// customize thresholdNLeafs for the test
thresholdNLeafs = 1024
var keys, values [][]byte
for i := 0; i < 3*thresholdNLeafs; i++ {
k := randomBytes(32)
v := randomBytes(32)
if err := tree1.Add(k, v); err != nil {
t.Fatal(err)
}
if i < thresholdNLeafs+1 {
if err := tree2.Add(k, v); err != nil {
t.Fatal(err)
}
}
// store for later addition through AddBatch
keys = append(keys, k)
values = append(values, v)
}
invalids, err := tree2.AddBatch(keys[thresholdNLeafs+1:], values[thresholdNLeafs+1:])
c.Assert(err, qt.IsNil)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree2)
// call directly the tree3.addBatchInDisk to ensure that is tested
wTx := tree3.db.WriteTx()
defer wTx.Discard()
invalids, err = tree3.addBatchInDisk(wTx, keys, values)
c.Assert(err, qt.IsNil)
err = wTx.Commit()
c.Assert(err, qt.IsNil)
c.Check(len(invalids), qt.Equals, 0)
// check that both trees roots are equal
checkRoots(c, tree1, tree3)
}
// TODO test adding batch with multiple invalid keys
// TODO for tests of AddBatch, if the root does not match the Add root, bulk
// all the leafs of both trees into a log file to later be able to debug and

+ 2
- 1
circomproofs_test.go

@ -6,12 +6,13 @@ import (
"testing"
qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
)
func TestCircomVerifierProof(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 4, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)

+ 2
- 2
go.mod

@ -5,6 +5,6 @@ go 1.16
require (
github.com/frankban/quicktest v1.13.0
github.com/iden3/go-iden3-crypto v0.0.6-0.20210308142348-8f85683b2cef
go.vocdoni.io/dvote v1.0.4-0.20210806163627-9494efbc5382
golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97
go.vocdoni.io/dvote v1.0.4-0.20211025120558-83c64f440044
golang.org/x/crypto v0.0.0-20210920023735-84f357641f63
)

+ 288
- 5
go.sum
File diff suppressed because it is too large
View File


+ 3
- 2
helpers_test.go

@ -9,6 +9,7 @@ import (
"time"
qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
)
@ -87,12 +88,12 @@ func TestReadTreeDBG(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)

+ 2
- 1
testvectors/circom/go-data-generator/generator_test.go

@ -8,12 +8,13 @@ import (
qt "github.com/frankban/quicktest"
"github.com/vocdoni/arbo"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
)
func TestGenerator(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := arbo.NewTree(database, 4, arbo.HashFunctionPoseidon)
c.Assert(err, qt.IsNil)

+ 202
- 5
tree.go

@ -17,6 +17,7 @@ import (
"fmt"
"io"
"math"
"runtime"
"sync"
"go.vocdoni.io/dvote/db"
@ -45,6 +46,12 @@ const (
)
var (
// thresholdNLeafs defines the threshold number of leafs in the tree
// that determines if AddBatch will work in memory or in disk. Is
// defined as a var in order to have the ability to modify it for
// testing purposes.
thresholdNLeafs = 1024 // TODO define a reasonable value
dbKeyRoot = []byte("root")
dbKeyNLeafs = []byte("nleafs")
emptyValue = []byte{0}
@ -197,11 +204,6 @@ func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
return nil, ErrSnapshotNotEditable
}
vt, err := t.loadVT(wTx)
if err != nil {
return nil, err
}
e := []byte{}
// equal the number of keys & values
if len(keys) > len(values) {
@ -214,6 +216,201 @@ func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
values = values[:len(keys)]
}
nLeafs, err := t.GetNLeafsWithTx(wTx)
if err != nil {
return nil, err
}
if nLeafs > thresholdNLeafs {
return t.addBatchInDisk(wTx, keys, values)
}
return t.addBatchInMemory(wTx, keys, values)
}
func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
nCPU := flp2(runtime.NumCPU())
if nCPU == 1 || len(keys) < nCPU {
var invalids []Invalid
for i := 0; i < len(keys); i++ {
if err := t.AddWithTx(wTx, keys[i], values[i]); err != nil {
invalids = append(invalids, Invalid{i, err})
}
}
return invalids, nil
}
kvs, invalids, err := keysValuesToKvs(t.maxLevels, keys, values)
if err != nil {
return nil, err
}
buckets := splitInBuckets(kvs, nCPU)
root, err := t.RootWithTx(wTx)
if err != nil {
return nil, err
}
l := int(math.Log2(float64(nCPU)))
subRoots, err := t.getSubRootsAtLevel(wTx, root, l+1)
if err != nil {
return nil, err
}
if len(subRoots) != nCPU {
// Already populated Tree but Unbalanced.
// add one key at each bucket, and then continue with the flow
for i := 0; i < len(buckets); i++ {
// add one leaf of the bucket, if there is an error when
// adding the k-v, try to add the next one of the bucket
// (until one is added)
inserted := -1
for j := 0; j < len(buckets[i]); j++ {
if newRoot, err := t.add(wTx, root, 0, buckets[i][j].k, buckets[i][j].v); err == nil {
inserted = j
root = newRoot
break
}
}
// remove the inserted element from buckets[i]
if inserted != -1 {
buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
}
}
subRoots, err = t.getSubRootsAtLevel(wTx, root, l+1)
if err != nil {
return nil, err
}
}
if len(subRoots) != nCPU {
return nil, fmt.Errorf("This error should not be reached."+
" len(subRoots) != nCPU, len(subRoots)=%d, nCPU=%d."+
" Please report it in a new issue:"+
" https://github.com/vocdoni/arbo/issues/new", len(subRoots), nCPU)
}
invalidsInBucket := make([][]Invalid, nCPU)
txs := make([]db.WriteTx, nCPU)
for i := 0; i < nCPU; i++ {
txs[i] = t.db.WriteTx()
err := txs[i].Apply(wTx)
if err != nil {
return nil, err
}
}
var wg sync.WaitGroup
wg.Add(nCPU)
for i := 0; i < nCPU; i++ {
go func(cpu int) {
// use different wTx for each cpu, after once all
// are done, iter over the cpuWTxs and copy their
// content into the main wTx
for j := 0; j < len(buckets[cpu]); j++ {
subRoots[cpu], err = t.add(txs[cpu], subRoots[cpu], l, buckets[cpu][j].k, buckets[cpu][j].v)
if err != nil {
invalidsInBucket[cpu] = append(invalidsInBucket[cpu], Invalid{buckets[cpu][j].pos, err})
}
}
wg.Done()
}(i)
}
wg.Wait()
for i := 0; i < nCPU; i++ {
if err := wTx.Apply(txs[i]); err != nil {
return nil, err
}
txs[i].Discard()
}
for i := 0; i < len(invalidsInBucket); i++ {
invalids = append(invalids, invalidsInBucket[i]...)
}
newRoot, err := t.upFromSubRoots(wTx, subRoots)
if err != nil {
return nil, err
}
// update dbKeyNLeafs
if err := t.SetRootWithTx(wTx, newRoot); err != nil {
return nil, err
}
// update nLeafs
if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
return nil, err
}
return invalids, nil
}
func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) {
// is a method of Tree just to get access to t.hashFunction and
// t.emptyHash.
// go up from subRoots to up, storing nodes in the given WriteTx
// once up at the root, store it in the WriteTx using the dbKeyRoot
if len(subRoots) == 1 {
return subRoots[0], nil
}
var newSubRoots [][]byte
// TODO store the nodes key-values into the wTx
for i := 0; i < len(subRoots); i += 2 {
if bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash) {
// TODO: || (ns[i].typ() == vtLeaf && ns[i+1].typ() == vtEmpty) {
// when both sub nodes are empty, the parent is also empty
// or
// (TODO WIP) when 1st sub node is a leaf but the 2nd is empty, the
// leaf is used as parent
newSubRoots = append(newSubRoots, subRoots[i])
continue
}
// TODO if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtLeaf {
// when 2nd sub node is a leaf but the 1st is empty, the
// leaf is used as 'parent'
k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
if err != nil {
return nil, err
}
// store k-v to db
if err = wTx.Set(k, v); err != nil {
return nil, err
}
newSubRoots = append(newSubRoots, k)
}
return t.upFromSubRoots(wTx, newSubRoots)
}
func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
// go at level l and return each node key, where each node key is the
// subRoot of the subTree that starts there
var subRoots [][]byte
err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
if currLvl == l && !bytes.Equal(k, t.emptyHash) {
subRoots = append(subRoots, k)
}
if currLvl >= l {
return true // to stop the iter from going down
}
return false
})
return subRoots, err
}
func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
vt, err := t.loadVT(wTx)
if err != nil {
return nil, err
}
invalids, err := vt.addBatch(keys, values)
if err != nil {
return nil, err

+ 35
- 25
tree_test.go

@ -11,6 +11,7 @@ import (
qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
"go.vocdoni.io/dvote/db/pebbledb"
)
func checkRootBIString(c *qt.C, tree *Tree, expected string) {
@ -23,11 +24,20 @@ func checkRootBIString(c *qt.C, tree *Tree, expected string) {
func TestDBTx(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
dbBadger, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
testDBTx(c, dbBadger)
dbPebble, err := pebbledb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
testDBTx(c, dbPebble)
}
func testDBTx(c *qt.C, database db.Database) {
wTx := database.WriteTx()
defer wTx.Discard()
_, err = wTx.Get([]byte("a"))
_, err := wTx.Get([]byte("a"))
c.Assert(err, qt.Equals, db.ErrKeyNotFound)
err = wTx.Set([]byte("a"), []byte("b"))
@ -60,7 +70,7 @@ func TestAddTestVectors(t *testing.T) {
}
func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, hashFunc)
c.Assert(err, qt.IsNil)
@ -92,7 +102,7 @@ func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
func TestAddBatch(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -110,7 +120,7 @@ func TestAddBatch(t *testing.T) {
checkRootBIString(c, tree,
"296519252211642170490407814696803112091039265640052570497930797516015811235")
database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err = badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -133,7 +143,7 @@ func TestAddBatch(t *testing.T) {
func TestAddDifferentOrder(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -148,7 +158,7 @@ func TestAddDifferentOrder(t *testing.T) {
}
}
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -173,7 +183,7 @@ func TestAddDifferentOrder(t *testing.T) {
func TestAddRepeatedIndex(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -191,7 +201,7 @@ func TestAddRepeatedIndex(t *testing.T) {
func TestUpdate(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -244,7 +254,7 @@ func TestUpdate(t *testing.T) {
func TestAux(t *testing.T) { // TODO split in proper tests
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -283,7 +293,7 @@ func TestAux(t *testing.T) { // TODO split in proper tests
func TestGet(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -307,7 +317,7 @@ func TestGet(t *testing.T) {
func TestGenProofAndVerify(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -339,7 +349,7 @@ func TestGenProofAndVerify(t *testing.T) {
func TestDumpAndImportDump(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -357,7 +367,7 @@ func TestDumpAndImportDump(t *testing.T) {
e, err := tree1.Dump(nil)
c.Assert(err, qt.IsNil)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -376,7 +386,7 @@ func TestDumpAndImportDump(t *testing.T) {
func TestRWMutex(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -408,7 +418,7 @@ func TestRWMutex(t *testing.T) {
// TODO UPDATE
// func TestSetGetNLeafs(t *testing.T) {
// c := qt.New(t)
// database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
// database, err := badgerdb.New(db.Options{Path: c.TempDir()})
// c.Assert(err, qt.IsNil)
// tree, err := NewTree(database, 100, HashFunctionPoseidon)
// c.Assert(err, qt.IsNil)
@ -459,12 +469,12 @@ func TestRWMutex(t *testing.T) {
func TestAddBatchFullyUsed(t *testing.T) {
c := qt.New(t)
database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -518,7 +528,7 @@ func TestAddBatchFullyUsed(t *testing.T) {
func TestSetRoot(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -574,7 +584,7 @@ func TestSetRoot(t *testing.T) {
func TestSnapshot(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -624,7 +634,7 @@ func TestSnapshot(t *testing.T) {
func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 256, HashFunctionPoseidon)
c.Assert(err, qt.IsNil)
@ -644,7 +654,7 @@ func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
func TestKeyLen(t *testing.T) {
c := qt.New(t)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
// maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13
maxLevels := 100
@ -678,7 +688,7 @@ func TestKeyLen(t *testing.T) {
// expect errors when adding a key bigger than maximum capacity of the
// tree: ceil(maxLevels/8)
maxLevels = 32
database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err = badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err = NewTree(database, maxLevels, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -748,7 +758,7 @@ func TestKeyLen(t *testing.T) {
func TestKeyLenBiggerThan32(t *testing.T) {
c := qt.New(t)
maxLevels := 264
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)
@ -790,7 +800,7 @@ func BenchmarkAdd(b *testing.B) {
func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
c := qt.New(b)
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, 140, hashFunc)
c.Assert(err, qt.IsNil)

+ 3
- 3
vt.go

@ -37,7 +37,7 @@ type kv struct {
v []byte
}
func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, []Invalid, error) {
func keysValuesToKvs(maxLevels int, ks, vs [][]byte) ([]kv, []Invalid, error) {
if len(ks) != len(vs) {
return nil, nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
len(ks), len(vs))
@ -45,7 +45,7 @@ func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, []Invalid, error) {
var invalids []Invalid
var kvs []kv
for i := 0; i < len(ks); i++ {
keyPath, err := keyPathFromKey(p.maxLevels, ks[i])
keyPath, err := keyPathFromKey(maxLevels, ks[i])
if err != nil {
invalids = append(invalids, Invalid{i, err})
continue
@ -101,7 +101,7 @@ func (t *vt) addBatch(ks, vs [][]byte) ([]Invalid, error) {
l := int(math.Log2(float64(nCPU)))
kvs, invalids, err := t.params.keysValuesToKvs(ks, vs)
kvs, invalids, err := keysValuesToKvs(t.params.maxLevels, ks, vs)
if err != nil {
return invalids, err
}

+ 3
- 2
vt_test.go

@ -7,6 +7,7 @@ import (
"testing"
qt "github.com/frankban/quicktest"
"go.vocdoni.io/dvote/db"
"go.vocdoni.io/dvote/db/badgerdb"
)
@ -16,7 +17,7 @@ func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) {
c.Assert(len(keys), qt.Equals, len(values))
// normal tree, to have an expected root value
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, maxLevels, HashFunctionSha256)
c.Assert(err, qt.IsNil)
@ -121,7 +122,7 @@ func TestVirtualTreeAddBatch(t *testing.T) {
}
// normal tree, to have an expected root value
database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
database, err := badgerdb.New(db.Options{Path: c.TempDir()})
c.Assert(err, qt.IsNil)
tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
c.Assert(err, qt.IsNil)

Loading…
Cancel
Save