Update BatchBuilder & TxSelector to new StateDB

batchbuilder/batchbuilder.go

@@ -1,13 +1,18 @@
 package batchbuilder
 
 import (
+	"encoding/binary"
 	"math/big"
 
 	ethCommon "github.com/ethereum/go-ethereum/common"
 	"github.com/hermeznetwork/hermez-node/common"
 	"github.com/hermeznetwork/hermez-node/db/statedb"
+	"github.com/iden3/go-merkletree/db"
 )
 
+// KEYIDX is used as key in the db to store the current Idx
+var KEYIDX = []byte("idx")
+
 // ConfigCircuit contains the circuit configuration
 type ConfigCircuit struct {
 	TxsMax       uint64
@@ -31,14 +36,13 @@ type ConfigBatch struct {
 
 // NewBatchBuilder constructs a new BatchBuilder, and executes the bb.Reset
 // method
-func NewBatchBuilder(synchronizerStateDB *statedb.StateDB, configCircuits []ConfigCircuit, batchNum uint64, idx, nLevels uint64) (*BatchBuilder, error) {
+func NewBatchBuilder(dbpath string, synchronizerStateDB *statedb.StateDB, configCircuits []ConfigCircuit, batchNum uint64, nLevels uint64) (*BatchBuilder, error) {
-	localStateDB, err := statedb.NewLocalStateDB(synchronizerStateDB, true, int(nLevels))
+	localStateDB, err := statedb.NewLocalStateDB(dbpath, synchronizerStateDB, true, int(nLevels))
 	if err != nil {
 		return nil, err
 	}
 
 	bb := BatchBuilder{
-		idx:            idx,
 		localStateDB:   localStateDB,
 		configCircuits: configCircuits,
 	}
@@ -52,12 +56,17 @@ func NewBatchBuilder(synchronizerStateDB *statedb.StateDB, configCircuits []Conf
 // copy of the rollup state from the Synchronizer at that `batchNum`, otherwise
 // it can just roll back the internal copy.
 func (bb *BatchBuilder) Reset(batchNum uint64, fromSynchronizer bool) error {
+	if batchNum == 0 {
+		bb.idx = 0
+		return nil
+	}
 	err := bb.localStateDB.Reset(batchNum, fromSynchronizer)
 	if err != nil {
 		return err
 	}
-	// bb.idx = idx // TODO idx will be obtained from the statedb reset
+	// idx is obtained from the statedb reset
-	return nil
+	bb.idx, err = bb.getIdx()
+	return err
 }
 
 // BuildBatch takes the transactions and returns the common.ZKInputs of the next batch
@@ -139,8 +148,8 @@ func (bb *BatchBuilder) processL1Tx(tx common.L1Tx) error {
 	return nil
 }
 
-// applyCreateAccount creates a new account in the account of the depositer, it stores
+// applyCreateAccount creates a new account in the account of the depositer, it
-// the deposit value
+// stores the deposit value
 func (bb *BatchBuilder) applyCreateAccount(tx common.L1Tx) error {
 	account := &common.Account{
 		TokenID:   tx.TokenID,
@@ -156,7 +165,7 @@ func (bb *BatchBuilder) applyCreateAccount(tx common.L1Tx) error {
 	}
 
 	bb.idx = bb.idx + 1
-	return nil
+	return bb.setIdx(bb.idx)
 }
 
 // applyDeposit updates the balance in the account of the depositer, if
@@ -225,3 +234,31 @@ func (bb *BatchBuilder) applyTransfer(tx common.Tx) error {
 
 	return nil
 }
+
+// getIdx returns the stored Idx from the localStateDB, which is the last Idx used for an Account in the localStateDB.
+func (bb *BatchBuilder) getIdx() (uint64, error) {
+	idxBytes, err := bb.localStateDB.DB().Get(KEYIDX)
+	if err == db.ErrNotFound {
+		return 0, nil
+	}
+	if err != nil {
+		return 0, err
+	}
+	idx := binary.LittleEndian.Uint64(idxBytes[:8])
+	return idx, nil
+}
+
+// setIdx stores Idx in the localStateDB
+func (bb *BatchBuilder) setIdx(idx uint64) error {
+	tx, err := bb.localStateDB.DB().NewTx()
+	if err != nil {
+		return err
+	}
+	var idxBytes [8]byte
+	binary.LittleEndian.PutUint64(idxBytes[:], idx)
+	tx.Put(KEYIDX, idxBytes[:])
+	if err := tx.Commit(); err != nil {
+		return err
+	}
+	return nil
+}

batchbuilder/batchbuilder_test.go

@@ -14,10 +14,12 @@ func TestBatchBuilder(t *testing.T) {
 	dir, err := ioutil.TempDir("", "tmpdb")
 	require.Nil(t, err)
 
-	synchDB, err := statedb.NewStateDB(dir, false, false, 0)
+	synchDB, err := statedb.NewStateDB(dir, false, 0)
 	assert.Nil(t, err)
 
-	bb, err := NewBatchBuilder(synchDB, nil, 0, 0, 32)
+	bbDir, err := ioutil.TempDir("", "tmpBatchBuilderDB")
+	require.Nil(t, err)
+	bb, err := NewBatchBuilder(bbDir, synchDB, nil, 0, 32)
 	assert.Nil(t, err)
 	fmt.Println(bb)
 }

db/statedb/statedb.go

@@ -23,8 +23,10 @@ var ErrAccountAlreadyExists = errors.New("Can not CreateAccount because Account
 // KEYCURRENTBATCH is used as key in the db to store the current BatchNum
 var KEYCURRENTBATCH = []byte("currentbatch")
 
-// STATEDBPATH defines the subpath of the StateDB
+// PATHSTATEDB defines the subpath of the StateDB
-const STATEDBPATH = "/statedb"
+const PATHSTATEDB = "/statedb"
+const PATHBATCHNUM = "/BatchNum"
+const PATHCURRENT = "/current"
 
 // StateDB represents the StateDB object
 type StateDB struct {
@@ -39,7 +41,7 @@ type StateDB struct {
 func NewStateDB(path string, withMT bool, nLevels int) (*StateDB, error) {
 	var sto *pebble.PebbleStorage
 	var err error
-	sto, err = pebble.NewPebbleStorage(path+STATEDBPATH+"/current", false)
+	sto, err = pebble.NewPebbleStorage(path+PATHSTATEDB+PATHCURRENT, false)
 	if err != nil {
 		return nil, err
 	}
@@ -53,7 +55,7 @@ func NewStateDB(path string, withMT bool, nLevels int) (*StateDB, error) {
 	}
 
 	sdb := &StateDB{
-		path: path + STATEDBPATH,
+		path: path + PATHSTATEDB,
 		db:   sto,
 		mt:   mt,
 	}
@@ -100,14 +102,17 @@ func (s *StateDB) setCurrentBatch() error {
 	return nil
 }
 
-// CheckPointAt does a checkpoint at the given batchNum in the defined path
+// MakeCheckpoint does a checkpoint at the given batchNum in the defined path. Internally this advances & stores the current BatchNum, and then stores a Checkpoint of the current state of the StateDB.
 func (s *StateDB) MakeCheckpoint() error {
 	// advance currentBatch
 	s.currentBatch++
 
-	checkpointPath := s.path + "/BatchNum" + strconv.Itoa(int(s.currentBatch))
+	checkpointPath := s.path + PATHBATCHNUM + strconv.Itoa(int(s.currentBatch))
 
-	s.setCurrentBatch()
+	err := s.setCurrentBatch()
+	if err != nil {
+		return err
+	}
 
 	// if checkpoint BatchNum already exist in disk, delete it
 	if _, err := os.Stat(checkpointPath); !os.IsNotExist(err) {
@@ -118,7 +123,7 @@ func (s *StateDB) MakeCheckpoint() error {
 	}
 
 	// execute Checkpoint
-	err := s.db.Pebble().Checkpoint(checkpointPath)
+	err = s.db.Pebble().Checkpoint(checkpointPath)
 	if err != nil {
 		return err
 	}
@@ -128,7 +133,7 @@ func (s *StateDB) MakeCheckpoint() error {
 
 // DeleteCheckpoint removes if exist the checkpoint of the given batchNum
 func (s *StateDB) DeleteCheckpoint(batchNum uint64) error {
-	checkpointPath := s.path + "/BatchNum" + strconv.Itoa(int(batchNum))
+	checkpointPath := s.path + PATHBATCHNUM + strconv.Itoa(int(batchNum))
 
 	if _, err := os.Stat(checkpointPath); os.IsNotExist(err) {
 		return fmt.Errorf("Checkpoint with batchNum %d does not exist in DB", batchNum)
@@ -142,8 +147,8 @@ func (s *StateDB) DeleteCheckpoint(batchNum uint64) error {
 // those checkpoints will remain in the storage, and eventually will be
 // deleted when MakeCheckpoint overwrites them.
 func (s *StateDB) Reset(batchNum uint64) error {
-	checkpointPath := s.path + "/BatchNum" + strconv.Itoa(int(batchNum))
+	checkpointPath := s.path + PATHBATCHNUM + strconv.Itoa(int(batchNum))
-	currentPath := s.path + "/current"
+	currentPath := s.path + PATHCURRENT
 
 	// remove 'current'
 	err := os.RemoveAll(currentPath)
@@ -311,9 +316,9 @@ func NewLocalStateDB(path string, synchronizerDB *StateDB, withMT bool, nLevels
 // gets the state from LocalStateDB.synchronizerStateDB for the given batchNum. If fromSynchronizer is false, get the state from LocalStateDB checkpoints.
 func (l *LocalStateDB) Reset(batchNum uint64, fromSynchronizer bool) error {
 
-	synchronizerCheckpointPath := l.synchronizerStateDB.path + "/BatchNum" + strconv.Itoa(int(batchNum))
+	synchronizerCheckpointPath := l.synchronizerStateDB.path + PATHBATCHNUM + strconv.Itoa(int(batchNum))
-	checkpointPath := l.path + "/BatchNum" + strconv.Itoa(int(batchNum))
+	checkpointPath := l.path + PATHBATCHNUM + strconv.Itoa(int(batchNum))
-	currentPath := l.path + "/current"
+	currentPath := l.path + PATHCURRENT
 
 	if fromSynchronizer {
 		// use checkpoint from SynchronizerStateDB

db/statedb/statedb_test.go

@@ -251,9 +251,12 @@ func TestCheckpoints(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, uint64(5), cb)
 
-	// printCheckpoints(t, sdb.path)
+	debug := false
-	// printCheckpoints(t, ldb.path)
+	if debug {
-	// printCheckpoints(t, ldb2.path)
+		printCheckpoints(t, sdb.path)
+		printCheckpoints(t, ldb.path)
+		printCheckpoints(t, ldb2.path)
+	}
 }
 
 func printCheckpoints(t *testing.T, path string) {

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/dghubble/sling v1.3.0
 	github.com/ethereum/go-ethereum v1.9.17
 	github.com/gobuffalo/packr/v2 v2.8.0
+	github.com/iden3/go-iden3-core v0.0.8
 	github.com/iden3/go-iden3-crypto v0.0.6-0.20200806115047-327a8175d6eb
 	github.com/iden3/go-merkletree v0.0.0-20200815144208-1f1bd54b93ae
 	github.com/jinzhu/gorm v1.9.15

txselector/txselector.go

@@ -35,8 +35,8 @@ type TxSelector struct {
 }
 
 // NewTxSelector returns a *TxSelector
-func NewTxSelector(synchronizerStateDB *statedb.StateDB, l2 *l2db.L2DB, maxL1UserTxs, maxL1OperatorTxs, maxTxs uint64) (*TxSelector, error) {
+func NewTxSelector(dbpath string, synchronizerStateDB *statedb.StateDB, l2 *l2db.L2DB, maxL1UserTxs, maxL1OperatorTxs, maxTxs uint64) (*TxSelector, error) {
-	localAccountsDB, err := statedb.NewLocalStateDB(synchronizerStateDB, false, 0) // without merkletree
+	localAccountsDB, err := statedb.NewLocalStateDB(dbpath, synchronizerStateDB, false, 0) // without merkletree
 	if err != nil {
 		return nil, err
 	}

txselector/txselector_test.go

@@ -104,12 +104,14 @@ func initTestDB(l2 *l2db.L2DB, sdb *statedb.StateDB) *mock.MockDB {
 func TestGetL2TxSelection(t *testing.T) {
 	dir, err := ioutil.TempDir("", "tmpdb")
 	require.Nil(t, err)
-	sdb, err := statedb.NewStateDB(dir, false, false, 0)
+	sdb, err := statedb.NewStateDB(dir, false, 0)
 	assert.Nil(t, err)
 	testL2DB := &l2db.L2DB{}
 	// initTestDB(testL2DB, sdb)
 
-	txsel, err := NewTxSelector(sdb, testL2DB, 3, 3, 3)
+	txselDir, err := ioutil.TempDir("", "tmpTxSelDB")
+	require.Nil(t, err)
+	txsel, err := NewTxSelector(txselDir, sdb, testL2DB, 3, 3, 3)
 	assert.Nil(t, err)
 	fmt.Println(txsel)