github.com/Bytom/bytom@v1.1.2-0.20210127130405-ae40204c0b09/database/leveldb/mem_db.go (about)

     1  package db
     2  
     3  import (
     4  	"fmt"
     5  	"sort"
     6  	"strings"
     7  	"sync"
     8  )
     9  
    10  func init() {
    11  	registerDBCreator(MemDBBackendStr, func(name string, dir string) (DB, error) {
    12  		return NewMemDB(), nil
    13  	}, false)
    14  }
    15  
    16  type MemDB struct {
    17  	mtx sync.Mutex
    18  	db  map[string][]byte
    19  }
    20  
    21  func NewMemDB() *MemDB {
    22  	database := &MemDB{db: make(map[string][]byte)}
    23  	return database
    24  }
    25  
    26  func (db *MemDB) Get(key []byte) []byte {
    27  	db.mtx.Lock()
    28  	defer db.mtx.Unlock()
    29  	return db.db[string(key)]
    30  }
    31  
    32  func (db *MemDB) Set(key []byte, value []byte) {
    33  	db.mtx.Lock()
    34  	defer db.mtx.Unlock()
    35  	db.db[string(key)] = value
    36  }
    37  
    38  func (db *MemDB) SetSync(key []byte, value []byte) {
    39  	db.mtx.Lock()
    40  	defer db.mtx.Unlock()
    41  	db.db[string(key)] = value
    42  }
    43  
    44  func (db *MemDB) Delete(key []byte) {
    45  	db.mtx.Lock()
    46  	defer db.mtx.Unlock()
    47  	delete(db.db, string(key))
    48  }
    49  
    50  func (db *MemDB) DeleteSync(key []byte) {
    51  	db.mtx.Lock()
    52  	defer db.mtx.Unlock()
    53  	delete(db.db, string(key))
    54  }
    55  
    56  func (db *MemDB) Close() {
    57  	// Close is a noop since for an in-memory
    58  	// database, we don't have a destination
    59  	// to flush contents to nor do we want
    60  	// any data loss on invoking Close()
    61  	// See the discussion in https://github.com/tendermint/tmlibs/pull/56
    62  }
    63  
    64  func (db *MemDB) Print() {
    65  	db.mtx.Lock()
    66  	defer db.mtx.Unlock()
    67  	for key, value := range db.db {
    68  		fmt.Printf("[%X]:\t[%X]\n", []byte(key), value)
    69  	}
    70  }
    71  
    72  func (db *MemDB) Stats() map[string]string {
    73  	stats := make(map[string]string)
    74  	stats["database.type"] = "memDB"
    75  	return stats
    76  }
    77  
    78  type memDBIterator struct {
    79  	last int
    80  	keys []string
    81  	db   *MemDB
    82  }
    83  
    84  func newMemDBIterator() *memDBIterator {
    85  	return &memDBIterator{}
    86  }
    87  
    88  func (it *memDBIterator) Next() bool {
    89  	if it.last >= len(it.keys)-1 {
    90  		return false
    91  	}
    92  	it.last++
    93  	return true
    94  }
    95  
    96  func (it *memDBIterator) Key() []byte {
    97  	return []byte(it.keys[it.last])
    98  }
    99  
   100  func (it *memDBIterator) Value() []byte {
   101  	return it.db.Get(it.Key())
   102  }
   103  
   104  func (it *memDBIterator) Seek(point []byte) bool {
   105  	for i, key := range it.keys {
   106  		if key >= string(point) {
   107  			it.last = i
   108  			return true
   109  		}
   110  	}
   111  	return false
   112  }
   113  
   114  func (it *memDBIterator) Release() {
   115  	it.db = nil
   116  	it.keys = nil
   117  }
   118  
   119  func (it *memDBIterator) Error() error {
   120  	return nil
   121  }
   122  
   123  func (db *MemDB) Iterator() Iterator {
   124  	return db.IteratorPrefix([]byte{})
   125  }
   126  
   127  func (db *MemDB) IteratorPrefix(prefix []byte) Iterator {
   128  	it := newMemDBIterator()
   129  	it.db = db
   130  	it.last = -1
   131  
   132  	db.mtx.Lock()
   133  	defer db.mtx.Unlock()
   134  
   135  	// unfortunately we need a copy of all of the keys
   136  	for key, _ := range db.db {
   137  		if strings.HasPrefix(key, string(prefix)) {
   138  			it.keys = append(it.keys, key)
   139  		}
   140  	}
   141  	// and we need to sort them
   142  	sort.Strings(it.keys)
   143  	return it
   144  }
   145  
   146  func (db *MemDB) NewBatch() Batch {
   147  	return &memDBBatch{db, nil}
   148  }
   149  
   150  //--------------------------------------------------------------------------------
   151  
   152  type memDBBatch struct {
   153  	db  *MemDB
   154  	ops []operation
   155  }
   156  
   157  type opType int
   158  
   159  const (
   160  	opTypeSet    = 1
   161  	opTypeDelete = 2
   162  )
   163  
   164  type operation struct {
   165  	opType
   166  	key   []byte
   167  	value []byte
   168  }
   169  
   170  func (mBatch *memDBBatch) Set(key, value []byte) {
   171  	mBatch.ops = append(mBatch.ops, operation{opTypeSet, key, value})
   172  }
   173  
   174  func (mBatch *memDBBatch) Delete(key []byte) {
   175  	mBatch.ops = append(mBatch.ops, operation{opTypeDelete, key, nil})
   176  }
   177  
   178  func (mBatch *memDBBatch) Write() {
   179  	mBatch.db.mtx.Lock()
   180  	defer mBatch.db.mtx.Unlock()
   181  
   182  	for _, op := range mBatch.ops {
   183  		if op.opType == opTypeSet {
   184  			mBatch.db.db[string(op.key)] = op.value
   185  		} else if op.opType == opTypeDelete {
   186  			delete(mBatch.db.db, string(op.key))
   187  		}
   188  	}
   189  
   190  }