1package pebble
  2
  3import (
  4	"github.com/cockroachdb/pebble"
  5	"github.com/iden3/go-merkletree/db"
  6)
  7
  8// Storage implements the db.Storage interface
  9type Storage struct {
 10	pdb    *pebble.DB
 11	prefix []byte
 12}
 13
 14// StorageTx implements the db.Tx interface
 15type StorageTx struct {
 16	*Storage
 17	batch *pebble.Batch
 18}
 19
 20// NewPebbleStorage returns a new Storage
 21func NewPebbleStorage(path string, errorIfMissing bool) (*Storage, error) {
 22	o := &pebble.Options{
 23		ErrorIfNotExists: errorIfMissing,
 24	}
 25	rdb, err := pebble.Open(path, o)
 26	if err != nil {
 27		return nil, err
 28	}
 29	return &Storage{rdb, []byte{}}, nil
 30}
 31
 32// WithPrefix implements the method WithPrefix of the interface db.Storage
 33func (p *Storage) WithPrefix(prefix []byte) db.Storage {
 34	return &Storage{p.pdb, db.Concat(p.prefix, prefix)}
 35}
 36
 37// NewTx implements the method NewTx of the interface db.Storage
 38func (p *Storage) NewTx() (db.Tx, error) {
 39	return &StorageTx{p, p.pdb.NewIndexedBatch()}, nil
 40}
 41
 42// Get retreives a value from a key in the db.Storage
 43func (p *Storage) Get(key []byte) ([]byte, error) {
 44	v, closer, err := p.pdb.Get(db.Concat(p.prefix, key[:]))
 45	if err == pebble.ErrNotFound {
 46		return nil, db.ErrNotFound
 47	}
 48	if err != nil {
 49		return nil, err
 50	}
 51	err = closer.Close()
 52	return v, err
 53}
 54
 55//nolint:lll
 56// https://github.com/cockroachdb/pebble/pull/923/files#diff-c2ade2f386c41794d5ebc57ee49b57a5fca8082e03255e5bff13977cbc061287R39
 57func keyUpperBound(b []byte) []byte {
 58	end := make([]byte, len(b))
 59	copy(end, b)
 60	for i := len(end) - 1; i >= 0; i-- {
 61		end[i] = end[i] + 1
 62		if end[i] != 0 {
 63			return end[:i+1]
 64		}
 65	}
 66	return nil // no upper-bound
 67}
 68func prefixIterOptions(prefix []byte) *pebble.IterOptions {
 69	return &pebble.IterOptions{
 70		LowerBound: prefix,
 71		UpperBound: keyUpperBound(prefix),
 72	}
 73}
 74
 75// Iterate implements the method Iterate of the interface db.Storage
 76func (p *Storage) Iterate(f func([]byte, []byte) (bool, error)) (err error) {
 77	// NewIter already provides a point-in-time view of the current DB
 78	// state, but if is used for long term (is not the case), should use an
 79	// iterator over an snapshot:
 80	// snapshot := p.pdb.NewSnapshot()
 81	// defer snapshot.Close()
 82	// iter := snapshot.NewIter(nil)
 83	iter := p.pdb.NewIter(prefixIterOptions(p.prefix))
 84	defer func() {
 85		err1 := iter.Close()
 86		if err != nil {
 87			return
 88		}
 89		err = err1
 90	}()
 91
 92	for iter.First(); iter.Valid(); iter.Next() {
 93		localKey := iter.Key()[len(p.prefix):]
 94		if cont, err := f(localKey, iter.Value()); err != nil {
 95			return err
 96		} else if !cont {
 97			break
 98		}
 99	}
100	return iter.Error()
101}
102
103// Get retreives a value from a key in the interface db.Tx
104func (tx *StorageTx) Get(key []byte) ([]byte, error) {
105	var err error
106
107	fullkey := db.Concat(tx.prefix, key)
108
109	v, closer, err := tx.batch.Get(fullkey)
110	if err == pebble.ErrNotFound {
111		return nil, db.ErrNotFound
112	}
113	if err != nil {
114		return nil, err
115	}
116	err = closer.Close()
117	return v, err
118}
119
120// Put saves a key:value into the db.Storage
121func (tx *StorageTx) Put(k, v []byte) error {
122	return tx.batch.Set(db.Concat(tx.prefix, k[:]), v, nil)
123}
124
125// Add implements the method Add of the interface db.Tx
126func (tx *StorageTx) Add(atx db.Tx) error {
127	patx := atx.(*StorageTx)
128	return tx.batch.Apply(patx.batch, nil)
129}
130
131// Commit implements the method Commit of the interface db.Tx
132func (tx *StorageTx) Commit() error {
133	return tx.batch.Commit(nil)
134}
135
136// Close implements the method Close of the interface db.Tx
137func (tx *StorageTx) Close() {
138	_ = tx.batch.Close()
139}
140
141// Close implements the method Close of the interface db.Storage
142func (p *Storage) Close() {
143	if err := p.pdb.Close(); err != nil {
144		panic(err)
145	}
146}
147
148// Pebble is an extra method that returns the *pebble.DB
149func (p *Storage) Pebble() *pebble.DB {
150	return p.pdb
151}
152
153// List implements the method List of the interface db.Storage
154func (p *Storage) List(limit int) ([]db.KV, error) {
155	ret := []db.KV{}
156	err := p.Iterate(func(key []byte, value []byte) (bool, error) {
157		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
158		if len(ret) == limit {
159			return false, nil
160		}
161		return true, nil
162	})
163	return ret, err
164}