Merge pull request #12 from absolutelightning/keys-with-cursor

Keys with cursor and limit - Faster Regex Scans

README.md

@@ -86,8 +86,8 @@ This graph shows the performance comparison between Treds - ZRangeScoreKeys and
 * `DBSIZE` - Get number of keys in the db
 * `SCANKEYS cursor prefix count` - Returns the count number of keys matching prefix starting from an index in lex order only present in Key/Value Store. Last element is the next cursor
 * `SCANKVS cursor prefix count` - Returns the count number of keys/value pair in which keys match prefix starting from an index in lex order only present in Key/Value Store. Last element is the next cursor
-* `KEYS regex` - Returns all keys matching a regex in lex order - (Not suitable to production use cases with huge number of keys)
-* `KVS regex` - Returns all keys/values in which keys match a regex in lex order - (Not suitable to production use cases with huge number of keys)
+* `KEYS cursor regex count` - Returns count number of keys matching a regex in lex order starting with cursor. Count is optional. Last element is the next cursor
+* `KVS cursor regex count` - Returns count number of keys/values in which keys match a regex in lex order starting with cursor. Count is optional. Last element is the next cursor
 * `ZADD key score member_key member_value [score member_key member_value ....]` - Add member_key with member value with score to a sorted map in key
 * `ZREM key member [member ...]` - Removes a member from sorted map in key
 * `ZCARD key` - Returns the count of key/value pairs in sorted map in key

client/main.go

@@ -95,8 +95,8 @@ func completer(d prompt.Document) []prompt.Suggest {
 		{Text: "HLEN", Description: "HLEN key - Returns the size of hash at the key"},
 		{Text: "HSET", Description: "HSET key field value [field value ...] - Sets field value pairs in the hash with key"},
 		{Text: "HVALS", Description: "HVALS key - Returns all values present in the hash at key"},
-		{Text: "KEYS", Description: "KEYS regex - Returns all keys matching a regex in lex order"},
-		{Text: "KVS", Description: "KVS regex - Returns all keys/values in which keys match a regex in lex order"},
+		{Text: "KEYS", Description: "Returns count number of keys matching a regex in lex order starting with cursor. Count is optional. Last element is the next cursor"},
+		{Text: "KVS", Description: "Returns count number of keys/values in which keys match a regex in lex order starting with cursor. Count is optional. Last element is the next cursor"},
 		{Text: "LINDEX", Description: "LINDEX key index - Returns the element at index of list with key"},
 		{Text: "LLEN", Description: "LLEN key - Returns the length of list with key"},
 		{Text: "LNGPREFIX", Description: "LNGPREFIX string - Returns the key value pair in which key is the longest prefix of given string"},

commands/keys.go

@@ -2,7 +2,9 @@ package commands
 
 import (
 	"fmt"
+	"math"
 	"regexp"
+	"strconv"
 
 	"treds/store"
 )
@@ -19,23 +21,34 @@ func RegisterKeysCommand(r CommandRegistry) {
 
 func validateKeys() ValidationHook {
 	return func(args []string) error {
-		if len(args) == 1 {
-			_, err := regexp.Compile(args[0])
+		if len(args) < 2 {
+			return fmt.Errorf("expected minimum 2 argument, got %d", len(args))
+		}
+		if len(args) == 3 {
+			_, err := strconv.Atoi(args[2])
 			if err != nil {
 				return err
 			}
 		}
+		_, err := regexp.Compile(args[0])
+		if err != nil {
+			return err
+		}
 		return nil
 	}
 }
 
 func executeKeys() ExecutionHook {
 	return func(args []string, store store.Store) (string, error) {
 		regex := ""
-		if len(args) == 1 {
-			regex = args[0]
+		count := math.MaxInt64
+		if len(args) >= 2 {
+			regex = args[1]
+		}
+		if len(args) == 3 {
+			count, _ = strconv.Atoi(args[2])
 		}
-		v, err := store.Keys(regex)
+		v, err := store.Keys(args[0], regex, count)
 		if err != nil {
 			return "", err
 		}

commands/keys_test.go

@@ -14,20 +14,6 @@ func TestRegisterKeysCommand(t *testing.T) {
 	}
 }
 
-// TestValidateKeys tests the validateKeys function.
-func TestValidateKeys(t *testing.T) {
-	validationHook := validateKeys()
-
-	// The validateKeys function always returns nil, so we simply check that it does.
-	if err := validationHook([]string{}); err != nil {
-		t.Errorf("expected no error, got: %v", err)
-	}
-
-	if err := validationHook([]string{"someArg"}); err != nil {
-		t.Errorf("expected no error, got: %v", err)
-	}
-}
-
 // TestExecuteKeys tests the executeKeys function.
 func TestExecuteKeys(t *testing.T) {
 	tests := []struct {
@@ -39,21 +25,21 @@ func TestExecuteKeys(t *testing.T) {
 	}{
 		{
 			name:        "retrieve all keys",
-			args:        []string{},
+			args:        []string{"0"},
 			store:       &MockStore{data: map[string]string{"key1": "value1", "key2": "value2"}},
 			expectErr:   false,
 			expectedMsg: "key1\nkey2\n",
 		},
 		{
 			name:        "retrieve keys with matching prefix",
-			args:        []string{"key"},
+			args:        []string{"0", "^key"},
 			store:       &MockStore{data: map[string]string{"key1": "value1", "key2": "value2", "other": "value3"}},
 			expectErr:   false,
 			expectedMsg: "key1\nkey2\n",
 		},
 		{
 			name:        "no matching keys",
-			args:        []string{"nomatch"},
+			args:        []string{"0", "nomatch"},
 			store:       &MockStore{data: map[string]string{"key1": "value1", "key2": "value2"}},
 			expectErr:   false,
 			expectedMsg: "",

commands/kvs.go

@@ -2,7 +2,9 @@ package commands
 
 import (
 	"fmt"
+	"math"
 	"regexp"
+	"strconv"
 
 	"treds/store"
 )
@@ -19,23 +21,34 @@ func RegisterKVSCommand(r CommandRegistry) {
 
 func validateKVS() ValidationHook {
 	return func(args []string) error {
-		if len(args) == 1 {
-			_, err := regexp.Compile(args[0])
+		if len(args) < 2 {
+			return fmt.Errorf("expected minimum 2 argument, got %d", len(args))
+		}
+		if len(args) == 3 {
+			_, err := strconv.Atoi(args[2])
 			if err != nil {
 				return err
 			}
 		}
+		_, err := regexp.Compile(args[0])
+		if err != nil {
+			return err
+		}
 		return nil
 	}
 }
 
 func executeKVS() ExecutionHook {
 	return func(args []string, store store.Store) (string, error) {
 		regex := ""
-		if len(args) == 1 {
-			regex = args[0]
+		count := math.MaxInt64
+		if len(args) >= 2 {
+			regex = args[1]
+		}
+		if len(args) == 3 {
+			count, _ = strconv.Atoi(args[2])
 		}
-		v, err := store.KVS(regex)
+		v, err := store.KVS(args[0], regex, count)
 		if err != nil {
 			return "", err
 		}

commands/kvs_test.go

@@ -14,20 +14,6 @@ func TestRegisterKVSCommand(t *testing.T) {
 	}
 }
 
-// TestValidateKVS tests the validateKVS function.
-func TestValidateKVS(t *testing.T) {
-	validationHook := validateKVS()
-
-	// The validateKVS function always returns nil, so we simply check that it does.
-	if err := validationHook([]string{}); err != nil {
-		t.Errorf("expected no error, got: %v", err)
-	}
-
-	if err := validationHook([]string{"someArg"}); err != nil {
-		t.Errorf("expected no error, got: %v", err)
-	}
-}
-
 // TestExecuteKVS tests the executeKVS function.
 func TestExecuteKVS(t *testing.T) {
 	tests := []struct {
@@ -39,21 +25,21 @@ func TestExecuteKVS(t *testing.T) {
 	}{
 		{
 			name:        "retrieve all key-value pairs",
-			args:        []string{},
+			args:        []string{"0"},
 			store:       &MockStore{data: map[string]string{"key1": "value1", "key2": "value2"}},
 			expectErr:   false,
 			expectedMsg: "key1\nvalue1\nkey2\nvalue2\n",
 		},
 		{
 			name:        "retrieve key-value pairs with matching prefix",
-			args:        []string{"key"},
+			args:        []string{"0", "^key"},
 			store:       &MockStore{data: map[string]string{"key1": "value1", "key2": "value2", "other": "value3"}},
 			expectErr:   false,
 			expectedMsg: "key1\nvalue1\nkey2\nvalue2\n",
 		},
 		{
 			name:        "no matching key-value pairs",
-			args:        []string{"nomatch"},
+			args:        []string{"0", "nomatch"},
 			store:       &MockStore{data: map[string]string{"key1": "value1", "key2": "value2"}},
 			expectErr:   false,
 			expectedMsg: "",

commands/mock_store.go

@@ -96,7 +96,7 @@ func (m *MockStore) DeletePrefix(prefix string) (int, error) {
 	return 0, nil
 }
 
-func (m *MockStore) Keys(regex string) (string, error) {
+func (m *MockStore) Keys(cursor, regex string, count int) (string, error) {
 	res := ""
 	keys := make([]string, 0)
 	for key, _ := range m.data {
@@ -112,7 +112,7 @@ func (m *MockStore) Keys(regex string) (string, error) {
 	return res, nil
 }
 
-func (m *MockStore) KVS(regex string) (string, error) {
+func (m *MockStore) KVS(cursor, regex string, count int) (string, error) {
 	res := ""
 	keys := make([]string, 0)
 	for key, _ := range m.data {

store/store.go

@@ -11,8 +11,8 @@ type Store interface {
 	PrefixScan(string, string, string) (string, error)
 	PrefixScanKeys(string, string, string) (string, error)
 	DeletePrefix(string) (int, error)
-	Keys(string) (string, error)
-	KVS(string) (string, error)
+	Keys(string, string, int) (string, error)
+	KVS(string, string, int) (string, error)
 	Size() (string, error)
 	ZAdd([]string) error
 	ZRem([]string) error

store/tred_store.go

@@ -323,12 +323,21 @@ func (rs *TredsStore) DeletePrefix(prefix string) (int, error) {
 	return numDel, nil
 }
 
-func (rs *TredsStore) Keys(regex string) (string, error) {
+func (rs *TredsStore) Keys(cursor, regex string, count int) (string, error) {
+	startHash, err := strconv.Atoi(cursor)
+	if err != nil {
+		return "", err
+	}
 	iterator := rs.tree.Root().Iterator()
 	rx := regexp.MustCompile(regex)
 	iterator.PatternMatch(rx)
 
 	var result strings.Builder
+	seenHash := false
+	if cursor == "0" {
+		seenHash = true
+	}
+	nextCursor := uint32(0)
 
 	for {
 		key, _, found := iterator.Next()
@@ -338,18 +347,48 @@ func (rs *TredsStore) Keys(regex string) (string, error) {
 		if rs.hasExpired(string(key)) {
 			continue
 		}
-		result.WriteString(fmt.Sprintf("%v\n", string(key)))
+		hashKey, herr := hash(string(key))
+		if herr != nil {
+			return "", herr
+		}
+		if !seenHash && hashKey == uint32(startHash) {
+			seenHash = true
+			continue
+		}
+		if seenHash && count > 0 {
+			result.WriteString(fmt.Sprintf("%v\n", string(key)))
+			nextCursor, herr = hash(string(key))
+			if herr != nil {
+				return "", herr
+			}
+			count--
+		}
+		if count == 0 {
+			break
+		}
 	}
-
+	if count != 0 {
+		nextCursor = uint32(0)
+	}
+	result.WriteString(strconv.Itoa(int(nextCursor)) + "\n")
 	return result.String(), nil
 }
 
-func (rs *TredsStore) KVS(regex string) (string, error) {
+func (rs *TredsStore) KVS(cursor, regex string, count int) (string, error) {
+	startHash, err := strconv.Atoi(cursor)
+	if err != nil {
+		return "", err
+	}
 	iterator := rs.tree.Root().Iterator()
 	rx := regexp.MustCompile(regex)
 	iterator.PatternMatch(rx)
 
 	var result strings.Builder
+	seenHash := false
+	if cursor == "0" {
+		seenHash = true
+	}
+	nextCursor := uint32(0)
 
 	for {
 		key, value, found := iterator.Next()
@@ -359,9 +398,30 @@ func (rs *TredsStore) KVS(regex string) (string, error) {
 		if rs.hasExpired(string(key)) {
 			continue
 		}
-		result.WriteString(fmt.Sprintf("%v\n%v\n", string(key), value.(string)))
+		hashKey, herr := hash(string(key))
+		if herr != nil {
+			return "", herr
+		}
+		if !seenHash && hashKey == uint32(startHash) {
+			seenHash = true
+			continue
+		}
+		if seenHash && count > 0 {
+			result.WriteString(fmt.Sprintf("%v\n%v\n", string(key), value.(string)))
+			nextCursor, herr = hash(string(key))
+			if herr != nil {
+				return "", herr
+			}
+			count--
+		}
+		if count == 0 {
+			break
+		}
 	}
-
+	if count != 0 {
+		nextCursor = uint32(0)
+	}
+	result.WriteString(strconv.Itoa(int(nextCursor)) + "\n")
 	return result.String(), nil
 }
 

store/tred_store_test.go

@@ -162,11 +162,11 @@ func TestTredsStore_Keys(t *testing.T) {
 	store.Set("otherkey", "value3")
 
 	// Test retrieving keys matching a regex
-	result, err := store.Keys("^key.*")
+	result, err := store.Keys("0", "^key.*", 100000000)
 	if err != nil {
 		t.Fatalf("expected no error, got %v", err)
 	}
-	expected := "key1\nkey2\n"
+	expected := "key1\nkey2\n0\n"
 	if result != expected {
 		t.Fatalf("expected %s, got %s", expected, result)
 	}
@@ -180,11 +180,11 @@ func TestTredsStore_KVS(t *testing.T) {
 	store.Set("otherkey", "value3")
 
 	// Test retrieving key-value pairs matching a regex
-	result, err := store.KVS("^key.*")
+	result, err := store.KVS("0", "^key.*", 10000000000000)
 	if err != nil {
 		t.Fatalf("expected no error, got %v", err)
 	}
-	expected := "key1\nvalue1\nkey2\nvalue2\n"
+	expected := "key1\nvalue1\nkey2\nvalue2\n0\n"
 	if result != expected {
 		t.Fatalf("expected %s, got %s", expected, result)
 	}