vendor/github.com/tendermint/tmlibs/db/mem_db.go

   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) Release() {
 105         it.db = nil
 106         it.keys = nil
 107 }
 108
 109 func (it *memDBIterator) Error() error {
 110         return nil
 111 }
 112
 113 func (db *MemDB) Iterator() Iterator {
 114         return db.IteratorPrefix([]byte{})
 115 }
 116
 117 func (db *MemDB) IteratorPrefix(prefix []byte) Iterator {
 118         it := newMemDBIterator()
 119         it.db = db
 120         it.last = -1
 121
 122         db.mtx.Lock()
 123         defer db.mtx.Unlock()
 124
 125         // unfortunately we need a copy of all of the keys
 126         for key, _ := range db.db {
 127                 if strings.HasPrefix(key, string(prefix)) {
 128                         it.keys = append(it.keys, key)
 129                 }
 130         }
 131         // and we need to sort them
 132         sort.Strings(it.keys)
 133         return it
 134 }
 135
 136 func (db *MemDB) NewBatch() Batch {
 137         return &memDBBatch{db, nil}
 138 }
 139
 140 //--------------------------------------------------------------------------------
 141
 142 type memDBBatch struct {
 143         db  *MemDB
 144         ops []operation
 145 }
 146
 147 type opType int
 148
 149 const (
 150         opTypeSet    = 1
 151         opTypeDelete = 2
 152 )
 153
 154 type operation struct {
 155         opType
 156         key   []byte
 157         value []byte
 158 }
 159
 160 func (mBatch *memDBBatch) Set(key, value []byte) {
 161         mBatch.ops = append(mBatch.ops, operation{opTypeSet, key, value})
 162 }
 163
 164 func (mBatch *memDBBatch) Delete(key []byte) {
 165         mBatch.ops = append(mBatch.ops, operation{opTypeDelete, key, nil})
 166 }
 167
 168 func (mBatch *memDBBatch) Write() {
 169         mBatch.db.mtx.Lock()
 170         defer mBatch.db.mtx.Unlock()
 171
 172         for _, op := range mBatch.ops {
 173                 if op.opType == opTypeSet {
 174                         mBatch.db.db[string(op.key)] = op.value
 175                 } else if op.opType == opTypeDelete {
 176                         delete(mBatch.db.db, string(op.key))
 177                 }
 178         }
 179
 180 }