/
bloomfilter_redis.go
93 lines (86 loc) · 2.03 KB
/
bloomfilter_redis.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
package bloomfilter
import (
"context"
"fmt"
"github.com/redis/go-redis/v9"
)
type RedisClient interface {
Pipeline() redis.Pipeliner
}
type redisBloomFilter struct {
client RedisClient
key string
bucketCount uint64
bucketMaxLen uint64
hashFuncs []func([]byte) uint64
}
func NewRedisBloomFilter(redisCli RedisClient, redisKey string, size uint64) BloomFilter {
if redisKey == "" {
redisKey = fmt.Sprintf("_bloomfilter_%d", size)
}
var bitMaxLen uint64 = 1000000
if size <= 0 {
size = bitMaxLen
}
bucketCount := size / bitMaxLen
if size%bitMaxLen > 0 {
bucketCount += 1
}
bucketCount *= 5
return &redisBloomFilter{
client: redisCli,
key: redisKey,
bucketCount: bucketCount,
bucketMaxLen: bitMaxLen,
hashFuncs: []func([]byte) uint64{hashFunc, hashFunc1, hashFunc2},
}
}
func (filter *redisBloomFilter) Put(b []byte) error {
pipe := filter.client.Pipeline()
ctx := context.Background()
var val uint64
var bucketIdx uint64
var key string
for i, f := range filter.hashFuncs {
val = f(b)
if i == 0 {
bucketIdx = val % filter.bucketCount //array index
key = fmt.Sprintf("%s_%d", filter.key, bucketIdx)
}
idx := val % filter.bucketMaxLen //bit index
_ = pipe.SetBit(ctx, key, int64(idx), 1)
}
_, err := pipe.Exec(ctx)
return err
}
func (filter *redisBloomFilter) MightContain(b []byte) (bool, error) {
pipe := filter.client.Pipeline()
ctx := context.Background()
var results []*redis.IntCmd
var val uint64
var bucketIdx uint64
var key string
for i, f := range filter.hashFuncs {
val = f(b)
if i == 0 {
bucketIdx = val % filter.bucketCount //array index
key = fmt.Sprintf("%s_%d", filter.key, bucketIdx)
}
idx := val % filter.bucketMaxLen //bit index
cmd := pipe.GetBit(ctx, key, int64(idx))
results = append(results, cmd)
}
if _, err := pipe.Exec(ctx); err != nil {
return false, err
}
for _, res := range results {
bit, err := res.Result()
if err != nil {
return false, err
}
if bit == 0 {
return false, nil
}
}
return true, nil
}