asmt/db/memorydb.go
db/memorydb.goRaw
1/**
2 * @file
3 * @copyright defined in aergo/LICENSE.txt
4 */
5
6package db
7
8import (
9 "bytes"
10 "container/list"
11 "encoding/gob"
12 "os"
13 "path"
14 "sort"
15 "sync"
16)
17
18// This function is always called first
19func init() {
20 dbConstructor := func(dir string) (DB, error) {
21 return newMemoryDB(dir)
22 }
23 registorDBConstructor(MemoryImpl, dbConstructor)
24}
25
26func newMemoryDB(dir string) (DB, error) {
27 var db map[string][]byte
28
29 filePath := path.Join(dir, "database")
30
31 file, err := os.Open(filePath)
32 if err == nil {
33 decoder := gob.NewDecoder(file) //
34 err = decoder.Decode(&db)
35
36 if err != nil {
37 return nil, err
38 }
39 }
40
41 file.Close()
42
43 if db == nil {
44 db = make(map[string][]byte)
45 }
46
47 database := &memorydb{
48 db: db,
49 dir: filePath,
50 }
51
52 return database, nil
53}
54
55//=========================================================
56// DB Implementation
57//=========================================================
58
59// Enforce database and transaction implements interfaces
60var _ DB = (*memorydb)(nil)
61
62type memorydb struct {
63 lock sync.Mutex
64 db map[string][]byte
65 dir string
66}
67
68func (db *memorydb) Type() string {
69 return "memorydb"
70}
71
72func (db *memorydb) Set(key, value []byte) {
73 db.lock.Lock()
74 defer db.lock.Unlock()
75
76 key = convNilToBytes(key)
77 value = convNilToBytes(value)
78
79 db.db[string(key)] = value
80}
81
82func (db *memorydb) Delete(key []byte) {
83 db.lock.Lock()
84 defer db.lock.Unlock()
85
86 key = convNilToBytes(key)
87
88 delete(db.db, string(key))
89}
90
91func (db *memorydb) Get(key []byte) []byte {
92 db.lock.Lock()
93 defer db.lock.Unlock()
94
95 key = convNilToBytes(key)
96
97 return db.db[string(key)]
98}
99
100func (db *memorydb) Exist(key []byte) bool {
101 db.lock.Lock()
102 defer db.lock.Unlock()
103
104 key = convNilToBytes(key)
105
106 _, ok := db.db[string(key)]
107
108 return ok
109}
110
111func (db *memorydb) Close() {
112 db.lock.Lock()
113 defer db.lock.Unlock()
114
115 file, err := os.OpenFile(db.dir, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
116 if err == nil {
117 encoder := gob.NewEncoder(file)
118 encoder.Encode(db.db)
119 }
120 file.Close()
121}
122
123func (db *memorydb) NewTx() Transaction {
124
125 return &memoryTransaction{
126 db: db,
127 opList: list.New(),
128 isDiscard: false,
129 isCommit: false,
130 }
131}
132
133func (db *memorydb) NewBulk() Bulk {
134
135 return &memoryBulk{
136 db: db,
137 opList: list.New(),
138 isDiscard: false,
139 isCommit: false,
140 }
141}
142
143//=========================================================
144// Transaction Implementation
145//=========================================================
146
147type memoryTransaction struct {
148 txLock sync.Mutex
149 db *memorydb
150 opList *list.List
151 isDiscard bool
152 isCommit bool
153}
154
155type txOp struct {
156 isSet bool
157 key []byte
158 value []byte
159}
160
161func (transaction *memoryTransaction) Set(key, value []byte) {
162 transaction.txLock.Lock()
163 defer transaction.txLock.Unlock()
164
165 key = convNilToBytes(key)
166 value = convNilToBytes(value)
167
168 transaction.opList.PushBack(&txOp{true, key, value})
169}
170
171func (transaction *memoryTransaction) Delete(key []byte) {
172 transaction.txLock.Lock()
173 defer transaction.txLock.Unlock()
174
175 key = convNilToBytes(key)
176
177 transaction.opList.PushBack(&txOp{false, key, nil})
178}
179
180func (transaction *memoryTransaction) Commit() {
181 transaction.txLock.Lock()
182 defer transaction.txLock.Unlock()
183
184 if transaction.isDiscard {
185 panic("Commit after dicard tx is not allowed")
186 } else if transaction.isCommit {
187 panic("Commit occures two times")
188 }
189
190 db := transaction.db
191
192 db.lock.Lock()
193 defer db.lock.Unlock()
194
195 for e := transaction.opList.Front(); e != nil; e = e.Next() {
196 op := e.Value.(*txOp)
197 if op.isSet {
198 db.db[string(op.key)] = op.value
199 } else {
200 delete(db.db, string(op.key))
201 }
202 }
203
204 transaction.isCommit = true
205}
206
207func (transaction *memoryTransaction) Discard() {
208 transaction.txLock.Lock()
209 defer transaction.txLock.Unlock()
210
211 transaction.isDiscard = true
212}
213
214//=========================================================
215// Bulk Implementation
216//=========================================================
217
218type memoryBulk struct {
219 txLock sync.Mutex
220 db *memorydb
221 opList *list.List
222 isDiscard bool
223 isCommit bool
224}
225
226func (bulk *memoryBulk) Set(key, value []byte) {
227 bulk.txLock.Lock()
228 defer bulk.txLock.Unlock()
229
230 key = convNilToBytes(key)
231 value = convNilToBytes(value)
232
233 bulk.opList.PushBack(&txOp{true, key, value})
234}
235
236func (bulk *memoryBulk) Delete(key []byte) {
237 bulk.txLock.Lock()
238 defer bulk.txLock.Unlock()
239
240 key = convNilToBytes(key)
241
242 bulk.opList.PushBack(&txOp{false, key, nil})
243}
244
245func (bulk *memoryBulk) Flush() {
246 bulk.txLock.Lock()
247 defer bulk.txLock.Unlock()
248
249 if bulk.isDiscard {
250 panic("Commit after dicard tx is not allowed")
251 } else if bulk.isCommit {
252 panic("Commit occures two times")
253 }
254
255 db := bulk.db
256
257 db.lock.Lock()
258 defer db.lock.Unlock()
259
260 for e := bulk.opList.Front(); e != nil; e = e.Next() {
261 op := e.Value.(*txOp)
262 if op.isSet {
263 db.db[string(op.key)] = op.value
264 } else {
265 delete(db.db, string(op.key))
266 }
267 }
268
269 bulk.isCommit = true
270}
271
272func (bulk *memoryBulk) DiscardLast() {
273 bulk.txLock.Lock()
274 defer bulk.txLock.Unlock()
275
276 bulk.isDiscard = true
277}
278
279//=========================================================
280// Iterator Implementation
281//=========================================================
282
283type memoryIterator struct {
284 start []byte
285 end []byte
286 reverse bool
287 keys []string
288 isInvalid bool
289 cursor int
290 db *memorydb
291}
292
293func isKeyInRange(key []byte, start []byte, end []byte, reverse bool) bool {
294 if reverse {
295 if start != nil && bytes.Compare(start, key) < 0 {
296 return false
297 }
298 if end != nil && bytes.Compare(key, end) <= 0 {
299 return false
300 }
301 return true
302 }
303
304 if bytes.Compare(key, start) < 0 {
305 return false
306 }
307 if end != nil && bytes.Compare(end, key) <= 0 {
308 return false
309 }
310 return true
311
312}
313
314func (db *memorydb) Iterator(start, end []byte) Iterator {
315 db.lock.Lock()
316 defer db.lock.Unlock()
317
318 var reverse bool
319
320 // if end is bigger then start, then reverse order
321 if bytes.Compare(start, end) == 1 {
322 reverse = true
323 } else {
324 reverse = false
325 }
326
327 var keys sort.StringSlice
328
329 for key := range db.db {
330 if isKeyInRange([]byte(key), start, end, reverse) {
331 keys = append(keys, key)
332 }
333 }
334 if reverse {
335 sort.Sort(sort.Reverse(keys))
336 } else {
337 sort.Strings(keys)
338 }
339
340 return &memoryIterator{
341 start: start,
342 end: end,
343 reverse: reverse,
344 isInvalid: false,
345 keys: keys,
346 cursor: 0,
347 db: db,
348 }
349}
350
351func (iter *memoryIterator) Next() {
352 if !iter.Valid() {
353 panic("Iterator is Invalid")
354 }
355
356 iter.cursor++
357}
358
359func (iter *memoryIterator) Valid() bool {
360 // Once invalid, forever invalid.
361 if iter.isInvalid {
362 return false
363 }
364
365 return 0 <= iter.cursor && iter.cursor < len(iter.keys)
366}
367
368func (iter *memoryIterator) Key() (key []byte) {
369 if !iter.Valid() {
370 panic("Iterator is Invalid")
371 }
372
373 return []byte(iter.keys[iter.cursor])
374}
375
376func (iter *memoryIterator) Value() (value []byte) {
377 if !iter.Valid() {
378 panic("Iterator is Invalid")
379 }
380
381 key := []byte(iter.keys[iter.cursor])
382
383 return iter.db.Get(key)
384}