1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112
113. 113
114. 114
115. 115
116. 116
117. 117
118. 118
119. 119

// 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 {
	set.lock.Lock()
	defer set.lock.Unlock()
	if !set.closeUnderlyingDB {
		set.db = nil
		return nil
	}
	err := set.db.Close()
	set.db = nil
	return err
}