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
|
// Copyright 2020 Andrew Thornton. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
package levelqueue
import (
"sync"
"github.com/syndtr/goleveldb/leveldb"
"github.com/syndtr/goleveldb/leveldb/errors"
"github.com/syndtr/goleveldb/leveldb/util"
)
const (
setPrefixStr = "set"
)
// Set defines a set struct
type Set struct {
db *leveldb.DB
closeUnderlyingDB bool
lock sync.Mutex
prefix []byte
}
// OpenSet opens a set from the db path or creates a set if it doesn't exist.
// The keys will be prefixed with "set-" by default
func OpenSet(dataDir string) (*Set, error) {
db, err := leveldb.OpenFile(dataDir, nil)
if err != nil {
if !errors.IsCorrupted(err) {
return nil, err
}
db, err = leveldb.RecoverFile(dataDir, nil)
if err != nil {
return nil, err
}
}
return NewSet(db, []byte(setPrefixStr), true)
}
// NewSet creates a set from a db. The keys will be prefixed with prefix
// and at close the db will be closed as per closeUnderlyingDB
func NewSet(db *leveldb.DB, prefix []byte, closeUnderlyingDB bool) (*Set, error) {
set := &Set{
db: db,
closeUnderlyingDB: closeUnderlyingDB,
}
set.prefix = make([]byte, len(prefix))
copy(set.prefix, prefix)
return set, nil
}
// Add adds a member string to a key set, returns true if the member was not already present
func (set *Set) Add(value []byte) (bool, error) {
set.lock.Lock()
defer set.lock.Unlock()
setKey := withPrefix(set.prefix, value)
has, err := set.db.Has(setKey, nil)
if err != nil || has {
return !has, err
}
return !has, set.db.Put(setKey, []byte(""), nil)
}
// Members returns the current members of the set
func (set *Set) Members() ([][]byte, error) {
set.lock.Lock()
defer set.lock.Unlock()
var members [][]byte
prefix := withPrefix(set.prefix, []byte{})
iter := set.db.NewIterator(util.BytesPrefix(prefix), nil)
for iter.Next() {
slice := iter.Key()[len(prefix):]
value := make([]byte, len(slice))
copy(value, slice)
members = append(members, value)
}
iter.Release()
return members, iter.Error()
}
// Has returns if the member is in the set
func (set *Set) Has(value []byte) (bool, error) {
set.lock.Lock()
defer set.lock.Unlock()
setKey := withPrefix(set.prefix, value)
return set.db.Has(setKey, nil)
}
// Remove removes a member from the set, returns true if the member was present
func (set *Set) Remove(value []byte) (bool, error) {
set.lock.Lock()
defer set.lock.Unlock()
setKey := withPrefix(set.prefix, value)
has, err := set.db.Has(setKey, nil)
if err != nil || !has {
return has, err
}
return has, set.db.Delete(setKey, nil)
}
// Close closes the set (and the underlying db if set to closeUnderlyingDB)
func (set *Set) Close() error {
if !set.closeUnderlyingDB {
set.db = nil
return nil
}
err := set.db.Close()
set.db = nil
return err
}
|
