forked from lonelyevil/kook
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sn_store.go
49 lines (40 loc) · 873 Bytes
/
sn_store.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
package khl
import (
"encoding/binary"
"sync"
"github.com/bits-and-blooms/bloom/v3"
)
// SnStore is the interface for storing sequence numbers.
type SnStore interface {
TestAndInsert(int64) bool
Lock()
Unlock()
Clear()
}
type bloomSnStore struct {
filter *bloom.BloomFilter
lock *sync.Mutex
}
func newBloomSnStore() bloomSnStore {
return bloomSnStore{
filter: bloom.NewWithEstimates(1000000, 0.01),
lock: &sync.Mutex{},
}
}
// TestAndInsert insert a number and check if it exists.
func (b bloomSnStore) TestAndInsert(i int64) bool {
buf := make([]byte, 8)
binary.BigEndian.PutUint64(buf, uint64(i))
return b.filter.TestAndAdd(buf)
}
// Lock locks the store.
func (b bloomSnStore) Lock() {
b.lock.Lock()
}
// Unlock unlocks the store.
func (b bloomSnStore) Unlock() {
b.lock.Unlock()
}
func (b bloomSnStore) Clear() {
b.filter.ClearAll()
}