1package pebble
  2
  3import (
  4	"bytes"
  5	"encoding/json"
  6
  7	"github.com/cockroachdb/pebble"
  8	"github.com/iden3/go-merkletree/db"
  9	log "github.com/sirupsen/logrus"
 10)
 11
 12// PebbleStorage implements the db.Storage interface
 13type PebbleStorage struct {
 14	pdb    *pebble.DB
 15	prefix []byte
 16}
 17
 18// PebbleStorageTx implements the db.Tx interface
 19type PebbleStorageTx struct {
 20	// FUTURE currently Tx is using the same strategy than in MemoryDB and
 21	// LevelDB, in next iteration can be moved to Pebble Batch strategy
 22	*PebbleStorage
 23	cache db.KvMap
 24}
 25
 26// NewPebbleStorage returns a new PebbleStorage
 27func NewPebbleStorage(path string, errorIfMissing bool) (*PebbleStorage, error) {
 28	o := &pebble.Options{
 29		ErrorIfNotExists: errorIfMissing,
 30	}
 31	rdb, err := pebble.Open(path, o)
 32	if err != nil {
 33		return nil, err
 34	}
 35	return &PebbleStorage{rdb, []byte{}}, nil
 36}
 37
 38type storageInfo struct {
 39	KeyCount   int
 40	ClaimCount int
 41}
 42
 43// Info implements the method Info of the interface db.Storage
 44func (p *PebbleStorage) Info() string {
 45	keycount := 0
 46	claimcount := 0
 47	err := p.Iterate(func(key []byte, value []byte) (bool, error) {
 48		if value[0] == byte(1) {
 49			claimcount++
 50		}
 51
 52		keycount++
 53		return true, nil
 54	})
 55	if err != nil {
 56		return err.Error()
 57	}
 58	json, _ := json.MarshalIndent(
 59		storageInfo{
 60			KeyCount:   keycount,
 61			ClaimCount: claimcount,
 62		},
 63		"", "  ",
 64	)
 65	return string(json)
 66}
 67
 68// WithPrefix implements the method WithPrefix of the interface db.Storage
 69func (p *PebbleStorage) WithPrefix(prefix []byte) db.Storage {
 70	return &PebbleStorage{p.pdb, db.Concat(p.prefix, prefix)}
 71}
 72
 73// NewTx implements the method NewTx of the interface db.Storage
 74func (p *PebbleStorage) NewTx() (db.Tx, error) {
 75	return &PebbleStorageTx{p, make(db.KvMap)}, nil
 76}
 77
 78// Get retreives a value from a key in the db.Storage
 79func (p *PebbleStorage) Get(key []byte) ([]byte, error) {
 80	v, closer, err := p.pdb.Get(db.Concat(p.prefix, key[:]))
 81	if err == pebble.ErrNotFound {
 82		return nil, db.ErrNotFound
 83	}
 84	closer.Close()
 85	return v, err
 86}
 87
 88// Iterate implements the method Iterate of the interface db.Storage
 89func (p *PebbleStorage) Iterate(f func([]byte, []byte) (bool, error)) error {
 90	// NewIter already provides a point-in-time view of the current DB
 91	// state, but if is used for long term (is not the case), should use an
 92	// iterator over an snapshot:
 93	// snapshot := p.pdb.NewSnapshot()
 94	// defer snapshot.Close()
 95	// iter := snapshot.NewIter(nil)
 96	iter := p.pdb.NewIter(nil)
 97	defer iter.Close()
 98
 99	iter.First() // move the iterator to the first key/value pair
100	if len(iter.Key()) < len(p.prefix) || !bytes.Equal(iter.Key()[:len(p.prefix)], p.prefix) {
101	} else {
102		localKey := iter.Key()[len(p.prefix):]
103		if _, err := f(localKey, iter.Value()); err != nil {
104			return err
105		}
106	}
107	for iter.Next() {
108		if len(iter.Key()) < len(p.prefix) || !bytes.Equal(iter.Key()[:len(p.prefix)], p.prefix) {
109			continue
110		}
111		localKey := iter.Key()[len(p.prefix):]
112		if cont, err := f(localKey, iter.Value()); err != nil {
113			return err
114		} else if !cont {
115			break
116		}
117	}
118	return iter.Error()
119}
120
121// Get retreives a value from a key in the interface db.Tx
122func (tx *PebbleStorageTx) Get(key []byte) ([]byte, error) {
123	var err error
124
125	fullkey := db.Concat(tx.prefix, key)
126
127	if value, ok := tx.cache.Get(fullkey); ok {
128		return value, nil
129	}
130
131	value, closer, err := tx.pdb.Get(fullkey)
132	if err == pebble.ErrNotFound {
133		return nil, db.ErrNotFound
134	}
135	closer.Close()
136
137	return value, err
138}
139
140// Put saves a key:value into the db.Storage
141func (tx *PebbleStorageTx) Put(k, v []byte) {
142	tx.cache.Put(db.Concat(tx.prefix, k[:]), v)
143}
144
145// Add implements the method Add of the interface db.Tx
146func (tx *PebbleStorageTx) Add(atx db.Tx) {
147	ldbtx := atx.(*PebbleStorageTx)
148	for _, v := range ldbtx.cache {
149		tx.cache.Put(v.K, v.V)
150	}
151}
152
153// Commit implements the method Commit of the interface db.Tx
154func (tx *PebbleStorageTx) Commit() error {
155	batch := tx.PebbleStorage.pdb.NewBatch()
156	for _, v := range tx.cache {
157		batch.Set(v.K, v.V, nil)
158	}
159
160	tx.cache = nil
161	return batch.Commit(nil)
162}
163
164// Close implements the method Close of the interface db.Tx
165func (tx *PebbleStorageTx) Close() {
166	tx.cache = nil
167}
168
169// Close implements the method Close of the interface db.Storage
170func (p *PebbleStorage) Close() {
171	if err := p.pdb.Close(); err != nil {
172		panic(err)
173	}
174	log.Info("Database closed")
175}
176
177// Pebble is an extra method that returns the *pebble.DB
178func (p *PebbleStorage) Pebble() *pebble.DB {
179	return p.pdb
180}
181
182// List implements the method List of the interface db.Storage
183func (p *PebbleStorage) List(limit int) ([]db.KV, error) {
184	ret := []db.KV{}
185	err := p.Iterate(func(key []byte, value []byte) (bool, error) {
186		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
187		if len(ret) == limit {
188			return false, nil
189		}
190		return true, nil
191	})
192	return ret, err
193}