using TCC for MSetNX in cluster

README.md

@@ -142,8 +142,8 @@ I suggest focusing on the following directories:
     - set: a hash set based on map
     - sortedset: a sorted set implements based on skiplist
 - database: the core of storage engine
-    - server.go: a standalone redis server, 
-    - db.go: the basement function of database
+    - server.go: a standalone redis server, with multiple database 
+    - db.go: data structure and base functions of single database
     - exec.go: the gateway of database
     - router.go: the command table
     - keys.go: handlers for keys commands

README_CN.md

@@ -132,7 +132,7 @@ MSET (10 keys): 65487.89 requests per second
     - set： 基于hash表的集合
     - sortedset: 基于跳表实现的有序集合
 - database: 存储引擎核心
-    - db.go: 引擎的基础功能
+    - db.go: 单个 database 的数据结构和基本功能
     - router.go: 将命令路由给响应的处理函数
     - keys.go: del、ttl、expire 等通用命令实现
     - string.go: get、set 等字符串命令实现

cluster/cluster.go

@@ -76,7 +76,7 @@ func MakeCluster() *Cluster {
 }
 
 // CmdFunc represents the handler of a redis command
-type CmdFunc func(cluster *Cluster, c redis.Connection, cmdAndArgs [][]byte) redis.Reply
+type CmdFunc func(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply
 
 // Close stops current node of cluster
 func (cluster *Cluster) Close() {

cluster/mset.go

@@ -3,18 +3,21 @@ package cluster
 import (
 	"fmt"
 	"github.com/hdt3213/godis/interface/redis"
+	"github.com/hdt3213/godis/lib/utils"
 	"github.com/hdt3213/godis/redis/reply"
 	"strconv"
 )
 
+const keyExistsErr = "key exists"
+
 // MGet atomically get multi key-value from cluster, writeKeys can be distributed on any node
-func MGet(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
-	if len(args) < 2 {
+func MGet(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
+	if len(cmdLine) < 2 {
 		return reply.MakeErrReply("ERR wrong number of arguments for 'mget' command")
 	}
-	keys := make([]string, len(args)-1)
-	for i := 1; i < len(args); i++ {
-		keys[i-1] = string(args[i])
+	keys := make([]string, len(cmdLine)-1)
+	for i := 1; i < len(cmdLine); i++ {
+		keys[i-1] = string(cmdLine[i])
 	}
 
 	resultMap := make(map[string][]byte)
@@ -39,8 +42,8 @@ func MGet(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
 }
 
 // MSet atomically sets multi key-value in cluster, writeKeys can be distributed on any node
-func MSet(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
-	argCount := len(args) - 1
+func MSet(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
+	argCount := len(cmdLine) - 1
 	if argCount%2 != 0 || argCount < 1 {
 		return reply.MakeErrReply("ERR wrong number of arguments for 'mset' command")
 	}
@@ -49,14 +52,14 @@ func MSet(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
 	keys := make([]string, size)
 	valueMap := make(map[string]string)
 	for i := 0; i < size; i++ {
-		keys[i] = string(args[2*i+1])
-		valueMap[keys[i]] = string(args[2*i+2])
+		keys[i] = string(cmdLine[2*i+1])
+		valueMap[keys[i]] = string(cmdLine[2*i+2])
 	}
 
 	groupMap := cluster.groupBy(keys)
 	if len(groupMap) == 1 && allowFastTransaction { // do fast
 		for peer := range groupMap {
-			return cluster.relay(peer, c, args)
+			return cluster.relay(peer, c, cmdLine)
 		}
 	}
 
@@ -97,23 +100,92 @@ func MSet(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
 }
 
 // MSetNX sets multi key-value in database, only if none of the given writeKeys exist and all given writeKeys are on the same node
-func MSetNX(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
-	argCount := len(args) - 1
+func MSetNX(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
+	argCount := len(cmdLine) - 1
 	if argCount%2 != 0 || argCount < 1 {
-		return reply.MakeErrReply("ERR wrong number of arguments for 'mset' command")
+		return reply.MakeErrReply("ERR wrong number of arguments for 'msetnx' command")
 	}
-	var peer string
+
 	size := argCount / 2
+	keys := make([]string, size)
+	valueMap := make(map[string]string)
 	for i := 0; i < size; i++ {
-		key := string(args[2*i])
-		currentPeer := cluster.peerPicker.PickNode(key)
-		if peer == "" {
-			peer = currentPeer
+		keys[i] = string(cmdLine[2*i+1])
+		valueMap[keys[i]] = string(cmdLine[2*i+2])
+	}
+
+	groupMap := cluster.groupBy(keys)
+	if len(groupMap) == 1 && allowFastTransaction { // do fast
+		for peer := range groupMap {
+			return cluster.relay(peer, c, cmdLine)
+		}
+	}
+
+	// prepare procedure:
+	// 1. Normal tcc preparation (undo log and lock related keys)
+	// 2. Peer checks whether any key already exists, If so it will return keyExistsErr. Then coordinator will request rollback over all participated nodes
+	var errReply redis.Reply
+	txID := cluster.idGenerator.NextID()
+	txIDStr := strconv.FormatInt(txID, 10)
+	rollback := false
+	for peer, group := range groupMap {
+		peerArgs := []string{txIDStr, "MSETNX"}
+		for _, k := range group {
+			peerArgs = append(peerArgs, k, valueMap[k])
+		}
+		var resp redis.Reply
+		if peer == cluster.self {
+			resp = execPrepare(cluster, c, makeArgs("Prepare", peerArgs...))
 		} else {
-			if peer != currentPeer {
-				return reply.MakeErrReply("ERR msetnx must within one slot in cluster mode")
+			resp = cluster.relay(peer, c, makeArgs("Prepare", peerArgs...))
+		}
+		if reply.IsErrorReply(resp) {
+			re := resp.(reply.ErrorReply)
+			if re.Error() == keyExistsErr {
+				errReply = reply.MakeIntReply(0)
+			} else {
+				errReply = resp
+			}
+			rollback = true
+			break
 		}
 	}
+	if rollback {
+		// rollback
+		requestRollback(cluster, c, txID, groupMap)
+		return errReply
 	}
-	return cluster.relay(peer, c, args)
+	_, errReply = requestCommit(cluster, c, txID, groupMap)
+	rollback = errReply != nil
+	if !rollback {
+		return reply.MakeIntReply(1)
+	}
+	return errReply
+}
+
+func prepareMSetNx(cluster *Cluster, conn redis.Connection, cmdLine CmdLine) redis.Reply {
+	args := cmdLine[1:]
+	if len(args)%2 != 0 {
+		return reply.MakeSyntaxErrReply()
+	}
+	size := len(args) / 2
+	values := make([][]byte, size)
+	keys := make([]string, size)
+	for i := 0; i < size; i++ {
+		keys[i] = string(args[2*i])
+		values[i] = args[2*i+1]
+	}
+	re := cluster.db.ExecWithLock(conn, utils.ToCmdLine2("ExistIn", keys...))
+	if reply.IsErrorReply(re) {
+		return re
+	}
+	_, ok := re.(*reply.EmptyMultiBulkReply)
+	if !ok {
+		return reply.MakeErrReply(keyExistsErr)
+	}
+	return reply.MakeOkReply()
+}
+
+func init() {
+	prepareFuncMap["msetnx"] = prepareMSetNx
 }

cluster/mset_test.go

@@ -21,4 +21,8 @@ func TestMSetNx(t *testing.T) {
 	FlushAll(testCluster, conn, toArgs("FLUSHALL"))
 	ret := MSetNX(testCluster, conn, toArgs("MSETNX", "a", "a", "b", "b"))
 	asserts.AssertNotError(t, ret)
+	ret = MSetNX(testCluster, conn, toArgs("MSETNX", "a", "a", "c", "c"))
+	asserts.AssertNotError(t, ret)
+	ret = testCluster.Exec(conn, toArgs("MGET", "a", "b", "c"))
+	asserts.AssertMultiBulkReply(t, ret, []string{"a", "b", ""})
 }

cluster/tcc.go

@@ -8,10 +8,15 @@ import (
 	"github.com/hdt3213/godis/lib/timewheel"
 	"github.com/hdt3213/godis/redis/reply"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 )
 
+// prepareFuncMap executed after related key locked, and use additional logic to determine whether the transaction can be committed
+// For example, prepareMSetNX  will return error to prevent MSetNx transaction from committing if any related key already exists
+var prepareFuncMap = make(map[string]CmdFunc)
+
 // Transaction stores state and data for a try-commit-catch distributed transaction
 type Transaction struct {
 	id      string   // transaction id
@@ -72,7 +77,7 @@ func (tx *Transaction) unLockKeys() {
 	}
 }
 
-// t should contains Keys and Id field
+// t should contain Keys and ID field
 func (tx *Transaction) prepare() error {
 	tx.mu.Lock()
 	defer tx.mu.Unlock()
@@ -117,15 +122,20 @@ func (tx *Transaction) rollback() error {
 // cmdLine: Prepare id cmdName args...
 func execPrepare(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
 	if len(cmdLine) < 3 {
-		return reply.MakeErrReply("ERR wrong number of arguments for 'preparedel' command")
+		return reply.MakeErrReply("ERR wrong number of arguments for 'prepare' command")
 	}
 	txID := string(cmdLine[1])
+	cmdName := strings.ToLower(string(cmdLine[2]))
 	tx := NewTransaction(cluster, c, txID, cmdLine[2:])
 	cluster.transactions.Put(txID, tx)
 	err := tx.prepare()
 	if err != nil {
 		return reply.MakeErrReply(err.Error())
 	}
+	prepareFunc, ok := prepareFuncMap[cmdName]
+	if ok {
+		return prepareFunc(cluster, c, cmdLine[2:])
+	}
 	return &reply.OkReply{}
 }
 

database/keys.go

@@ -36,7 +36,7 @@ func undoDel(db *DB, args [][]byte) []CmdLine {
 	return rollbackGivenKeys(db, keys...)
 }
 
-// execExists checks if a is existed in db
+// execExists checks if given key is existed in db
 func execExists(db *DB, args [][]byte) redis.Reply {
 	result := int64(0)
 	for _, arg := range args {
@@ -49,6 +49,24 @@ func execExists(db *DB, args [][]byte) redis.Reply {
 	return reply.MakeIntReply(result)
 }
 
+// execExistIn returns existing key in given keys
+// example: ExistIn key1 key2 key3..., returns [key1, key2]
+// custom command for MSetNX tcc transaction
+func execExistIn(db *DB, args [][]byte) redis.Reply {
+	var result [][]byte
+	for _, arg := range args {
+		key := string(arg)
+		_, exists := db.GetEntity(key)
+		if exists {
+			result = append(result, []byte(key))
+		}
+	}
+	if len(result) == 0 {
+		return reply.MakeEmptyMultiBulkReply()
+	}
+	return reply.MakeMultiBulkReply(result)
+}
+
 // execFlushDB removes all data in current db
 func execFlushDB(db *DB, args [][]byte) redis.Reply {
 	db.Flush()
@@ -318,6 +336,7 @@ func init() {
 	RegisterCommand("PTTL", execPTTL, readFirstKey, nil, 2)
 	RegisterCommand("Persist", execPersist, writeFirstKey, undoExpire, 2)
 	RegisterCommand("Exists", execExists, readAllKeys, nil, -2)
+	RegisterCommand("ExistIn", execExistIn, readAllKeys, nil, -1)
 	RegisterCommand("Type", execType, readFirstKey, nil, 2)
 	RegisterCommand("Rename", execRename, prepareRename, undoRename, 3)
 	RegisterCommand("RenameNx", execRenameNx, prepareRename, undoRename, 3)

database/keys_test.go

@@ -22,6 +22,19 @@ func TestExists(t *testing.T) {
 	asserts.AssertIntReply(t, result, 0)
 }
 
+func TestExistIn(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	value := utils.RandString(10)
+	key2 := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine("set", key, value))
+	result := testDB.Exec(nil, utils.ToCmdLine("ExistIn", key, key2))
+	asserts.AssertMultiBulkReply(t, result, []string{key})
+	key3 := utils.RandString(10)
+	result = testDB.Exec(nil, utils.ToCmdLine("ExistIn", key2, key3))
+	asserts.AssertMultiBulkReplySize(t, result, 0)
+}
+
 func TestType(t *testing.T) {
 	testDB.Flush()
 	key := utils.RandString(10)

interface/database/db.go

@@ -10,7 +10,7 @@ type CmdLine = [][]byte
 
 // DB is the interface for redis style storage engine
 type DB interface {
-	Exec(client redis.Connection, args [][]byte) redis.Reply
+	Exec(client redis.Connection, cmdLine [][]byte) redis.Reply
 	AfterClientClose(c redis.Connection)
 	Close()
 }
@@ -18,7 +18,7 @@ type DB interface {
 // EmbedDB is the embedding storage engine exposing more methods for complex application
 type EmbedDB interface {
 	DB
-	ExecWithLock(conn redis.Connection, args [][]byte) redis.Reply
+	ExecWithLock(conn redis.Connection, cmdLine [][]byte) redis.Reply
 	ExecMulti(conn redis.Connection, watching map[string]uint32, cmdLines []CmdLine) redis.Reply
 	GetUndoLogs(dbIndex int, cmdLine [][]byte) []CmdLine
 	ForEach(dbIndex int, cb func(key string, data *DataEntity, expiration *time.Time) bool)