add scan command

database/keys.go

@@ -412,6 +412,146 @@ func execCopy(mdb *Server, conn redis.Connection, args [][]byte) redis.Reply {
 	return protocol.MakeIntReply(1)
 }
 
+//// execScan1 iteratively output all keys in the current db
+//func execScan1(db *DB, args [][]byte) redis.Reply {
+//	argsNum := len(args)
+//	if argsNum < 3 && argsNum > 5 {
+//		return protocol.MakeArgNumErrReply("scan")
+//	}
+//
+//	if argsNum == 1 {
+//		return scanWithArg(db, args, false)
+//	} else if argsNum == 3 {
+//		firstArg := strings.ToLower(string(args[1]))
+//		if firstArg == "match" {
+//			return scanWithArg(db, args, true)
+//		} else if firstArg == "count" {
+//			return scanWithArg(db, args, false)
+//		} else {
+//			return protocol.MakeSyntaxErrReply()
+//		}
+//	} else if argsNum == 5 {
+//		return scanWithArg(db, args, true)
+//	}
+//	return protocol.MakeNullBulkReply()
+//}
+
+// execScan iteratively output all keys in the current db
+func execScan(db *DB, args [][]byte) redis.Reply {
+	const (
+		noArgs = iota
+		count
+		match
+		countAndMatch = 5
+	)
+	commandArg := noArgs
+	argsNum := len(args)
+	if argsNum > 5 {
+		return protocol.MakeArgNumErrReply("scan")
+	}
+	if argsNum == 1 {
+		commandArg = noArgs
+	} else if argsNum > 2 {
+		for i := 1; i < argsNum; i++ {
+			arg := strings.ToLower(string(args[i]))
+			if arg == "count" {
+				commandArg = count
+			} else if arg == "match" {
+				commandArg = match
+			} else if i == 4 {
+				if string(args[1]) == "count" && string(args[3]) == "match" ||
+					(string(args[1]) == "match" && string(args[3]) == "count") {
+					commandArg = countAndMatch
+				} else {
+					return protocol.MakeSyntaxErrReply()
+				}
+				break
+			} else {
+				i++
+			}
+		}
+	}
+	switch commandArg {
+	case noArgs:
+		return execScanWithArg(db, args, noArgs)
+	case count:
+		return execScanWithArg(db, args, count)
+	case match:
+		return execScanWithArg(db, args, match)
+	case countAndMatch:
+		return execScanWithArg(db, args, countAndMatch)
+	}
+	return protocol.MakeNullBulkReply()
+}
+
+// execScanWithArg execute scan command based on cli args
+func execScanWithArg(db *DB, args [][]byte, argType int) redis.Reply {
+	result := make([]redis.Reply, 2)
+	var multiBulkReply [][]byte
+	var scanReturnKeys []string
+	var nextCursor int
+	cursor, err := strconv.Atoi(string(args[0]))
+	if err != nil {
+		return &protocol.SyntaxErrReply{}
+	}
+	// no args
+	if argType == 0 {
+		scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, 10, "*")
+	} else if argType >= 1 {
+		commandAndArgs := parseScanCommandArgs(args)
+		// only count
+		if argType == 1 {
+			count, err := strconv.Atoi(commandAndArgs["count"])
+			if err != nil || count < 0 {
+				return &protocol.SyntaxErrReply{}
+			}
+			scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, count, "*")
+		} else if argType == 2 {
+			// only match
+			// if there is only the match parameter, then count defaults to 10
+			scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, 10, commandAndArgs["match"])
+		} else if argType == 5 {
+			// count and match
+			count, err := strconv.Atoi(commandAndArgs["count"])
+			if err != nil || count < 0 {
+				return &protocol.SyntaxErrReply{}
+			}
+			scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, count, commandAndArgs["match"])
+		}
+	}
+	if nextCursor == -1 {
+		return protocol.MakeErrReply(scanReturnKeys[0])
+	}
+	nextCursorTobyte := strconv.FormatInt(int64(nextCursor), 10)
+	result[0] = protocol.MakeBulkReply([]byte(nextCursorTobyte))
+	for _, s := range scanReturnKeys {
+		if s != "" {
+			multiBulkReply = append(multiBulkReply, []byte(s))
+		}
+	}
+
+	result[1] = protocol.MakeMultiBulkReply(multiBulkReply)
+	println()
+	return protocol.MakeMultiRawReply(result)
+}
+
+// parseScanCommandArgs parse the parameters of the scan args
+func parseScanCommandArgs(args [][]byte) map[string]string {
+	// solving the order problem of count and match parameters
+	arg := make(map[string]string)
+	argNum := len(args)
+	if argNum == 3 {
+		firstCommand := strings.ToLower(string(args[1]))
+		arg[firstCommand] = string(args[2])
+	} else if argNum == 5 {
+		firstCommand := strings.ToLower(string(args[1]))
+		arg[firstCommand] = string(args[2])
+		secondCommand := strings.ToLower(string(args[3]))
+		arg[secondCommand] = string(args[4])
+	}
+	return arg
+}
+
 func init() {
 	registerCommand("Del", execDel, writeAllKeys, undoDel, -2, flagWrite).
 		attachCommandExtra([]string{redisFlagWrite}, 1, -1, 1)
@@ -443,4 +583,6 @@ func init() {
 		attachCommandExtra([]string{redisFlagWrite, redisFlagFast}, 1, 1, 1)
 	registerCommand("Keys", execKeys, noPrepare, nil, 2, flagReadOnly).
 		attachCommandExtra([]string{redisFlagReadonly, redisFlagSortForScript}, 1, 1, 1)
+	registerCommand("Scan", execScan, readAllKeys, nil, -2, flagReadOnly).
+		attachCommandExtra([]string{redisFlagReadonly, redisFlagSortForScript}, 1, 1, 1)
 }

database/keys_test.go

@@ -315,3 +315,39 @@ func TestCopy(t *testing.T) {
 	result = testMDB.Exec(conn, utils.ToCmdLine("ttl", destKey))
 	asserts.AssertIntReplyGreaterThan(t, result, 0)
 }
+
+func TestScan(t *testing.T) {
+	testDB.Flush()
+	for i := 0; i < 3; i++ {
+		key := string(rune(i))
+		value := key
+		testDB.Exec(nil, utils.ToCmdLine("set", "a:"+key, value))
+	}
+
+	result := testDB.Exec(nil, utils.ToCmdLine("scan", "0"))
+	expected := "*2\r\n$1\r\n0\r\n*3\r\n$3\r\na:\u0000\r\n$3\r\na:\u0001\r\n$3\r\na:\u0002\r\n"
+	if string(result.ToBytes()) != expected {
+		t.Error("test failed")
+	}
+	result = testDB.Exec(nil, utils.ToCmdLine("scan", "0", "match", "a*"))
+	if string(result.ToBytes()) != expected {
+		t.Error("test failed")
+	}
+	testDB.Exec(nil, utils.ToCmdLine("scan", "0", "match", "*"))
+	result = testDB.Exec(nil, utils.ToCmdLine("scan", "0", "count", "2"))
+	expected = "*2\r\n$5\r\n28194\r\n*2\r\n$3\r\na:\u0000\r\n$3\r\na:\u0001\r\n"
+	if string(result.ToBytes()) != expected {
+		t.Error("test failed")
+	}
+	result = testDB.Exec(nil, utils.ToCmdLine("scan", "0", "count", "2", "match", "a*"))
+	if string(result.ToBytes()) != expected {
+		t.Error("test failed")
+	}
+
+	result = testDB.Exec(nil, utils.ToCmdLine("scan", "0", "match", "b*"))
+	expected = "*2\r\n$1\r\n0\r\n*0\r\n"
+	if string(result.ToBytes()) != expected {
+		t.Error("test failed")
+	}
+
+}

datastruct/dict/concurrent.go

@@ -1,6 +1,7 @@
 package dict
 
 import (
+	"github.com/hdt3213/godis/lib/wildcard"
 	"math"
 	"math/rand"
 	"sort"
@@ -435,3 +436,60 @@ func (dict *ConcurrentDict) RWUnLocks(writeKeys []string, readKeys []string) {
 		}
 	}
 }
+
+// ScanKeys iteratively output all keys in the current db
+func (dict *ConcurrentDict) ScanKeys(cursor, count int, pattern string) ([]string, int) {
+	nextCursor := 0
+	errReturnCode := -1
+	size := dict.Len()
+	result := make([]string, count)
+	storeScanKeysTemp := make(map[string]struct{})
+	if pattern == "*" && count >= size {
+		return dict.Keys(), nextCursor
+	}
+
+	remainingShards := dict.table[cursor:]
+
+	matchKey, err := wildcard.CompilePattern(pattern)
+	if err != nil {
+		return []string{err.Error()}, errReturnCode
+	}
+
+	i := 0
+	for shardIndex, s := range remainingShards {
+		s.mutex.RLock()
+		f := func(m map[string]struct{}) bool {
+			defer s.mutex.RUnlock()
+			for key, _ := range s.m {
+				if key != "" {
+					if pattern != "*" {
+						if matchKey.IsMatch(key) {
+							m[key] = struct{}{}
+							i++
+						}
+					} else {
+						m[key] = struct{}{}
+						i++
+					}
+				}
+				if len(m) >= count {
+					return false
+				}
+			}
+			return true
+		}
+		nextCursor = shardIndex + cursor + 1
+		if !f(storeScanKeysTemp) {
+			break
+		}
+	}
+	j := 0
+	for k := range storeScanKeysTemp {
+		result[j] = k
+		j++
+	}
+	if nextCursor >= dict.shardCount {
+		nextCursor = 0
+	}
+	return result, nextCursor
+}

datastruct/dict/concurrent_test.go

@@ -465,7 +465,7 @@ func TestConcurrentRemoveWithLock(t *testing.T) {
 	}
 }
 
-//change t.Error remove->forEach
+// change t.Error remove->forEach
 func TestConcurrentForEach(t *testing.T) {
 	d := MakeConcurrent(0)
 	size := 100
@@ -524,3 +524,31 @@ func TestConcurrentDict_Keys(t *testing.T) {
 		t.Errorf("expect %d keys, actual: %d", size, len(d.Keys()))
 	}
 }
+
+func TestScanKeys(t *testing.T) {
+	d := MakeConcurrent(0)
+	count := 100
+	for i := 0; i < count; i++ {
+		key := "k" + strconv.Itoa(i)
+		d.Put(key, i)
+	}
+	cursor := 0
+	matchKey := "*"
+	c := 20
+	returnKeys, _ := d.ScanKeys(cursor, c, matchKey)
+	if len(returnKeys) != c {
+		t.Errorf("scan command count error: %d, should be %d ", len(returnKeys), c)
+	}
+	matchKey = "k*"
+	returnKeys, _ = d.ScanKeys(cursor, c, matchKey)
+	if len(returnKeys) != c {
+		t.Errorf("scan command count error: %d, should be %d ", len(returnKeys), c)
+	}
+	matchKey = "s*"
+	returnKeys, _ = d.ScanKeys(cursor, c, matchKey)
+	for _, key := range returnKeys {
+		if key != "" {
+			t.Errorf("returnKeys should be empty")
+		}
+	}
+}

datastruct/dict/dict.go

@@ -16,4 +16,5 @@ type Dict interface {
 	RandomKeys(limit int) []string
 	RandomDistinctKeys(limit int) []string
 	Clear()
+	ScanKeys(cursor, count int, matchKey string) ([]string, int)
 }

datastruct/dict/simple.go

@@ -120,3 +120,25 @@ func (dict *SimpleDict) RandomDistinctKeys(limit int) []string {
 func (dict *SimpleDict) Clear() {
 	*dict = *MakeSimple()
 }
+
+// ScanKeys randomly returns keys of the given number, may contain duplicated key
+func (dict *SimpleDict) ScanKeys(cursor, count int, matchKey string) ([]string, int) {
+	nextCursor := 0
+	size := dict.Len()
+	if count >= size {
+		return dict.Keys(), nextCursor
+	}
+
+	if count == 0 {
+		count = 10
+	}
+
+	result := make([]string, count)
+	for i := cursor; i < count; i++ {
+		for k := range dict.m {
+			result[i] = k
+		}
+	}
+
+	return result, nextCursor
+}

datastruct/dict/simple_test.go

@@ -2,6 +2,7 @@ package dict
 
 import (
 	"github.com/hdt3213/godis/lib/utils"
+	"github.com/hdt3213/godis/lib/wildcard"
 	"sort"
 	"testing"
 )
@@ -53,3 +54,31 @@ func TestSimpleDict_PutIfExists(t *testing.T) {
 		return
 	}
 }
+
+func TestSimpleDict_ScanKeys(t *testing.T) {
+	count := 1
+	d := MakeSimple()
+	for i := 0; i < 5; i++ {
+		value := utils.RandString(5)
+		key := "k" + value
+		ret := d.PutIfExists(key, value)
+		if ret != 0 {
+			t.Error("expect 0")
+			return
+		}
+	}
+	result, _ := d.ScanKeys(0, count, "k*")
+
+	pattern, err := wildcard.CompilePattern("k*")
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	for _, s := range result {
+		if !pattern.IsMatch(s) {
+			t.Error("Scan command execution error")
+		}
+	}
+
+}