kvdb/leveldb.go

package kvdb

import (
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/util"
	"io"
	"sync/atomic"
)

type LevelDb struct {
	filename string
	openFlag int32
	engine   *leveldb.DB
}

type levelDbIterator struct {
	v iterator.Iterator
}

func (i *levelDbIterator) Next() bool {
	return i.v.Next()
}

func (i *levelDbIterator) Key() string {
	return string(i.v.Key())
}

func (i *levelDbIterator) Value() []byte {
	return i.v.Value()
}

func (i *levelDbIterator) Close() error {
	i.v.Release()
	return nil
}

func (db *LevelDb) Open() (err error) {
	db.engine, err = leveldb.OpenFile(db.filename, nil)
	if err == nil {
		atomic.StoreInt32(&db.openFlag, 1)
	} else {
		if db.engine, err = leveldb.RecoverFile(db.filename, nil); err == nil {
			atomic.StoreInt32(&db.openFlag, 1)
		}
	}
	return
}

func (db *LevelDb) Put(key string, value []byte) (err error) {
	if atomic.LoadInt32(&db.openFlag) == 0 {
		err = io.ErrClosedPipe
		return
	}
	err = db.engine.Put([]byte(key), value, nil)
	return
}

func (db *LevelDb) Get(key string) (b []byte, err error) {
	if atomic.LoadInt32(&db.openFlag) == 0 {
		err = io.ErrClosedPipe
		return
	}
	b, err = db.engine.Get([]byte(key), nil)
	return
}

func (db *LevelDb) Delete(key string) (err error) {
	if atomic.LoadInt32(&db.openFlag) == 0 {
		err = io.ErrClosedPipe
		return
	}
	err = db.engine.Delete([]byte(key), nil)
	return
}

func (db *LevelDb) Iterator(key string) (iterator Iterator, err error) {
	if atomic.LoadInt32(&db.openFlag) == 0 {
		err = io.ErrClosedPipe
		return
	}
	v := db.engine.NewIterator(util.BytesPrefix([]byte(key)), nil)
	iterator = &levelDbIterator{
		v: v,
	}
	return
}

func (db *LevelDb) Close() error {
	if atomic.LoadInt32(&db.openFlag) == 1 {
		atomic.StoreInt32(&db.openFlag, 0)
		return db.engine.Close()
	}
	return nil
}

func NewLevelDb(filename string) *LevelDb {
	return &LevelDb{
		filename: filename,
	}
}