Merge pull request #101 from ceshihao/rump_filter_key

support filter by key in restore and rump mode

conf/redis-shake.conf

@@ -111,7 +111,7 @@ filter.db =
 # filter key with prefix string. multiple keys are separated by ';'.
 # e.g., a;b;c
 # default is all.
-# used in `restore` and `sync`.
+# used in `restore`, `sync` and `rump`.
 # 支持过滤key，只让指定的key通过，分号分隔
 filter.key =
 # filter given slot, multiple slots are separated by ';'.

src/redis-shake/common/utils.go

@@ -1003,3 +1003,13 @@ var defaultDialFunction = func(addr string) (redigo.Conn, error) {
 	}
 	return c, nil
 }
+
+// HasAtLeastOnePrefix checks whether the key begins with at least one of prefixes.
+func HasAtLeastOnePrefix(key string, prefixes []string) bool {
+	for _, prefix := range prefixes {
+		if strings.HasPrefix(key, prefix) {
+			return true
+		}
+	}
+	return false
+}

src/redis-shake/common/utils_test.go

@@ -1,9 +1,9 @@
 package utils
 
 import (
-	"testing"
 	"fmt"
 	"sort"
+	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -14,8 +14,8 @@ func TestGetAllClusterNode(t *testing.T) {
 		fmt.Printf("TestGetAllClusterNode case %d.\n", nr)
 		nr++
 
-		client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false)
-		ret, err := GetAllClusterNode(client, "master")
+		client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false, false)
+		ret, err := GetAllClusterNode(client, "master", "")
 		sort.Strings(ret)
 		assert.Equal(t, nil, err, "should be equal")
 		assert.Equal(t, 3, len(ret), "should be equal")
@@ -28,8 +28,8 @@ func TestGetAllClusterNode(t *testing.T) {
 		fmt.Printf("TestGetAllClusterNode case %d.\n", nr)
 		nr++
 
-		client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false)
-		ret, err := GetAllClusterNode(client, "slave")
+		client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false, false)
+		ret, err := GetAllClusterNode(client, "slave", "")
 		sort.Strings(ret)
 		assert.Equal(t, nil, err, "should be equal")
 		assert.Equal(t, 3, len(ret), "should be equal")
@@ -42,8 +42,8 @@ func TestGetAllClusterNode(t *testing.T) {
 		fmt.Printf("TestGetAllClusterNode case %d.\n", nr)
 		nr++
 
-		client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false)
-		ret, err := GetAllClusterNode(client, "all")
+		client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false, false)
+		ret, err := GetAllClusterNode(client, "all", "")
 		sort.Strings(ret)
 		assert.Equal(t, nil, err, "should be equal")
 		assert.Equal(t, 6, len(ret), "should be equal")
@@ -55,3 +55,35 @@ func TestGetAllClusterNode(t *testing.T) {
 		assert.Equal(t, "10.1.1.1:21336", ret[5], "should be equal")
 	}
 }
+
+func TestHasAtLeastOnePrefix(t *testing.T) {
+	cases := []struct {
+		key          string
+		prefixes     []string
+		expectResult bool
+	}{
+		{
+			// no prefix provided
+			"a",
+			[]string{},
+			false,
+		},
+		{
+			// has prefix
+			"abc",
+			[]string{"ab"},
+			true,
+		},
+		{
+			// does NOT have prefix
+			"abc",
+			[]string{"edf", "wab"},
+			false,
+		},
+	}
+
+	for _, c := range cases {
+		result := HasAtLeastOnePrefix(c.key, c.prefixes)
+		assert.Equal(t, c.expectResult, result)
+	}
+}

src/redis-shake/restore.go

@@ -137,7 +137,6 @@ func (dr *dbRestorer) restore() {
 	}
 }
 
-
 func (dr *dbRestorer) restoreRDBFile(reader *bufio.Reader, target []string, auth_type, passwd string, nsize int64,
 	tlsEnable bool) {
 	pipe := utils.NewRDBLoader(reader, &dr.rbytes, base.RDBPipeSize)
@@ -168,6 +167,9 @@ func (dr *dbRestorer) restoreRDBFile(reader *bufio.Reader, target []string, auth
 								utils.SelectDB(c, lastdb)
 							}
 						}
+						if len(conf.Options.FilterKey) != 0 && !utils.HasAtLeastOnePrefix(string(e.Key), conf.Options.FilterKey) {
+							continue
+						}
 						utils.RestoreRdbEntry(c, e)
 					}
 				}

src/redis-shake/rump.go

@@ -1,18 +1,18 @@
 package run
 
 import (
-	"strconv"
-	"sync"
 	"fmt"
-	"reflect"
 	"math"
+	"reflect"
+	"strconv"
+	"sync"
 
-	"pkg/libs/log"
 	"pkg/libs/atomic2"
+	"pkg/libs/log"
 	"redis-shake/common"
 	"redis-shake/configure"
-	"redis-shake/scanner"
 	"redis-shake/metric"
+	"redis-shake/scanner"
 
 	"github.com/garyburd/redigo/redis"
 )
@@ -320,6 +320,9 @@ func (dre *dbRumperExecutor) writer() {
 	bucket := utils.StartQoS(conf.Options.Qps)
 	preDb := 0
 	for ele := range dre.keyChan {
+		if len(conf.Options.FilterKey) != 0 && !utils.HasAtLeastOnePrefix(ele.key, conf.Options.FilterKey) {
+			continue
+		}
 		// QoS, limit the qps
 		<-bucket
 

src/redis-shake/sync.go

@@ -426,11 +426,8 @@ func (ds *dbSyncer) syncRDBFile(reader *bufio.Reader, target []string, auth_type
 						}
 
 						if len(conf.Options.FilterKey) != 0 {
-							for i := 0; i < len(conf.Options.FilterKey); i++ {
-								if strings.HasPrefix(string(e.Key), conf.Options.FilterKey[i]) {
+							if utils.HasAtLeastOnePrefix(string(e.Key), conf.Options.FilterKey) {
 								utils.RestoreRdbEntry(c, e)
-									break
-								}
 							}
 						} else if len(conf.Options.FilterSlot) > 0 {
 							for _, slot := range conf.Options.FilterSlot {