Elizabeth/trie db (#107)

* first implementation of writing to DB

* implement batch write to DB

* implement dirty flag for nodes; add Close to database interface

* update gitignore

* change test to appease golangcibot

* change test again to appease golangcibot

* update gitignore and Iterable

.gitignore

@@ -24,4 +24,6 @@
 # Sublime
 *.sublime*
 
-/bin
+trie/gossamer_data
+
+/bin

polkadb/database.go

@@ -30,8 +30,8 @@ type BadgerDB struct {
 	db   *badger.DB
 }
 
-// Iterator struct contains a transaction, iterator and context fields released, initialized
-type Iterate struct {
+// Iterable struct contains a transaction, iterator and context fields released, initialized
+type Iterable struct {
 	txn      *badger.Txn
 	iter     *badger.Iterator
 	released bool
@@ -148,11 +148,11 @@ func (db *BadgerDB) Close() {
 }
 
 // NewIterator returns a new iterator within the Iterator struct along with a new transaction
-func (db *BadgerDB) NewIterator() Iterate {
+func (db *BadgerDB) NewIterator() Iterable {
 	txn := db.db.NewTransaction(false)
 	opts := badger.DefaultIteratorOptions
 	iter := txn.NewIterator(opts)
-	return Iterate{
+	return Iterable{
 		txn:      txn,
 		iter:     iter,
 		released: false,
@@ -161,20 +161,20 @@ func (db *BadgerDB) NewIterator() Iterate {
 }
 
 // Release closes the iterator, discards the created transaction and sets released value to true
-func (i *Iterate) Release() {
+func (i *Iterable) Release() {
 	i.iter.Close()
 	i.txn.Discard()
 	i.released = true
 }
 
 // Released returns the boolean indicating whether the iterator and transaction was successfully released
-func (i *Iterate) Released() bool {
+func (i *Iterable) Released() bool {
 	return i.released
 }
 
 // Next rewinds the iterator to the zero-th position if uninitialized, and then will advance the iterator by one
 // returns bool to ensure access to the item
-func (i *Iterate) Next() bool {
+func (i *Iterable) Next() bool {
 	if !i.init {
 		i.iter.Rewind()
 		i.init = true
@@ -185,12 +185,12 @@ func (i *Iterate) Next() bool {
 
 // Seek will look for the provided key if present and go to that position. If
 // absent, it would seek to the next smallest key
-func (i *Iterate) Seek(key []byte) {
+func (i *Iterable) Seek(key []byte) {
 	i.iter.Seek(snappy.Encode(nil, key))
 }
 
 // Key returns an item key
-func (i *Iterate) Key() []byte {
+func (i *Iterable) Key() []byte {
 	ret, err := snappy.Decode(nil, i.iter.Item().Key())
 	if err != nil {
 		log.Printf("%+v", errors.Wrap(err, "key retrieval error"))
@@ -199,7 +199,7 @@ func (i *Iterate) Key() []byte {
 }
 
 // Value returns a copy of the value of the item
-func (i *Iterate) Value() []byte {
+func (i *Iterable) Value() []byte {
 	val, err := i.iter.Item().ValueCopy(nil)
 	if err != nil {
 		log.Printf("%+v", errors.Wrap(err, "value retrieval error"))
@@ -284,6 +284,10 @@ func (dt *table) Del(key []byte) error {
 	return dt.db.Del(append([]byte(dt.prefix), key...))
 }
 
+func (dt *table) Close() {
+	dt.db.Close()
+}
+
 // NewTableBatch returns a Batch object which prefixes all keys with a given string.
 func NewTableBatch(db Database, prefix string) Batch {
 	return &tableBatch{db.NewBatch(), prefix}

polkadb/interface.go

@@ -28,6 +28,7 @@ type Database interface {
 	Has(key []byte) (bool, error)
 	Del(key []byte) error
 	NewBatch() Batch
+	Close()
 }
 
 // Batch is a write-only operation
@@ -50,5 +51,5 @@ type Iterator interface {
 
 // Iteratee wraps the NewIterator methods of BadgerDB
 type Iteratee interface {
-	NewIterator() Iterate
+	NewIterator() Iterable
 }

polkadb/memoryDB.go

@@ -75,7 +75,7 @@ func (db *MemDatabase) Keys() [][]byte {
 	return keys
 }
 
-// Delete removes the key from the mapping
+// Del removes the key from the mapping
 func (db *MemDatabase) Del(key []byte) error {
 	db.lock.Lock()
 	defer db.lock.Unlock()
@@ -84,6 +84,11 @@ func (db *MemDatabase) Del(key []byte) error {
 	return nil
 }
 
+// Close ...
+func (db *MemDatabase) Close() {
+	// do nothing
+}
+
 // NewBatch ...
 func (db *MemDatabase) NewBatch() Batch {
 	return nil

trie/database.go

@@ -18,11 +18,75 @@ package trie
 
 import (
 	"github.com/ChainSafe/gossamer/polkadb"
-	//"sync"
+	"sync"
 )
 
 // Database is a wrapper around a polkadb
 type Database struct {
-	db polkadb.Database
-	//lock sync.RWMutex
+	db     polkadb.Database
+	batch  polkadb.Batch
+	lock   sync.RWMutex
+	hasher *Hasher
+}
+
+// WriteToDB writes the trie to the underlying database batch writer
+// Stores the merkle value of the node as the key and the encoded node as the value
+// This does not actually write to the db, just to the batch writer
+// Commit must be called afterwards to finish writing to the db
+func (t *Trie) WriteToDB() error {
+	t.db.batch = t.db.db.NewBatch()
+	return t.writeToDB(t.root)
+}
+
+// writeToDB recursively attempts to write each node in the trie to the db batch writer
+func (t *Trie) writeToDB(n node) error {
+	_, err := t.writeNodeToDB(n)
+	if err != nil {
+		return err
+	}
+
+	switch n := n.(type) {
+	case *branch:
+		for _, child := range n.children {
+			if child != nil {
+				err = t.writeToDB(child)
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
+
+	return nil
+}
+
+// writeNodeToDB returns true if node is written to db batch writer, false otherwise
+// if node is clean, it will not attempt to be written to the db
+// otherwise if it's dirty, try to write it to db
+func (t *Trie) writeNodeToDB(n node) (bool, error) {
+	if !n.isDirty() {
+		return false, nil
+	}
+
+	encRoot, err := Encode(n)
+	if err != nil {
+		return false, err
+	}
+
+	hash, err := t.db.hasher.Hash(n)
+	if err != nil {
+		return false, err
+	}
+
+	t.db.lock.Lock()
+	err = t.db.batch.Put(hash, encRoot)
+	t.db.lock.Unlock()
+
+	n.setDirty(false)
+	return true, err
+}
+
+// Commit writes the contents of the db's batch writer to the db
+func (t *Trie) Commit() error {
+	return t.db.batch.Write()
 }

trie/database_test.go

@@ -0,0 +1,90 @@
+package trie
+
+import (
+	"bytes"
+	"testing"
+
+	db "github.com/ChainSafe/gossamer/polkadb"
+)
+
+func newTrie() (*Trie, error) {
+	hasher, err := newHasher()
+	if err != nil {
+		return nil, err
+	}
+
+	db, err := db.NewBadgerDB("./gossamer_data")
+	if err != nil {
+		return nil, err
+	}
+
+	trie := &Trie{
+		db: &Database{
+			db:     db,
+			hasher: hasher,
+		},
+		root: nil,
+	}
+
+	trie.db.batch = trie.db.db.NewBatch()
+
+	return trie, nil
+}
+
+func TestWriteToDB(t *testing.T) {
+	trie, err := newTrie()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rt := generateRandTest(20000)
+	var val []byte
+	for _, test := range rt {
+		err = trie.Put(test.key, test.value)
+		if err != nil {
+			t.Errorf("Fail to put with key %x and value %x: %s", test.key, test.value, err.Error())
+		}
+
+		val, err = trie.Get(test.key)
+		if err != nil {
+			t.Errorf("Fail to get key %x: %s", test.key, err.Error())
+		} else if !bytes.Equal(val, test.value) {
+			t.Errorf("Fail to get key %x with value %x: got %x", test.key, test.value, val)
+		}
+	}
+
+	err = trie.WriteToDB()
+	if err != nil {
+		t.Errorf("Fail: could not write to batch writer: %s", err)
+	}
+
+	err = trie.Commit()
+	if err != nil {
+		t.Errorf("Fail: could not commit (batch write) to DB: %s", err)
+	}
+
+	trie.db.db.Close()
+}
+
+func TestWriteDirty(t *testing.T) {
+	trie, err := newTrie()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	dirtyNode := &leaf{key: generateRandBytes(10), value: generateRandBytes(10), dirty: true}
+	written, err := trie.writeNodeToDB(dirtyNode)
+	if err != nil {
+		t.Errorf("Fail: could not write to db: %s", err)
+	} else if !written {
+		t.Errorf("Fail: did not write dirty node to db")
+	}
+
+	cleanNode := &leaf{key: generateRandBytes(10), value: generateRandBytes(10), dirty: false}
+	written, err = trie.writeNodeToDB(cleanNode)
+	if err != nil {
+		t.Errorf("Fail: could not write to db: %s", err)
+	} else if written {
+		t.Errorf("Fail: wrote clean node to db")
+	}
+}

trie/hash.go

@@ -24,6 +24,7 @@ import (
 	"golang.org/x/crypto/blake2s"
 )
 
+// Hasher is a wrapper around a hash function
 type Hasher struct {
 	hash hash.Hash
 }

trie/node.go

@@ -25,17 +25,21 @@ import (
 
 type node interface {
 	Encode() ([]byte, error)
+	isDirty() bool
+	setDirty(dirty bool)
 }
 
 type (
 	branch struct {
 		key      []byte // partial key
 		children [16]node
 		value    []byte
+		dirty    bool
 	}
 	leaf struct {
 		key   []byte // partial key
 		value []byte
+		dirty bool
 	}
 )
 
@@ -50,6 +54,22 @@ func (b *branch) childrenBitmap() uint16 {
 	return bitmap
 }
 
+func (l *leaf) isDirty() bool {
+	return l.dirty
+}
+
+func (b *branch) isDirty() bool {
+	return b.dirty
+}
+
+func (l *leaf) setDirty(dirty bool) {
+	l.dirty = dirty
+}
+
+func (b *branch) setDirty(dirty bool) {
+	b.dirty = dirty
+}
+
 // Encode is the high-level function wrapping the encoding for different node types
 // encoding has the following format:
 // NodeHeader | Extra partial key length | Partial Key | Value