forked from influxdata/influxdb
/
series_set.go
156 lines (130 loc) · 3.61 KB
/
series_set.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
package tsdb
import (
"io"
"sync"
"github.com/RoaringBitmap/roaring"
)
// SeriesIDSet represents a lockable bitmap of series ids.
type SeriesIDSet struct {
sync.RWMutex
bitmap *roaring.Bitmap
}
// NewSeriesIDSet returns a new instance of SeriesIDSet.
func NewSeriesIDSet() *SeriesIDSet {
return &SeriesIDSet{
bitmap: roaring.NewBitmap(),
}
}
// Add adds the series id to the set.
func (s *SeriesIDSet) Add(id uint64) {
s.Lock()
defer s.Unlock()
s.AddNoLock(id)
}
// AddNoLock adds the series id to the set. Add is not safe for use from multiple
// goroutines. Callers must manage synchronization.
func (s *SeriesIDSet) AddNoLock(id uint64) {
s.bitmap.Add(uint32(id))
}
// Contains returns true if the id exists in the set.
func (s *SeriesIDSet) Contains(id uint64) bool {
s.RLock()
x := s.ContainsNoLock(id)
s.RUnlock()
return x
}
// ContainsNoLock returns true if the id exists in the set. ContainsNoLock is
// not safe for use from multiple goroutines. The caller must manage synchronization.
func (s *SeriesIDSet) ContainsNoLock(id uint64) bool {
return s.bitmap.Contains(uint32(id))
}
// Remove removes the id from the set.
func (s *SeriesIDSet) Remove(id uint64) {
s.Lock()
defer s.Unlock()
s.RemoveNoLock(id)
}
// RemoveNoLock removes the id from the set. RemoveNoLock is not safe for use
// from multiple goroutines. The caller must manage synchronization.
func (s *SeriesIDSet) RemoveNoLock(id uint64) {
s.bitmap.Remove(uint32(id))
}
// Cardinality returns the cardinality of the SeriesIDSet.
func (s *SeriesIDSet) Cardinality() uint64 {
s.RLock()
defer s.RUnlock()
return s.bitmap.GetCardinality()
}
// Merge merged the contents of others into s. The caller does not need to
// provide s as an argument, and the contents of s will always be present in s
// after Merge returns.
func (s *SeriesIDSet) Merge(others ...*SeriesIDSet) {
bms := make([]*roaring.Bitmap, 0, len(others)+1)
s.RLock()
bms = append(bms, s.bitmap) // Add ourself.
// Add other bitsets.
for _, other := range others {
other.RLock()
defer other.RUnlock() // Hold until we have merged all the bitmaps
bms = append(bms, other.bitmap)
}
result := roaring.FastOr(bms...)
s.RUnlock()
s.Lock()
s.bitmap = result
s.Unlock()
}
// Equals returns true if other and s are the same set of ids.
func (s *SeriesIDSet) Equals(other *SeriesIDSet) bool {
if s == other {
return true
}
s.RLock()
defer s.RUnlock()
other.RLock()
defer other.RUnlock()
return s.bitmap.Equals(other.bitmap)
}
// AndNot returns a new SeriesIDSet containing elements that were present in s,
// but not present in other.
func (s *SeriesIDSet) AndNot(other *SeriesIDSet) *SeriesIDSet {
s.RLock()
defer s.RUnlock()
other.RLock()
defer other.RUnlock()
return &SeriesIDSet{bitmap: roaring.AndNot(s.bitmap, other.bitmap)}
}
// ForEach calls f for each id in the set.
func (s *SeriesIDSet) ForEach(f func(id uint64)) {
s.RLock()
defer s.RUnlock()
itr := s.bitmap.Iterator()
for itr.HasNext() {
f(uint64(itr.Next()))
}
}
func (s *SeriesIDSet) String() string {
s.RLock()
defer s.RUnlock()
return s.bitmap.String()
}
// Diff removes from s any elements also present in other.
func (s *SeriesIDSet) Diff(other *SeriesIDSet) {
other.RLock()
defer other.RUnlock()
s.Lock()
defer s.Unlock()
s.bitmap = roaring.AndNot(s.bitmap, other.bitmap)
}
// UnmarshalBinary unmarshals data into the set.
func (s *SeriesIDSet) UnmarshalBinary(data []byte) error {
s.Lock()
defer s.Unlock()
return s.bitmap.UnmarshalBinary(data)
}
// WriteTo writes the set to w.
func (s *SeriesIDSet) WriteTo(w io.Writer) (int64, error) {
s.RLock()
defer s.RUnlock()
return s.bitmap.WriteTo(w)
}