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