feat(bloom): bloom support Ctx API (#3089)

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
 	}

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}))
 }