database/leveldb/mem_db.go

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