6 changed files with 297 additions and 5 deletions
Unified View
Diff Options
-
+193 -0db/pebble/pebble.go
-
+57 -0db/pebble/pebble_test.go
-
+0 -1db/test/test.go
-
+2 -1go.mod
-
+42 -0go.sum
-
+3 -3merkletree.go
@ -0,0 +1,193 @@ |
|||||
|
package pebble |
||||
|
|
||||
|
import ( |
||||
|
"bytes" |
||||
|
"encoding/json" |
||||
|
|
||||
|
"github.com/cockroachdb/pebble" |
||||
|
"github.com/iden3/go-merkletree/db" |
||||
|
log "github.com/sirupsen/logrus" |
||||
|
) |
||||
|
|
||||
|
// PebbleStorage implements the db.Storage interface
|
||||
|
type PebbleStorage struct { |
||||
|
pdb *pebble.DB |
||||
|
prefix []byte |
||||
|
} |
||||
|
|
||||
|
// PebbleStorageTx implements the db.Tx interface
|
||||
|
type PebbleStorageTx struct { |
||||
|
// FUTURE currently Tx is using the same strategy than in MemoryDB and
|
||||
|
// LevelDB, in next iteration can be moved to Pebble Batch strategy
|
||||
|
*PebbleStorage |
||||
|
cache db.KvMap |
||||
|
} |
||||
|
|
||||
|
// NewPebbleStorage returns a new PebbleStorage
|
||||
|
func NewPebbleStorage(path string, errorIfMissing bool) (*PebbleStorage, error) { |
||||
|
o := &pebble.Options{ |
||||
|
ErrorIfNotExists: errorIfMissing, |
||||
|
} |
||||
|
rdb, err := pebble.Open(path, o) |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
return &PebbleStorage{rdb, []byte{}}, nil |
||||
|
} |
||||
|
|
||||
|
type storageInfo struct { |
||||
|
KeyCount int |
||||
|
ClaimCount int |
||||
|
} |
||||
|
|
||||
|
// Info implements the method Info of the interface db.Storage
|
||||
|
func (p *PebbleStorage) Info() string { |
||||
|
keycount := 0 |
||||
|
claimcount := 0 |
||||
|
err := p.Iterate(func(key []byte, value []byte) (bool, error) { |
||||
|
if value[0] == byte(1) { |
||||
|
claimcount++ |
||||
|
} |
||||
|
|
||||
|
keycount++ |
||||
|
return true, nil |
||||
|
}) |
||||
|
if err != nil { |
||||
|
return err.Error() |
||||
|
} |
||||
|
json, _ := json.MarshalIndent( |
||||
|
storageInfo{ |
||||
|
KeyCount: keycount, |
||||
|
ClaimCount: claimcount, |
||||
|
}, |
||||
|
"", " ", |
||||
|
) |
||||
|
return string(json) |
||||
|
} |
||||
|
|
||||
|
// WithPrefix implements the method WithPrefix of the interface db.Storage
|
||||
|
func (p *PebbleStorage) WithPrefix(prefix []byte) db.Storage { |
||||
|
return &PebbleStorage{p.pdb, db.Concat(p.prefix, prefix)} |
||||
|
} |
||||
|
|
||||
|
// NewTx implements the method NewTx of the interface db.Storage
|
||||
|
func (p *PebbleStorage) NewTx() (db.Tx, error) { |
||||
|
return &PebbleStorageTx{p, make(db.KvMap)}, nil |
||||
|
} |
||||
|
|
||||
|
// Get retreives a value from a key in the db.Storage
|
||||
|
func (p *PebbleStorage) Get(key []byte) ([]byte, error) { |
||||
|
v, closer, err := p.pdb.Get(db.Concat(p.prefix, key[:])) |
||||
|
if err == pebble.ErrNotFound { |
||||
|
return nil, db.ErrNotFound |
||||
|
} |
||||
|
closer.Close() |
||||
|
return v, err |
||||
|
} |
||||
|
|
||||
|
// Iterate implements the method Iterate of the interface db.Storage
|
||||
|
func (p *PebbleStorage) Iterate(f func([]byte, []byte) (bool, error)) error { |
||||
|
// NewIter already provides a point-in-time view of the current DB
|
||||
|
// state, but if is used for long term (is not the case), should use an
|
||||
|
// iterator over an snapshot:
|
||||
|
// snapshot := p.pdb.NewSnapshot()
|
||||
|
// defer snapshot.Close()
|
||||
|
// iter := snapshot.NewIter(nil)
|
||||
|
iter := p.pdb.NewIter(nil) |
||||
|
defer iter.Close() |
||||
|
|
||||
|
iter.First() // move the iterator to the first key/value pair
|
||||
|
if len(iter.Key()) < len(p.prefix) || !bytes.Equal(iter.Key()[:len(p.prefix)], p.prefix) { |
||||
|
} else { |
||||
|
localKey := iter.Key()[len(p.prefix):] |
||||
|
if _, err := f(localKey, iter.Value()); err != nil { |
||||
|
return err |
||||
|
} |
||||
|
} |
||||
|
for iter.Next() { |
||||
|
if len(iter.Key()) < len(p.prefix) || !bytes.Equal(iter.Key()[:len(p.prefix)], p.prefix) { |
||||
|
continue |
||||
|
} |
||||
|
localKey := iter.Key()[len(p.prefix):] |
||||
|
if cont, err := f(localKey, iter.Value()); err != nil { |
||||
|
return err |
||||
|
} else if !cont { |
||||
|
break |
||||
|
} |
||||
|
} |
||||
|
return iter.Error() |
||||
|
} |
||||
|
|
||||
|
// Get retreives a value from a key in the interface db.Tx
|
||||
|
func (tx *PebbleStorageTx) Get(key []byte) ([]byte, error) { |
||||
|
var err error |
||||
|
|
||||
|
fullkey := db.Concat(tx.prefix, key) |
||||
|
|
||||
|
if value, ok := tx.cache.Get(fullkey); ok { |
||||
|
return value, nil |
||||
|
} |
||||
|
|
||||
|
value, closer, err := tx.pdb.Get(fullkey) |
||||
|
if err == pebble.ErrNotFound { |
||||
|
return nil, db.ErrNotFound |
||||
|
} |
||||
|
closer.Close() |
||||
|
|
||||
|
return value, err |
||||
|
} |
||||
|
|
||||
|
// Put saves a key:value into the db.Storage
|
||||
|
func (tx *PebbleStorageTx) Put(k, v []byte) { |
||||
|
tx.cache.Put(db.Concat(tx.prefix, k[:]), v) |
||||
|
} |
||||
|
|
||||
|
// Add implements the method Add of the interface db.Tx
|
||||
|
func (tx *PebbleStorageTx) Add(atx db.Tx) { |
||||
|
ldbtx := atx.(*PebbleStorageTx) |
||||
|
for _, v := range ldbtx.cache { |
||||
|
tx.cache.Put(v.K, v.V) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// Commit implements the method Commit of the interface db.Tx
|
||||
|
func (tx *PebbleStorageTx) Commit() error { |
||||
|
batch := tx.PebbleStorage.pdb.NewBatch() |
||||
|
for _, v := range tx.cache { |
||||
|
batch.Set(v.K, v.V, nil) |
||||
|
} |
||||
|
|
||||
|
tx.cache = nil |
||||
|
return batch.Commit(nil) |
||||
|
} |
||||
|
|
||||
|
// Close implements the method Close of the interface db.Tx
|
||||
|
func (tx *PebbleStorageTx) Close() { |
||||
|
tx.cache = nil |
||||
|
} |
||||
|
|
||||
|
// Close implements the method Close of the interface db.Storage
|
||||
|
func (p *PebbleStorage) Close() { |
||||
|
if err := p.pdb.Close(); err != nil { |
||||
|
panic(err) |
||||
|
} |
||||
|
log.Info("Database closed") |
||||
|
} |
||||
|
|
||||
|
// Pebble is an extra method that returns the *pebble.DB
|
||||
|
func (p *PebbleStorage) Pebble() *pebble.DB { |
||||
|
return p.pdb |
||||
|
} |
||||
|
|
||||
|
// List implements the method List of the interface db.Storage
|
||||
|
func (p *PebbleStorage) List(limit int) ([]db.KV, error) { |
||||
|
ret := []db.KV{} |
||||
|
err := p.Iterate(func(key []byte, value []byte) (bool, error) { |
||||
|
ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)}) |
||||
|
if len(ret) == limit { |
||||
|
return false, nil |
||||
|
} |
||||
|
return true, nil |
||||
|
}) |
||||
|
return ret, err |
||||
|
} |
||||
@ -0,0 +1,57 @@ |
|||||
|
package pebble |
||||
|
|
||||
|
import ( |
||||
|
"io/ioutil" |
||||
|
"os" |
||||
|
"testing" |
||||
|
|
||||
|
"github.com/iden3/go-merkletree/db" |
||||
|
"github.com/iden3/go-merkletree/db/test" |
||||
|
"github.com/stretchr/testify/require" |
||||
|
) |
||||
|
|
||||
|
var rmDirs []string |
||||
|
|
||||
|
func pebbleStorage(t *testing.T) db.Storage { |
||||
|
dir, err := ioutil.TempDir("", "db") |
||||
|
rmDirs = append(rmDirs, dir) |
||||
|
if err != nil { |
||||
|
t.Fatal(err) |
||||
|
return nil |
||||
|
} |
||||
|
sto, err := NewPebbleStorage(dir, false) |
||||
|
if err != nil { |
||||
|
t.Fatal(err) |
||||
|
return nil |
||||
|
} |
||||
|
return sto |
||||
|
} |
||||
|
|
||||
|
func TestPebble(t *testing.T) { |
||||
|
test.TestReturnKnownErrIfNotExists(t, pebbleStorage(t)) |
||||
|
test.TestStorageInsertGet(t, pebbleStorage(t)) |
||||
|
test.TestStorageWithPrefix(t, pebbleStorage(t)) |
||||
|
test.TestConcatTx(t, pebbleStorage(t)) |
||||
|
test.TestIterate(t, pebbleStorage(t)) |
||||
|
test.TestList(t, pebbleStorage(t)) |
||||
|
} |
||||
|
|
||||
|
func TestPebbleInterface(t *testing.T) { |
||||
|
var db db.Storage //nolint:gosimple
|
||||
|
|
||||
|
dir, err := ioutil.TempDir("", "db") |
||||
|
require.Nil(t, err) |
||||
|
rmDirs = append(rmDirs, dir) |
||||
|
sto, err := NewPebbleStorage(dir, false) |
||||
|
require.Nil(t, err) |
||||
|
db = sto |
||||
|
require.NotNil(t, db) |
||||
|
} |
||||
|
|
||||
|
func TestMain(m *testing.M) { |
||||
|
result := m.Run() |
||||
|
for _, dir := range rmDirs { |
||||
|
os.RemoveAll(dir) |
||||
|
} |
||||
|
os.Exit(result) |
||||
|
} |
||||