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// PebbleStorage implements the db.Storage interface
 12type PebbleStorage struct {
 13	pdb    *pebble.DB
 14	prefix []byte
 15}
 16
 17// PebbleStorageTx implements the db.Tx interface
 18type PebbleStorageTx struct {
 19	*PebbleStorage
 20	batch *pebble.Batch
 21}
 22
 23// NewPebbleStorage returns a new PebbleStorage
 24func NewPebbleStorage(path string, errorIfMissing bool) (*PebbleStorage, 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 &PebbleStorage{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 *PebbleStorage) 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 *PebbleStorage) WithPrefix(prefix []byte) db.Storage {
 67	return &PebbleStorage{p.pdb, db.Concat(p.prefix, prefix)}
 68}
 69
 70// NewTx implements the method NewTx of the interface db.Storage
 71func (p *PebbleStorage) NewTx() (db.Tx, error) {
 72	return &PebbleStorageTx{p, p.pdb.NewIndexedBatch()}, nil
 73}
 74
 75// Get retreives a value from a key in the db.Storage
 76func (p *PebbleStorage) 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	closer.Close()
 82	return v, err
 83}
 84
 85// https://github.com/cockroachdb/pebble/pull/923/files#diff-c2ade2f386c41794d5ebc57ee49b57a5fca8082e03255e5bff13977cbc061287R39
 86func keyUpperBound(b []byte) []byte {
 87	end := make([]byte, len(b))
 88	copy(end, b)
 89	for i := len(end) - 1; i >= 0; i-- {
 90		end[i] = end[i] + 1
 91		if end[i] != 0 {
 92			return end[:i+1]
 93		}
 94	}
 95	return nil // no upper-bound
 96}
 97func prefixIterOptions(prefix []byte) *pebble.IterOptions {
 98	return &pebble.IterOptions{
 99		LowerBound: prefix,
100		UpperBound: keyUpperBound(prefix),
101	}
102}
103
104// Iterate implements the method Iterate of the interface db.Storage
105func (p *PebbleStorage) Iterate(f func([]byte, []byte) (bool, error)) error {
106	// NewIter already provides a point-in-time view of the current DB
107	// state, but if is used for long term (is not the case), should use an
108	// iterator over an snapshot:
109	// snapshot := p.pdb.NewSnapshot()
110	// defer snapshot.Close()
111	// iter := snapshot.NewIter(nil)
112	iter := p.pdb.NewIter(prefixIterOptions(p.prefix))
113	defer iter.Close()
114
115	for iter.First(); iter.Valid(); iter.Next() {
116		localKey := iter.Key()[len(p.prefix):]
117		if cont, err := f(localKey, iter.Value()); err != nil {
118			return err
119		} else if !cont {
120			break
121		}
122	}
123	return iter.Error()
124}
125
126// Get retreives a value from a key in the interface db.Tx
127func (tx *PebbleStorageTx) Get(key []byte) ([]byte, error) {
128	var err error
129
130	fullkey := db.Concat(tx.prefix, key)
131
132	v, closer, err := tx.batch.Get(fullkey)
133	if err == pebble.ErrNotFound {
134		return nil, db.ErrNotFound
135	}
136	closer.Close()
137
138	return v, err
139}
140
141// Put saves a key:value into the db.Storage
142func (tx *PebbleStorageTx) Put(k, v []byte) error {
143	return tx.batch.Set(db.Concat(tx.prefix, k[:]), v, nil)
144}
145
146// Add implements the method Add of the interface db.Tx
147func (tx *PebbleStorageTx) Add(atx db.Tx) error {
148	patx := atx.(*PebbleStorageTx)
149	return tx.batch.Apply(patx.batch, nil)
150}
151
152// Commit implements the method Commit of the interface db.Tx
153func (tx *PebbleStorageTx) Commit() error {
154	return tx.batch.Commit(nil)
155}
156
157// Close implements the method Close of the interface db.Tx
158func (tx *PebbleStorageTx) Close() {
159	_ = tx.batch.Close()
160}
161
162// Close implements the method Close of the interface db.Storage
163func (p *PebbleStorage) Close() {
164	if err := p.pdb.Close(); err != nil {
165		panic(err)
166	}
167	log.Info("Database closed")
168}
169
170// Pebble is an extra method that returns the *pebble.DB
171func (p *PebbleStorage) Pebble() *pebble.DB {
172	return p.pdb
173}
174
175// List implements the method List of the interface db.Storage
176func (p *PebbleStorage) List(limit int) ([]db.KV, error) {
177	ret := []db.KV{}
178	err := p.Iterate(func(key []byte, value []byte) (bool, error) {
179		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
180		if len(ret) == limit {
181			return false, nil
182		}
183		return true, nil
184	})
185	return ret, err
186}