1// Package pebbledb implements arbo database interfaces using Pebble.
  2package pebbledb
  3
  4import (
  5	"bytes"
  6	"errors"
  7	"fmt"
  8	"os"
  9
 10	"github.com/arnaucube/arbo/db"
 11	"github.com/cockroachdb/pebble"
 12)
 13
 14// WriteTx implements the interface db.WriteTx
 15type WriteTx struct {
 16	batch *pebble.Batch
 17}
 18
 19// check that WriteTx implements the db.WriteTx interface
 20var _ db.WriteTx = (*WriteTx)(nil)
 21
 22func get(reader pebble.Reader, k []byte) ([]byte, error) {
 23	v, closer, err := reader.Get(k)
 24	if errors.Is(err, pebble.ErrNotFound) {
 25		return nil, db.ErrKeyNotFound
 26	}
 27	if err != nil {
 28		return nil, err
 29	}
 30
 31	// Note that the returned value slice is only valid until Close is called.
 32	// Make a copy so we can return it.
 33	// TODO(mvdan): write a dbtest test to ensure this property on all DBs.
 34	v2 := bytes.Clone(v)
 35
 36	if err := closer.Close(); err != nil {
 37		return nil, err
 38	}
 39	return v2, nil
 40}
 41
 42func iterate(reader pebble.Reader, prefix []byte, callback func(k, v []byte) bool) (err error) {
 43	iterOptions := &pebble.IterOptions{
 44		LowerBound: prefix,
 45		UpperBound: keyUpperBound(prefix),
 46	}
 47	iter, err := reader.NewIter(iterOptions)
 48	if err != nil {
 49		return err
 50	}
 51	defer func() {
 52		errC := iter.Close()
 53		if err != nil {
 54			return
 55		}
 56		err = errC
 57	}()
 58
 59	for iter.First(); iter.Valid(); iter.Next() {
 60		localKey := iter.Key()[len(prefix):]
 61		if cont := callback(localKey, iter.Value()); !cont {
 62			break
 63		}
 64	}
 65	return iter.Error()
 66}
 67
 68// Get implements the db.WriteTx.Get interface method
 69func (tx *WriteTx) Get(k []byte) ([]byte, error) {
 70	return get(tx.batch, k)
 71}
 72
 73// Iterate implements the db.WriteTx.Iterate interface method.
 74func (tx *WriteTx) Iterate(prefix []byte, callback func(k, v []byte) bool) (err error) {
 75	return iterate(tx.batch, prefix, callback)
 76}
 77
 78// Set implements the db.WriteTx.Set interface method
 79func (tx *WriteTx) Set(k, v []byte) error {
 80	return tx.batch.Set(k, v, nil)
 81}
 82
 83// Delete implements the db.WriteTx.Delete interface method
 84func (tx *WriteTx) Delete(k []byte) error {
 85	return tx.batch.Delete(k, nil)
 86}
 87
 88// Apply implements the db.WriteTx.Apply interface method
 89func (tx *WriteTx) Apply(other db.WriteTx) (err error) {
 90	otherPebble := db.UnwrapWriteTx(other).(*WriteTx)
 91	return tx.batch.Apply(otherPebble.batch, nil)
 92}
 93
 94// Commit implements the db.WriteTx.Commit interface method
 95func (tx *WriteTx) Commit() error {
 96	if tx.batch == nil {
 97		return fmt.Errorf("cannot commit pebble tx: already committed or discarded")
 98	}
 99	err := tx.batch.Commit(nil)
100	tx.batch = nil
101	return err
102}
103
104// Discard implements the db.WriteTx.Discard interface method
105func (tx *WriteTx) Discard() {
106	if tx.batch == nil {
107		// Silently allow discarding twice or after a commit,
108		// since it can help for the sake of defers.
109		// Note that upstream pebbledb doesn't really allow calling Close twice;
110		// the object gets put in a shared pool, so a later call to Close races
111		// with another goroutine having already reused the same object from the pool.
112		return
113	}
114	// Discard cannot return errors; closing the batch only releases its resources.
115	_ = tx.batch.Close()
116	tx.batch = nil
117}
118
119// PebbleDB implements db.Database interface
120type PebbleDB struct {
121	db *pebble.DB
122}
123
124// check that PebbleDB implements the db.Database interface
125var _ db.Database = (*PebbleDB)(nil)
126
127// New returns a PebbleDB using the given Options, which implements the
128// db.Database interface
129func New(opts db.Options) (*PebbleDB, error) {
130	if err := os.MkdirAll(opts.Path, os.ModePerm); err != nil {
131		return nil, err
132	}
133	o := &pebble.Options{
134		Levels: []pebble.LevelOptions{
135			{
136				Compression: pebble.SnappyCompression,
137			},
138		},
139	}
140	db, err := pebble.Open(opts.Path, o)
141	if err != nil {
142		return nil, err
143	}
144
145	return &PebbleDB{
146		db: db,
147	}, nil
148}
149
150// Get implements the db.WriteTx.Get interface method
151func (db *PebbleDB) Get(k []byte) ([]byte, error) {
152	return get(db.db, k)
153}
154
155// WriteTx returns a db.WriteTx
156func (db *PebbleDB) WriteTx() db.WriteTx {
157	return &WriteTx{
158		batch: db.db.NewIndexedBatch(),
159	}
160}
161
162// Close closes the PebbleDB
163func (db *PebbleDB) Close() error {
164	return db.db.Close()
165}
166
167func keyUpperBound(b []byte) []byte {
168	// See keyUpperBound in Pebble's iterator_example_test.go.
169	end := bytes.Clone(b)
170	for i := len(end) - 1; i >= 0; i-- {
171		end[i] = end[i] + 1
172		if end[i] != 0 {
173			return end[:i+1]
174		}
175	}
176	return nil // no upper-bound
177}
178
179// Iterate implements the db.Database.Iterate interface method
180func (db *PebbleDB) Iterate(prefix []byte, callback func(k, v []byte) bool) (err error) {
181	return iterate(db.db, prefix, callback)
182}
183
184// Compact implements the db.Database.Compact interface method
185func (db *PebbleDB) Compact() error {
186	// from https://github.com/cockroachdb/pebble/issues/1474#issuecomment-1022313365
187	iter, err := db.db.NewIter(nil)
188	if err != nil {
189		return err
190	}
191	var first, last []byte
192	if iter.First() {
193		first = append(first, iter.Key()...)
194	}
195	if iter.Last() {
196		last = append(last, iter.Key()...)
197	}
198	if err := iter.Close(); err != nil {
199		return err
200	}
201	return db.db.Compact(first, last, true)
202}