diff --git a/core/bloom/bloom.go b/core/bloom/bloom.go index 8646b017..897f4bc8 100644 --- a/core/bloom/bloom.go +++ b/core/bloom/bloom.go @@ -1,6 +1,7 @@ package bloom import ( + "context" "errors" "strconv" @@ -8,16 +9,15 @@ import ( "github.com/zeromicro/go-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 -) +// 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 +const maps = 14 var ( // ErrTooLargeOffset indicates the offset is too large in bitset. ErrTooLargeOffset = errors.New("too large offset") - setScript = redis.NewScript(` + + setScript = redis.NewScript(` for _, offset in ipairs(ARGV) do redis.call("setbit", KEYS[1], offset, 1) end @@ -40,8 +40,8 @@ type ( } bitSetProvider interface { - check([]uint) (bool, error) - set([]uint) error + check(ctx context.Context, offsets []uint) (bool, error) + set(ctx context.Context, offsets []uint) error } ) @@ -60,14 +60,24 @@ func New(store *redis.Redis, key string, bits uint) *Filter { // Add adds data into f. func (f *Filter) Add(data []byte) error { + return f.AddCtx(context.Background(), data) +} + +// AddCtx adds data into f with context. +func (f *Filter) AddCtx(ctx context.Context, data []byte) error { locations := f.getLocations(data) - return f.bitSet.set(locations) + return f.bitSet.set(ctx, locations) } // Exists checks if data is in f. func (f *Filter) Exists(data []byte) (bool, error) { + return f.ExistsCtx(context.Background(), data) +} + +// ExistsCtx checks if data is in f with context. +func (f *Filter) ExistsCtx(ctx context.Context, data []byte) (bool, error) { locations := f.getLocations(data) - isSet, err := f.bitSet.check(locations) + isSet, err := f.bitSet.check(ctx, locations) if err != nil { return false, err } @@ -113,13 +123,13 @@ func (r *redisBitSet) buildOffsetArgs(offsets []uint) ([]string, error) { return args, nil } -func (r *redisBitSet) check(offsets []uint) (bool, error) { +func (r *redisBitSet) check(ctx context.Context, offsets []uint) (bool, error) { args, err := r.buildOffsetArgs(offsets) if err != nil { return false, err } - resp, err := r.store.ScriptRun(testScript, []string{r.key}, args) + resp, err := r.store.ScriptRunCtx(ctx, testScript, []string{r.key}, args) if err == redis.Nil { return false, nil } else if err != nil { @@ -134,22 +144,24 @@ func (r *redisBitSet) check(offsets []uint) (bool, error) { return exists == 1, nil } +// del only use for testing. func (r *redisBitSet) del() error { _, err := r.store.Del(r.key) return err } +// expire only use for testing. func (r *redisBitSet) expire(seconds int) error { return r.store.Expire(r.key, seconds) } -func (r *redisBitSet) set(offsets []uint) error { +func (r *redisBitSet) set(ctx context.Context, offsets []uint) error { args, err := r.buildOffsetArgs(offsets) if err != nil { return err } - _, err = r.store.ScriptRun(setScript, []string{r.key}, args) + _, err = r.store.ScriptRunCtx(ctx, setScript, []string{r.key}, args) if err == redis.Nil { return nil } diff --git a/core/bloom/bloom_test.go b/core/bloom/bloom_test.go index 6362e390..0ecd622c 100644 --- a/core/bloom/bloom_test.go +++ b/core/bloom/bloom_test.go @@ -1,6 +1,7 @@ package bloom import ( + "context" "testing" "github.com/stretchr/testify/assert" @@ -10,20 +11,21 @@ import ( func TestRedisBitSet_New_Set_Test(t *testing.T) { store := redistest.CreateRedis(t) + ctx := context.Background() bitSet := newRedisBitSet(store, "test_key", 1024) - isSetBefore, err := bitSet.check([]uint{0}) + isSetBefore, err := bitSet.check(ctx, []uint{0}) if err != nil { t.Fatal(err) } if isSetBefore { t.Fatal("Bit should not be set") } - err = bitSet.set([]uint{512}) + err = bitSet.set(ctx, []uint{512}) if err != nil { t.Fatal(err) } - isSetAfter, err := bitSet.check([]uint{512}) + isSetAfter, err := bitSet.check(ctx, []uint{512}) if err != nil { t.Fatal(err) } @@ -66,33 +68,35 @@ func TestFilter_Exists(t *testing.T) { func TestRedisBitSet_check(t *testing.T) { store, clean := redistest.CreateRedisWithClean(t) + ctx := context.Background() rbs := newRedisBitSet(store, "test", 0) - assert.Error(t, rbs.set([]uint{0, 1, 2})) - _, err := rbs.check([]uint{0, 1, 2}) + assert.Error(t, rbs.set(ctx, []uint{0, 1, 2})) + _, err := rbs.check(ctx, []uint{0, 1, 2}) assert.Error(t, err) rbs = newRedisBitSet(store, "test", 64) - _, err = rbs.check([]uint{0, 1, 2}) + _, err = rbs.check(ctx, []uint{0, 1, 2}) assert.NoError(t, err) clean() rbs = newRedisBitSet(store, "test", 64) - _, err = rbs.check([]uint{0, 1, 2}) + _, err = rbs.check(ctx, []uint{0, 1, 2}) assert.Error(t, err) } func TestRedisBitSet_set(t *testing.T) { logx.Disable() store, clean := redistest.CreateRedisWithClean(t) + ctx := context.Background() rbs := newRedisBitSet(store, "test", 0) - assert.Error(t, rbs.set([]uint{0, 1, 2})) + assert.Error(t, rbs.set(ctx, []uint{0, 1, 2})) rbs = newRedisBitSet(store, "test", 64) - assert.NoError(t, rbs.set([]uint{0, 1, 2})) + assert.NoError(t, rbs.set(ctx, []uint{0, 1, 2})) clean() rbs = newRedisBitSet(store, "test", 64) - assert.Error(t, rbs.set([]uint{0, 1, 2})) + assert.Error(t, rbs.set(ctx, []uint{0, 1, 2})) }