github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/database/leveldb/mem_db.go

package leveldb

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"sync"
)

func init() {
	registerDBCreator(MemDBBackendStr, func(name string, dir string) (DB, error) {
		return NewMemDB(), nil
	}, false)
}

type MemDB struct {
	mtx sync.Mutex
	db  map[string][]byte
}

func NewMemDB() *MemDB {
	database := &MemDB{db: make(map[string][]byte)}
	return database
}

func (db *MemDB) Get(key []byte) []byte {
	db.mtx.Lock()
	defer db.mtx.Unlock()
	return db.db[string(key)]
}

func (db *MemDB) Set(key []byte, value []byte) {
	db.mtx.Lock()
	defer db.mtx.Unlock()
	db.db[string(key)] = value
}

func (db *MemDB) SetSync(key []byte, value []byte) {
	db.mtx.Lock()
	defer db.mtx.Unlock()
	db.db[string(key)] = value
}

func (db *MemDB) Delete(key []byte) {
	db.mtx.Lock()
	defer db.mtx.Unlock()
	delete(db.db, string(key))
}

func (db *MemDB) DeleteSync(key []byte) {
	db.mtx.Lock()
	defer db.mtx.Unlock()
	delete(db.db, string(key))
}

func (db *MemDB) Close() {
	// Close is a noop since for an in-memory
	// database, we don't have a destination
	// to flush contents to nor do we want
	// any data loss on invoking Close()
	// See the discussion in https://github.com/tendermint/tmlibs/pull/56
}

func (db *MemDB) Print() {
	db.mtx.Lock()
	defer db.mtx.Unlock()
	for key, value := range db.db {
		fmt.Printf("[%X]:\t[%X]\n", []byte(key), value)
	}
}

func (db *MemDB) Stats() map[string]string {
	stats := make(map[string]string)
	stats["database.type"] = "memDB"
	return stats
}

type memDBIterator struct {
	last int
	keys []string
	db   DB

	start []byte
}

func newMemDBIterator() *memDBIterator {
	return &memDBIterator{}
}

// Keys is expected to be in reverse order for reverse iterators.
func newMemDBIteratorWithArgs(db DB, keys []string, start []byte) *memDBIterator {
	itr := &memDBIterator{
		db:    db,
		keys:  keys,
		start: start,
		last:  -1,
	}
	if start != nil {
		itr.Seek(start)
	}
	return itr
}

func (it *memDBIterator) Next() bool {
	if it.last >= len(it.keys)-1 {
		return false
	}
	it.last++
	return true
}

func (it *memDBIterator) Key() []byte {
	if it.last < 0 {
		return []byte("")
	}
	return []byte(it.keys[it.last])
}

func (it *memDBIterator) Value() []byte {
	return it.db.Get(it.Key())
}

func (it *memDBIterator) Seek(point []byte) bool {
	for i, key := range it.keys {
		if key >= string(point) {
			it.last = i
			return true
		}
	}
	return false
}

func (it *memDBIterator) Release() {
	it.db = nil
	it.keys = nil
}

func (it *memDBIterator) Error() error {
	return nil
}

func (db *MemDB) Iterator() Iterator {
	return db.IteratorPrefix([]byte{})
}

func (db *MemDB) IteratorPrefix(prefix []byte) Iterator {
	it := newMemDBIterator()
	it.db = db
	it.last = -1

	db.mtx.Lock()
	defer db.mtx.Unlock()

	// unfortunately we need a copy of all of the keys
	for key, _ := range db.db {
		if strings.HasPrefix(key, string(prefix)) {
			it.keys = append(it.keys, key)
		}
	}
	// and we need to sort them
	sort.Strings(it.keys)
	return it
}

func (db *MemDB) IteratorPrefixWithStart(Prefix, start []byte, isReverse bool) Iterator {
	db.mtx.Lock()
	defer db.mtx.Unlock()

	keys := db.getSortedKeys(start, isReverse)
	return newMemDBIteratorWithArgs(db, keys, start)
}

func (db *MemDB) NewBatch() Batch {
	return &memDBBatch{db, nil}
}

func (db *MemDB) getSortedKeys(start []byte, reverse bool) []string {
	keys := []string{}
	for key := range db.db {
		if bytes.Compare([]byte(key), start) < 0 {
			continue
		}
		keys = append(keys, key)
	}
	sort.Strings(keys)
	if reverse {
		nkeys := len(keys)
		for i := 0; i < nkeys/2; i++ {
			temp := keys[i]
			keys[i] = keys[nkeys-i-1]
			keys[nkeys-i-1] = temp
		}
	}
	return keys
}

//--------------------------------------------------------------------------------

type memDBBatch struct {
	db  *MemDB
	ops []operation
}

type opType int

const (
	opTypeSet    = 1
	opTypeDelete = 2
)

type operation struct {
	opType
	key   []byte
	value []byte
}

func (mBatch *memDBBatch) Set(key, value []byte) {
	mBatch.ops = append(mBatch.ops, operation{opTypeSet, key, value})
}

func (mBatch *memDBBatch) Delete(key []byte) {
	mBatch.ops = append(mBatch.ops, operation{opTypeDelete, key, nil})
}

func (mBatch *memDBBatch) Write() {
	mBatch.db.mtx.Lock()
	defer mBatch.db.mtx.Unlock()

	for _, op := range mBatch.ops {
		if op.opType == opTypeSet {
			mBatch.db.db[string(op.key)] = op.value
		} else if op.opType == opTypeDelete {
			delete(mBatch.db.db, string(op.key))
		}
	}

}