sample_chain/database/leveldbprovider.go

package database

import (
	log "github.com/corgi-kx/logcustom"
	"github.com/pkg/errors"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	goleveldbutil "github.com/syndtr/goleveldb/leveldb/util"
)

const maxBatchSize = 1000000

// DBProvider is a core to a named db
type DBProvider struct {
	dbName string
	db     *DB
}

// Iterator extends actual leveldb iterator
type Iterator struct {
	dbName string
	iterator.Iterator
}

// NewDBProvider constructs a `DBProvider`
//
//dp:dbPath dn:dbname
func NewDBProvider(dp, dn string) *DBProvider {
	db := CreateDB(dp)
	db.Open()

	return &DBProvider{
		dbName: dn,
		db:     db,
	}
}

// Get returns the value for the given key
func (p *DBProvider) Get(key []byte) ([]byte, error) {
	return p.db.Get(constructLevelKey(p.dbName, key))
}

// Put saves the key/value
func (p *DBProvider) Put(key []byte, value []byte, sync bool) error {
	return p.db.Put(constructLevelKey(p.dbName, key), value, sync)
}

// Delete deletes the given key
func (p *DBProvider) Delete(key []byte, sync bool) error {
	return p.db.Delete(constructLevelKey(p.dbName, key), sync)
}

// DeleteAll deletes all the keys that belong to the channel (dbName).
func (p *DBProvider) DeleteAll() error {
	iter, err := p.GetIterator(nil, nil)
	if err != nil {
		return err
	}
	defer iter.Release()

	// use leveldb iterator directly to be more efficient
	dbIter := iter.Iterator

	// This is common code shared by all the leveldb instances. Because each leveldb has its own key size pattern,
	// each batch is limited by memory usage instead of number of keys. Once the batch memory usage reaches maxBatchSize,
	// the batch will be committed.
	numKeys := 0
	batchSize := 0
	batch := &leveldb.Batch{}
	for dbIter.Next() {
		if err := dbIter.Error(); err != nil {
			return errors.Wrap(err, "internal leveldb error while retrieving data from db iterator")
		}
		key := dbIter.Key()
		numKeys++
		batchSize = batchSize + len(key)
		batch.Delete(key)
		if batchSize >= maxBatchSize {
			if err := p.db.WriteBatch(batch, true); err != nil {
				return err
			}
			log.Infof("Have removed %d entries for channel %s in leveldb %s", numKeys, p.dbName, p.db.dbPath)
			batchSize = 0
			batch.Reset()
		}
	}
	if batch.Len() > 0 {
		return p.db.WriteBatch(batch, true)
	}
	return nil
}

// IsEmpty returns true if no data exists for the DBProvider
func (p *DBProvider) IsEmpty() (bool, error) {
	itr, err := p.GetIterator(nil, nil)
	if err != nil {
		return false, err
	}
	defer itr.Release()

	if err := itr.Error(); err != nil {
		return false, errors.WithMessagef(itr.Error(), "internal leveldb error while obtaining next entry from iterator")
	}

	return !itr.Next(), nil
}

// NewUpdateBatch returns a new UpdateBatch that can be used to update the db
func (p *DBProvider) NewUpdateBatch() *UpdateBatch {
	return &UpdateBatch{
		dbName:       p.dbName,
		leveldbBatch: &leveldb.Batch{},
	}
}

// WriteBatch writes a batch in an atomic way
func (p *DBProvider) WriteBatch(batch *UpdateBatch, sync bool) error {
	if batch == nil || batch.leveldbBatch.Len() == 0 {
		return nil
	}
	if err := p.db.WriteBatch(batch.leveldbBatch, sync); err != nil {
		return err
	}
	return nil
}

// GetIterator gets a core to iterator. The iterator should be released after the use.
// The resultset contains all the keys that are present in the db between the startKey (inclusive) and the endKey (exclusive).
// A nil startKey represents the first available key and a nil endKey represent a logical key after the last available key
func (p *DBProvider) GetIterator(startKey []byte, endKey []byte) (*Iterator, error) {
	var itr iterator.Iterator
	if startKey == nil && endKey == nil {
		itr = p.db.GetIterator(startKey, startKey)
	} else if endKey == nil {
		sKey := constructLevelKey(p.dbName, startKey)
		itr = p.db.GetIterator(sKey, startKey)
	} else {
		sKey := constructLevelKey(p.dbName, startKey)
		eKey := constructLevelKey(p.dbName, endKey)
		itr = p.db.GetIterator(sKey, eKey)
	}
	if err := itr.Error(); err != nil {
		itr.Release()
		return nil, errors.Wrapf(err, "internal leveldb error while obtaining db iterator")
	}
	return &Iterator{p.dbName, itr}, nil
}

func (p *DBProvider) PrefixQuery(key []byte) (*Iterator, error) {
	k := constructLevelKey(p.dbName, key)
	//log.Debugf("Getting iterator for range [%#v] - [%#v]", sKey, eKey)
	itr := p.db.PrefixQuery(k)
	if err := itr.Error(); err != nil {
		itr.Release()
		return nil, errors.Wrapf(err, "internal leveldb error while obtaining db iterator")
	}
	return &Iterator{p.dbName, itr}, nil
}
func (p *DBProvider) GetSnapshot() (*Snapshot, error) {
	itr, err := p.db.db.GetSnapshot()
	if err != nil {
		return nil, errors.Wrap(err, "error writing batch to leveldb")
	}

	return &Snapshot{
		dbName:   p.dbName,
		snapshot: itr,
	}, nil
}

func (sn *Snapshot) Get(key []byte) ([]byte, error) {
	k := constructLevelKey(sn.dbName, key)
	value, err := sn.snapshot.Get(k, sn.readOpts)
	if err != nil {
		return nil, errors.Wrap(err, "error writing batch to leveldb")
	}
	return value, nil
}
func (sn *Snapshot) Release() {
	sn.snapshot.Release()
}

func (sn *Snapshot) NewIterator(startKey []byte, endKey []byte) (*Iterator, error) {
	var itr iterator.Iterator
	if startKey == nil && endKey == nil {
		itr = sn.snapshot.NewIterator(&goleveldbutil.Range{Start: startKey, Limit: endKey}, sn.readOpts)
	} else {
		sKey := constructLevelKey(sn.dbName, startKey)
		eKey := constructLevelKey(sn.dbName, endKey)
		itr = sn.snapshot.NewIterator(&goleveldbutil.Range{Start: sKey, Limit: eKey}, sn.readOpts)
	}

	if err := itr.Error(); err != nil {
		itr.Release()
		return nil, errors.Wrapf(err, "internal leveldb error while obtaining db iterator")
	}

	return &Iterator{sn.dbName, itr}, nil
}

// Close closes the DBProvider after its db data have been deleted
func (p *DBProvider) Close() {
	p.db.Close()
}

// Put adds a KV
func (b *UpdateBatch) Put(key []byte, value []byte) {
	if value == nil {
		panic("Nil value not allowed")
	}
	k := constructLevelKey(b.dbName, key)
	b.leveldbBatch.Put(k, value)
	b.size += len(k) + len(value)
}

// Delete deletes a Key and associated value
func (b *UpdateBatch) Delete(key []byte) {
	k := constructLevelKey(b.dbName, key)
	b.size += len(k)
	b.leveldbBatch.Delete(k)
}

// Size returns the current size of the batch
func (b *UpdateBatch) Size() int {
	return b.size
}

// Len returns number of records in the batch
func (b *UpdateBatch) Len() int {
	return b.leveldbBatch.Len()
}

// Reset resets the batch
func (b *UpdateBatch) Reset() {
	b.leveldbBatch.Reset()
	b.size = 0
}

// Seek moves the iterator to the first key/value pair
// whose key is greater than or equal to the given key.
// It returns whether such pair exist.
func (itr *Iterator) Seek(key []byte) bool {
	levelKey := constructLevelKey(itr.dbName, key)
	return itr.Iterator.Seek(levelKey)
}

func constructLevelKey(dbName string, key []byte) []byte {
	return append([]byte(dbName), key...)
}