You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/core/bloom/bloom.go

157 lines
3.3 KiB
Go

4 years ago
package bloom
import (
"errors"
"strconv"
"github.com/zeromicro/go-zero/core/hash"
"github.com/zeromicro/go-zero/core/stores/redis"
4 years ago
)
const (
// for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
// maps as k in the error rate table
maps = 14
setScript = `
for _, offset in ipairs(ARGV) do
redis.call("setbit", KEYS[1], offset, 1)
4 years ago
end
`
testScript = `
for _, offset in ipairs(ARGV) do
if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
4 years ago
return false
end
end
return true
`
)
// ErrTooLargeOffset indicates the offset is too large in bitset.
4 years ago
var ErrTooLargeOffset = errors.New("too large offset")
type (
// A Filter is a bloom filter.
Filter struct {
bits uint
bitSet bitSetProvider
4 years ago
}
bitSetProvider interface {
check([]uint) (bool, error)
set([]uint) error
4 years ago
}
)
// New create a Filter, store is the backed redis, key is the key for the bloom filter,
4 years ago
// bits is how many bits will be used, maps is how many hashes for each addition.
// best practices:
// elements - means how many actual elements
// when maps = 14, formula: 0.7*(bits/maps), bits = 20*elements, the error rate is 0.000067 < 1e-4
// for detailed error rate table, see http://pages.cs.wisc.edu/~cao/papers/summary-cache/node8.html
func New(store *redis.Redis, key string, bits uint) *Filter {
return &Filter{
4 years ago
bits: bits,
bitSet: newRedisBitSet(store, key, bits),
}
}
// Add adds data into f.
func (f *Filter) Add(data []byte) error {
4 years ago
locations := f.getLocations(data)
return f.bitSet.set(locations)
4 years ago
}
// Exists checks if data is in f.
func (f *Filter) Exists(data []byte) (bool, error) {
4 years ago
locations := f.getLocations(data)
isSet, err := f.bitSet.check(locations)
if err != nil {
return false, err
}
return isSet, nil
4 years ago
}
func (f *Filter) getLocations(data []byte) []uint {
4 years ago
locations := make([]uint, maps)
for i := uint(0); i < maps; i++ {
hashValue := hash.Hash(append(data, byte(i)))
locations[i] = uint(hashValue % uint64(f.bits))
}
return locations
}
type redisBitSet struct {
store *redis.Redis
key string
bits uint
}
func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
return &redisBitSet{
store: store,
key: key,
bits: bits,
}
}
func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) {
var args []string
for _, offset := range offsets {
if offset >= r.bits {
return nil, ErrTooLargeOffset
}
args = append(args, strconv.FormatUint(uint64(offset), 10))
}
return args, nil
}
func (r *redisBitSet) check(offsets []uint) (bool, error) {
args, err := r.buildOffsetArgs(offsets)
if err != nil {
return false, err
}
resp, err := r.store.Eval(testScript, []string{r.key}, args)
if err == redis.Nil {
return false, nil
} else if err != nil {
return false, err
}
exists, ok := resp.(int64)
if !ok {
4 years ago
return false, nil
}
return exists == 1, nil
4 years ago
}
func (r *redisBitSet) del() error {
_, err := r.store.Del(r.key)
return err
}
func (r *redisBitSet) expire(seconds int) error {
return r.store.Expire(r.key, seconds)
}
func (r *redisBitSet) set(offsets []uint) error {
args, err := r.buildOffsetArgs(offsets)
if err != nil {
return err
}
_, err = r.store.Eval(setScript, []string{r.key}, args)
if err == redis.Nil {
return nil
}
return err
4 years ago
}