database/leveldb/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) 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 }