-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
boltdd.go
417 lines (343 loc) · 10.3 KB
/
boltdd.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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
// BOLTdd contains a wrapper around BoltDB to deduplicate writes and encode
// values using mgspack. (dd stands for DeDuplicate)
package boltdd
import (
"bytes"
"fmt"
"os"
"sync"
"github.com/boltdb/bolt"
"github.com/hashicorp/go-msgpack/codec"
"github.com/hashicorp/nomad/nomad/structs"
"golang.org/x/crypto/blake2b"
)
// ErrNotFound is returned when a key is not found.
type ErrNotFound struct {
name string
}
func (e *ErrNotFound) Error() string {
return fmt.Sprintf("key not found: %s", e.name)
}
// NotFound returns a new error for a key that was not found.
func NotFound(name string) error {
return &ErrNotFound{name}
}
// IsErrNotFound returns true if the error is an ErrNotFound error.
func IsErrNotFound(e error) bool {
if e == nil {
return false
}
_, ok := e.(*ErrNotFound)
return ok
}
// DB wraps an underlying bolt.DB to create write deduplicating buckets and
// msgpack encoded values.
type DB struct {
rootBuckets map[string]*bucketMeta
rootBucketsLock sync.Mutex
bdb *bolt.DB
}
// Open a bolt.DB and wrap it in a write-deduplicating msgpack-encoding
// implementation.
func Open(path string, mode os.FileMode, options *bolt.Options) (*DB, error) {
bdb, err := bolt.Open(path, mode, options)
if err != nil {
return nil, err
}
return New(bdb), nil
}
// New deduplicating wrapper for the given boltdb.
func New(bdb *bolt.DB) *DB {
return &DB{
rootBuckets: make(map[string]*bucketMeta),
bdb: bdb,
}
}
func (db *DB) bucket(btx *bolt.Tx, name []byte) *Bucket {
bb := btx.Bucket(name)
if bb == nil {
return nil
}
db.rootBucketsLock.Lock()
defer db.rootBucketsLock.Unlock()
if db.isClosed() {
return nil
}
b, ok := db.rootBuckets[string(name)]
if !ok {
b = newBucketMeta()
db.rootBuckets[string(name)] = b
}
return newBucket(b, bb)
}
func (db *DB) createBucket(btx *bolt.Tx, name []byte) (*Bucket, error) {
bb, err := btx.CreateBucket(name)
if err != nil {
return nil, err
}
db.rootBucketsLock.Lock()
defer db.rootBucketsLock.Unlock()
// While creating a bucket on a closed db would error, we must recheck
// after acquiring the lock to avoid races.
if db.isClosed() {
return nil, bolt.ErrDatabaseNotOpen
}
// Always create a new Bucket since CreateBucket above fails if the
// bucket already exists.
b := newBucketMeta()
db.rootBuckets[string(name)] = b
return newBucket(b, bb), nil
}
func (db *DB) createBucketIfNotExists(btx *bolt.Tx, name []byte) (*Bucket, error) {
bb, err := btx.CreateBucketIfNotExists(name)
if err != nil {
return nil, err
}
db.rootBucketsLock.Lock()
defer db.rootBucketsLock.Unlock()
// While creating a bucket on a closed db would error, we must recheck
// after acquiring the lock to avoid races.
if db.isClosed() {
return nil, bolt.ErrDatabaseNotOpen
}
b, ok := db.rootBuckets[string(name)]
if !ok {
b = newBucketMeta()
db.rootBuckets[string(name)] = b
}
return newBucket(b, bb), nil
}
func (db *DB) Update(fn func(*Tx) error) error {
return db.bdb.Update(func(btx *bolt.Tx) error {
tx := newTx(db, btx)
return fn(tx)
})
}
func (db *DB) View(fn func(*Tx) error) error {
return db.bdb.View(func(btx *bolt.Tx) error {
tx := newTx(db, btx)
return fn(tx)
})
}
// isClosed returns true if the database is closed and must be called while
// db.rootBucketsLock is acquired.
func (db *DB) isClosed() bool {
return db.rootBuckets == nil
}
// Close closes the underlying bolt.DB and clears all bucket hashes. DB is
// unusable after closing.
func (db *DB) Close() error {
db.rootBucketsLock.Lock()
db.rootBuckets = nil
db.rootBucketsLock.Unlock()
return db.bdb.Close()
}
// BoltDB returns the underlying bolt.DB.
func (db *DB) BoltDB() *bolt.DB {
return db.bdb
}
type Tx struct {
db *DB
btx *bolt.Tx
}
func newTx(db *DB, btx *bolt.Tx) *Tx {
return &Tx{
db: db,
btx: btx,
}
}
// Bucket returns a root bucket or nil if it doesn't exist.
func (tx *Tx) Bucket(name []byte) *Bucket {
return tx.db.bucket(tx.btx, name)
}
func (tx *Tx) CreateBucket(name []byte) (*Bucket, error) {
return tx.db.createBucket(tx.btx, name)
}
// CreateBucketIfNotExists returns a root bucket or creates a new one if it
// doesn't already exist.
func (tx *Tx) CreateBucketIfNotExists(name []byte) (*Bucket, error) {
return tx.db.createBucketIfNotExists(tx.btx, name)
}
// Writable wraps boltdb Tx.Writable.
func (tx *Tx) Writable() bool {
return tx.btx.Writable()
}
// BoltTx returns the underlying bolt.Tx.
func (tx *Tx) BoltTx() *bolt.Tx {
return tx.btx
}
// bucketMeta persists metadata -- such as key hashes and child buckets --
// about boltdb Buckets across transactions.
type bucketMeta struct {
// hashes holds all of the value hashes for keys in this bucket
hashes map[string][]byte
hashesLock sync.Mutex
// buckets holds all of the child buckets
buckets map[string]*bucketMeta
bucketsLock sync.Mutex
}
func newBucketMeta() *bucketMeta {
return &bucketMeta{
hashes: make(map[string][]byte),
buckets: make(map[string]*bucketMeta),
}
}
// getHash of last value written to a key or nil if no hash exists.
func (bm *bucketMeta) getHash(hashKey string) []byte {
bm.hashesLock.Lock()
lastHash := bm.hashes[hashKey]
bm.hashesLock.Unlock()
return lastHash
}
// setHash of last value written to key.
func (bm *bucketMeta) setHash(hashKey string, hashVal []byte) {
bm.hashesLock.Lock()
bm.hashes[hashKey] = hashVal
bm.hashesLock.Unlock()
}
// delHash deletes a hash value or does nothing if the hash key does not exist.
func (bm *bucketMeta) delHash(hashKey string) {
bm.hashesLock.Lock()
delete(bm.hashes, hashKey)
bm.hashesLock.Unlock()
}
// createBucket metadata entry for the given nested bucket. Overwrites any
// existing entry so caller should ensure bucket does not already exist.
func (bm *bucketMeta) createBucket(name []byte) *bucketMeta {
bm.bucketsLock.Lock()
defer bm.bucketsLock.Unlock()
// Always create a new Bucket since CreateBucket above fails if the
// bucket already exists.
b := newBucketMeta()
bm.buckets[string(name)] = b
return b
}
// deleteBucket metadata entry for the given nested bucket. Does nothing if
// nested bucket metadata does not exist.
func (bm *bucketMeta) deleteBucket(name []byte) {
bm.bucketsLock.Lock()
delete(bm.buckets, string(name))
bm.bucketsLock.Unlock()
}
// getOrCreateBucket metadata entry for the given nested bucket.
func (bm *bucketMeta) getOrCreateBucket(name []byte) *bucketMeta {
bm.bucketsLock.Lock()
defer bm.bucketsLock.Unlock()
b, ok := bm.buckets[string(name)]
if !ok {
b = newBucketMeta()
bm.buckets[string(name)] = b
}
return b
}
type Bucket struct {
bm *bucketMeta
boltBucket *bolt.Bucket
}
// newBucket creates a new view into a bucket backed by a boltdb
// transaction.
func newBucket(b *bucketMeta, bb *bolt.Bucket) *Bucket {
return &Bucket{
bm: b,
boltBucket: bb,
}
}
// Put into boltdb iff it has changed since the last write.
func (b *Bucket) Put(key []byte, val interface{}) error {
// buffer for writing serialized state to
var buf bytes.Buffer
// Serialize the object
if err := codec.NewEncoder(&buf, structs.MsgpackHandle).Encode(val); err != nil {
return fmt.Errorf("failed to encode passed object: %v", err)
}
// Hash for skipping unnecessary writes
hashKey := string(key)
hashVal := blake2b.Sum256(buf.Bytes())
// lastHash value or nil if it hasn't been hashed yet
lastHash := b.bm.getHash(hashKey)
// If the hashes are equal, skip the write
if bytes.Equal(hashVal[:], lastHash) {
return nil
}
// New value: write it to the underlying boltdb
if err := b.boltBucket.Put(key, buf.Bytes()); err != nil {
return fmt.Errorf("failed to write data at key %s: %v", key, err)
}
// New value written, store hash (bucket path map was created above)
b.bm.setHash(hashKey, hashVal[:])
return nil
}
// Get value by key from boltdb or return an ErrNotFound error if key not
// found.
func (b *Bucket) Get(key []byte, obj interface{}) error {
// Get the raw data from the underlying boltdb
data := b.boltBucket.Get(key)
if data == nil {
return NotFound(string(key))
}
// Deserialize the object
if err := codec.NewDecoderBytes(data, structs.MsgpackHandle).Decode(obj); err != nil {
return fmt.Errorf("failed to decode data into passed object: %v", err)
}
return nil
}
// Delete removes a key from the bucket. If the key does not exist then nothing
// is done and a nil error is returned. Returns an error if the bucket was
// created from a read-only transaction.
func (b *Bucket) Delete(key []byte) error {
err := b.boltBucket.Delete(key)
b.bm.delHash(string(key))
return err
}
// Bucket represents a boltdb Bucket and its associated metadata necessary for
// write deduplication. Like bolt.Buckets it is only valid for the duration of
// the transaction that created it.
func (b *Bucket) Bucket(name []byte) *Bucket {
bb := b.boltBucket.Bucket(name)
if bb == nil {
return nil
}
bmeta := b.bm.getOrCreateBucket(name)
return newBucket(bmeta, bb)
}
// CreateBucket creates a new bucket at the given key and returns the new
// bucket. Returns an error if the key already exists, if the bucket name is
// blank, or if the bucket name is too long. The bucket instance is only valid
// for the lifetime of the transaction.
func (b *Bucket) CreateBucket(name []byte) (*Bucket, error) {
bb, err := b.boltBucket.CreateBucket(name)
if err != nil {
return nil, err
}
bmeta := b.bm.createBucket(name)
return newBucket(bmeta, bb), nil
}
// CreateBucketIfNotExists creates a new bucket if it doesn't already exist and
// returns a reference to it. The bucket instance is only valid for the
// lifetime of the transaction.
func (b *Bucket) CreateBucketIfNotExists(name []byte) (*Bucket, error) {
bb, err := b.boltBucket.CreateBucketIfNotExists(name)
if err != nil {
return nil, err
}
bmeta := b.bm.getOrCreateBucket(name)
return newBucket(bmeta, bb), nil
}
// DeleteBucket deletes a child bucket. Returns an error if the bucket
// corresponds to a non-bucket key or another error is encountered. No error is
// returned if the bucket does not exist.
func (b *Bucket) DeleteBucket(name []byte) error {
// Delete the bucket from the underlying boltdb
err := b.boltBucket.DeleteBucket(name)
if err == bolt.ErrBucketNotFound {
err = nil
}
// Remove reference to child bucket
b.bm.deleteBucket(name)
return err
}
// BoltBucket returns the internal bolt.Bucket for this Bucket. Only valid
// for the duration of the current transaction.
func (b *Bucket) BoltBucket() *bolt.Bucket {
return b.boltBucket
}