arnaucube
arnaucube/arboPublic Self-hosted repository
arbo/db/goleveldb/goleveldb.go
db/goleveldb/goleveldb.goRaw
  1// Package goleveldb implements arbo database interfaces using LevelDB.
  2package goleveldb
  3
  4import (
  5	"bytes"
  6	"errors"
  7	"fmt"
  8	"sync"
  9
 10	"github.com/arnaucube/arbo/db"
 11	"github.com/syndtr/goleveldb/leveldb"
 12	"github.com/syndtr/goleveldb/leveldb/opt"
 13	"github.com/syndtr/goleveldb/leveldb/util"
 14)
 15
 16// LevelDB implements db.Database using LevelDB.
 17type LevelDB struct {
 18	db *leveldb.DB
 19}
 20
 21// Ensure that LevelDB implements the db.Database interface
 22var _ db.Database = (*LevelDB)(nil)
 23
 24// New returns a LevelDB which implements the db.Database interface
 25func New(opts db.Options) (*LevelDB, error) {
 26	// Open the LevelDB database
 27	db, err := leveldb.OpenFile(opts.Path, &opt.Options{})
 28	if err != nil {
 29		return nil, fmt.Errorf("could not open leveldb: %w", err)
 30	}
 31	return &LevelDB{
 32		db: db,
 33	}, nil
 34}
 35
 36// Close closes the database.
 37func (d *LevelDB) Close() error {
 38	return d.db.Close()
 39}
 40
 41// WriteTx creates a buffered write transaction.
 42func (d *LevelDB) WriteTx() db.WriteTx {
 43	return &WriteTx{
 44		db:    d.db,
 45		batch: new(leveldb.Batch),
 46	}
 47}
 48
 49// Get retrieves the value associated with key.
 50func (d *LevelDB) Get(key []byte) ([]byte, error) {
 51	val, err := d.db.Get(key, nil)
 52	if errors.Is(err, leveldb.ErrNotFound) {
 53		return nil, db.ErrKeyNotFound
 54	}
 55	if err != nil {
 56		return nil, err
 57	}
 58	return val, nil
 59}
 60
 61// Iterate visits keys with the given prefix in lexicographic order.
 62func (d *LevelDB) Iterate(prefix []byte, callback func(key, value []byte) bool) error {
 63	iter := d.db.NewIterator(util.BytesPrefix(prefix), nil)
 64	defer iter.Release()
 65	for iter.Next() {
 66		if !callback(iter.Key(), iter.Value()) {
 67			break
 68		}
 69	}
 70	return iter.Error()
 71}
 72
 73// Set stores a key-value pair.
 74func (d *LevelDB) Set(key, value []byte) error {
 75	return d.db.Put(key, value, nil)
 76}
 77
 78// Delete removes a key.
 79func (d *LevelDB) Delete(key []byte) error {
 80	return d.db.Delete(key, nil)
 81}
 82
 83// Commit writes a batch to the database.
 84func (d *LevelDB) Commit(batch *leveldb.Batch) error {
 85	return d.db.Write(batch, nil)
 86}
 87
 88// Compact implements the db.Database.Compact interface method.
 89func (d *LevelDB) Compact() error {
 90	return d.db.CompactRange(util.Range{})
 91}
 92
 93// WriteTx implements the interface db.WriteTx for goleveldb
 94type WriteTx struct {
 95	batch      *leveldb.Batch
 96	db         *leveldb.DB
 97	inMemBatch sync.Map
 98}
 99
100// check that WriteTx implements the db.WriteTx interface
101var _ db.WriteTx = (*WriteTx)(nil)
102
103// Get retrieves a value, including pending writes.
104func (tx *WriteTx) Get(k []byte) ([]byte, error) {
105	val, ok := tx.inMemBatch.Load(string(k))
106	if !ok {
107		val, err := tx.db.Get(k, nil)
108		if errors.Is(err, leveldb.ErrNotFound) {
109			return nil, db.ErrKeyNotFound
110		}
111		return val, err
112	}
113	return val.([]byte), nil
114}
115
116// Iterate visits pending writes and stored keys with the given prefix.
117func (tx *WriteTx) Iterate(prefix []byte, callback func(k, v []byte) bool) error {
118	inMemory := make(map[string]bool)
119	tx.inMemBatch.Range(func(k, v any) bool {
120		keyBytes := []byte(k.(string))
121		if bytes.HasPrefix(keyBytes, prefix) {
122			inMemory[string(keyBytes)] = true
123			return callback(keyBytes, v.([]byte))
124		}
125		return true
126	})
127	iter := tx.db.NewIterator(util.BytesPrefix(prefix), nil)
128	defer iter.Release()
129	for iter.Next() {
130		if inMemory[string(iter.Key())] {
131			continue
132		}
133		if !callback(iter.Key(), iter.Value()) {
134			break
135		}
136	}
137	return iter.Error()
138}
139
140// Set buffers a key-value pair.
141func (tx *WriteTx) Set(k, v []byte) error {
142	tx.batch.Put(k, v)
143	tx.inMemBatch.Store(string(k), v)
144	return nil
145}
146
147// Delete buffers the removal of a key.
148func (tx *WriteTx) Delete(k []byte) error {
149	tx.batch.Delete(k)
150	tx.inMemBatch.Delete(string(k))
151	return nil
152}
153
154// Apply copies the entries from another transaction.
155func (tx *WriteTx) Apply(otherTx db.WriteTx) error {
156	return otherTx.Iterate(nil, func(k, v []byte) bool {
157		tx.inMemBatch.Store(string(k), v)
158		tx.batch.Put(k, v)
159		return true
160	})
161}
162
163// Commit writes the buffered batch to the database.
164func (tx *WriteTx) Commit() error {
165	return tx.db.Write(tx.batch, nil)
166}
167
168// Discard clears the buffered writes.
169func (tx *WriteTx) Discard() {
170	tx.batch.Reset()
171	tx.inMemBatch = sync.Map{}
172}