add scan command

database/keys.go

@@ -412,6 +412,91 @@ func execCopy(mdb *Server, conn redis.Connection, args [][]byte) redis.Reply {
 	return protocol.MakeIntReply(1)
 }
 
+// execScan iteratively output all keys in the current db
+func execScan(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()
+}
+
+func scanWithArg(db *DB, args [][]byte, match bool) redis.Reply {
+	result := make([]redis.Reply, 2)
+	var m [][]byte
+	var scanReturnKeys []string
+	var nextCursor int
+	cursor, err := strconv.Atoi(string(args[0]))
+	if err != nil {
+		return &protocol.SyntaxErrReply{}
+	}
+
+	if len(args) == 1 {
+		scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, 10, "*")
+	} else if !match {
+		count, err := strconv.Atoi(string(args[2]))
+		if err != nil || count < 0 {
+			return &protocol.SyntaxErrReply{}
+		}
+		scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, count, "*")
+	} else {
+		commandAndArgs := parseScanCommandArgs(args)
+		numOfcommandAndArgs := len(commandAndArgs)
+		if numOfcommandAndArgs == 1 {
+			//if there is only the match parameter, then count defaults to 10
+			scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, 10, commandAndArgs["match"])
+		} else if numOfcommandAndArgs == 2 {
+			count, err := strconv.Atoi(commandAndArgs["count"])
+			if err != nil || count <= 0 {
+				return &protocol.SyntaxErrReply{}
+			}
+			scanReturnKeys, nextCursor = db.data.ScanKeys(cursor, count, commandAndArgs["match"])
+		} else {
+			return &protocol.SyntaxErrReply{}
+		}
+	}
+	nextCursorTobyte := strconv.FormatInt(int64(nextCursor), 10)
+	result[0] = protocol.MakeBulkReply([]byte(nextCursorTobyte))
+	for _, s := range scanReturnKeys {
+		if s != "" {
+			m = append(m, []byte(s))
+		}
+	}
+	result[1] = protocol.MakeMultiBulkReply(m)
+	return protocol.MakeMultiRawReply(result)
+}
+
+// parseScanCommandArgs parse the parameters of the scan args
+func parseScanCommandArgs(args [][]byte) map[string]string {
+	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 +528,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)
 }

datastruct/dict/concurrent.go

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