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