database/leveldb/mem_db.go

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