initial import

core/bloom/bloom.go

@@ -0,0 +1,161 @@
+package bloom
+
+import (
+	"errors"
+	"strconv"
+
+	"zero/core/hash"
+	"zero/core/stores/redis"
+)
+
+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 = `
+local key = KEYS[1]
+for _, offset in ipairs(ARGV) do
+	redis.call("setbit", key, offset, 1)
+end
+`
+	testScript = `
+local key = KEYS[1]
+for _, offset in ipairs(ARGV) do
+	if tonumber(redis.call("getbit", key, offset)) == 0 then
+		return false
+	end
+end
+return true
+`
+)
+
+var ErrTooLargeOffset = errors.New("too large offset")
+
+type (
+	BitSetProvider interface {
+		check([]uint) (bool, error)
+		set([]uint) error
+	}
+
+	BloomFilter struct {
+		bits   uint
+		maps   uint
+		bitSet BitSetProvider
+	}
+)
+
+// New create a BloomFilter, store is the backed redis, key is the key for the bloom filter,
+// 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) *BloomFilter {
+	return &BloomFilter{
+		bits:   bits,
+		bitSet: newRedisBitSet(store, key, bits),
+	}
+}
+
+func (f *BloomFilter) Add(data []byte) error {
+	locations := f.getLocations(data)
+	err := f.bitSet.set(locations)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (f *BloomFilter) Exists(data []byte) (bool, error) {
+	locations := f.getLocations(data)
+	isSet, err := f.bitSet.check(locations)
+	if err != nil {
+		return false, err
+	}
+	if !isSet {
+		return false, nil
+	}
+
+	return true, nil
+}
+
+func (f *BloomFilter) getLocations(data []byte) []uint {
+	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
+	}
+
+	if exists, ok := resp.(int64); !ok {
+		return false, nil
+	} else {
+		return exists == 1, nil
+	}
+}
+
+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
+	} else {
+		return err
+	}
+}

core/bloom/bloom_test.go

@@ -0,0 +1,63 @@
+package bloom
+
+import (
+	"testing"
+
+	"zero/core/stores/redis"
+
+	"github.com/alicebob/miniredis"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRedisBitSet_New_Set_Test(t *testing.T) {
+	s, err := miniredis.Run()
+	if err != nil {
+		t.Error("Miniredis could not start")
+	}
+	defer s.Close()
+
+	store := redis.NewRedis(s.Addr(), redis.NodeType)
+	bitSet := newRedisBitSet(store, "test_key", 1024)
+	isSetBefore, err := bitSet.check([]uint{0})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if isSetBefore {
+		t.Fatal("Bit should not be set")
+	}
+	err = bitSet.set([]uint{512})
+	if err != nil {
+		t.Fatal(err)
+	}
+	isSetAfter, err := bitSet.check([]uint{512})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !isSetAfter {
+		t.Fatal("Bit should be set")
+	}
+	err = bitSet.expire(3600)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = bitSet.del()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestRedisBitSet_Add(t *testing.T) {
+	s, err := miniredis.Run()
+	if err != nil {
+		t.Error("Miniredis could not start")
+	}
+	defer s.Close()
+
+	store := redis.NewRedis(s.Addr(), redis.NodeType)
+	filter := New(store, "test_key", 64)
+	assert.Nil(t, filter.Add([]byte("hello")))
+	assert.Nil(t, filter.Add([]byte("world")))
+	ok, err := filter.Exists([]byte("hello"))
+	assert.Nil(t, err)
+	assert.True(t, ok)
+}