using TCC for MSetNX in cluster

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
-		} else {
-			if peer != currentPeer {
-				return reply.MakeErrReply("ERR msetnx must within one slot in cluster mode")
-			}
+		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)
 		}
 	}
-	return cluster.relay(peer, c, args)
+
+	// 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 {
+			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
+	}
+	_, 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
 }