go-zero源码阅读之布隆过滤器实现代码

type (
	// A Filter is a bloom filter.
    // 结构体
	Filter struct {
		bits   uint
		bitSet bitSetProvider
	}
 
    // 位数组接口定义 
	bitSetProvider interface {
		check([]uint) (bool, error)
		set([]uint) error
	}
)

func New(store *redis.Redis, key string, bits uint) *Filter {
	return &Filter{
		bits:   bits,
		bitSet: newRedisBitSet(store, key, bits),
	}
}

func newRedisBitSet(store *redis.Redis, key string, bits uint) *redisBitSet {
	return &redisBitSet{
		store: store,
		key:   key,
		bits:  bits,
	}
}

func (f *Filter) Add(data []byte) error {
    // 获取数据多次hash后的各key
	locations := f.getLocations(data)
    // 插入数据
	return f.bitSet.set(locations)
}

func (r *redisBitSet) set(offsets []uint) error {
    // 将[]uint转为[]string
	args, err := r.buildOffsetArgs(offsets)
	if err != nil {
		return err
	}
    // 执行lua脚本
	_, err = r.store.Eval(setScript, []string{r.key}, args)
	if err == redis.Nil {
		return nil
	}
 
	return err
}

setScript = `
for _, offset in ipairs(ARGV) do
	redis.call("setbit", KEYS[1], offset, 1)
end
`

func (f *Filter) Exists(data []byte) (bool, error) {
    // 同数据set一致，获取数据多次hash后，偏移量切片
	locations := f.getLocations(data)
    // 调用check方法进行检测
	isSet, err := f.bitSet.check(locations)
	if err != nil {
		return false, err
	}
 
	return isSet, nil
}

func (r *redisBitSet) check(offsets []uint) (bool, error) {
    // []uint转为[]string,和set调用的一致
	args, err := r.buildOffsetArgs(offsets)
	if err != nil {
		return false, err
	}
 
    //执行lua脚本，检测各偏移量数据是否都存在 
	resp, err := r.store.Eval(testScript, []string{r.key}, args)
    // 根据返回值判断数据是否存在
   // key不存在特殊处理
	if err == redis.Nil {
		return false, nil
	} else if err != nil {
		return false, err
	}
 
	exists, ok := resp.(int64)
	if !ok {
		return false, nil
	}
   
	return exists == 1, nil
}

testScript = `
for _, offset in ipairs(ARGV) do
	if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then
		return false
	end
end
return true
`