forked from evilsocket/islazy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
unsortedkv.go
168 lines (146 loc) · 3.78 KB
/
unsortedkv.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
package data
import (
"bytes"
"encoding/gob"
"encoding/json"
"io/ioutil"
"os"
"sync"
"github.com/guozhaoyun/islazy/fs"
)
// UnsortedKV is a thread safe and unsorted key-value
// storage with optional persistency on disk.
type UnsortedKV struct {
sync.RWMutex
fileName string
m map[string]string
policy FlushPolicy
}
// NewUnsortedKV creates a new UnsortedKV with the given flush policy.
// If fileName already exists, it will be deserialized and loaded.
func NewUnsortedKV(fileName string, flushPolicy FlushPolicy) (*UnsortedKV, error) {
ukv := &UnsortedKV{
fileName: fileName,
m: make(map[string]string),
policy: flushPolicy,
}
if fileName != "" && fs.Exists(fileName) {
raw, err := ioutil.ReadFile(fileName)
if err != nil {
return nil, err
}
decoder := gob.NewDecoder(bytes.NewReader(raw))
if err = decoder.Decode(&ukv.m); err != nil {
return nil, err
}
}
return ukv, nil
}
// NewDiskUnsortedKV returns an UnsortedKV that flushed data on disk
// every time it gets updated.
func NewDiskUnsortedKV(fileName string) (*UnsortedKV, error) {
return NewUnsortedKV(fileName, FlushOnEdit)
}
// NewDiskUnsortedKVReader returns an UnsortedKV from disk as a reader
// but it doesn't flush any modifications on disk.
func NewDiskUnsortedKVReader(fileName string) (*UnsortedKV, error) {
return NewUnsortedKV(fileName, FlushNone)
}
// NewMemUnsortedKV returns an UnsortedKV that only lives in
// memory and never persists on disk.
func NewMemUnsortedKV() (*UnsortedKV, error) {
return NewUnsortedKV("", FlushNone)
}
// MarshalJSON is used to serialize the UnsortedKV data structure to
// JSON correctly.
func (u *UnsortedKV) MarshalJSON() ([]byte, error) {
u.RLock()
defer u.RUnlock()
return json.Marshal(u.m)
}
// Has return true if name exists in the store.
func (u *UnsortedKV) Has(name string) bool {
u.RLock()
defer u.RUnlock()
_, found := u.m[name]
return found
}
// Get return the value of the named object if present, or returns
// found as false otherwise.
func (u *UnsortedKV) Get(name string) (v string, found bool) {
u.RLock()
defer u.RUnlock()
v, found = u.m[name]
return
}
// GetOr will return the value of the named object if present,
// or a default value.
func (u *UnsortedKV) GetOr(name, or string) string {
if v, found := u.Get(name); found {
return v
}
return or
}
func (u *UnsortedKV) flushUnlocked() error {
buf := new(bytes.Buffer)
encoder := gob.NewEncoder(buf)
if err := encoder.Encode(u.m); err != nil {
return err
}
return ioutil.WriteFile(u.fileName, buf.Bytes(), os.ModePerm)
}
// Flush flushes the store to disk if the flush policy
// is different than FlushNone
func (u *UnsortedKV) Flush() error {
u.RLock()
defer u.RUnlock()
if u.policy != FlushNone {
return u.flushUnlocked()
}
return nil
}
func (u *UnsortedKV) onEdit() error {
if u.policy == FlushOnEdit {
return u.flushUnlocked()
}
return nil
}
// Set sets a value for a named object.
func (u *UnsortedKV) Set(name, value string) error {
u.Lock()
defer u.Unlock()
u.m[name] = value
return u.onEdit()
}
// Del deletes a named object from the store.
func (u *UnsortedKV) Del(name string) error {
u.Lock()
defer u.Unlock()
delete(u.m, name)
return u.onEdit()
}
// Clear deletes every named object from the store.
func (u *UnsortedKV) Clear() error {
u.Lock()
defer u.Unlock()
u.m = make(map[string]string)
return u.onEdit()
}
// Each iterates each named object in the store by
// executing the callback cb on them, if the callback
// returns true the iteration is interrupted.
func (u *UnsortedKV) Each(cb func(k, v string) bool) {
u.Lock()
defer u.Unlock()
for k, v := range u.m {
if stop := cb(k, v); stop {
return
}
}
}
// Empty returns bool if the store is empty.
func (u *UnsortedKV) Empty() bool {
u.RLock()
defer u.RUnlock()
return len(u.m) == 0
}