support mset in cluster

This commit is contained in:
finley
2025-03-12 08:25:41 +08:00
parent 43dc28d335
commit 360de97479
9 changed files with 137 additions and 15 deletions

2
cluster/commands/mget.go Normal file
View File

@@ -0,0 +1,2 @@
package commands

View File

@@ -8,13 +8,28 @@ import (
)
func init() {
core.RegisterCmd("mset_", execMSet_)
core.RegisterCmd("mset_", execMSetInLocal)
core.RegisterCmd("mset", execMSet)
}
type CmdLine = [][]byte
// execMSet_ executes msets in local node
func execMSet_(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
// node -> keys on the node
type RouteMap map[string][]string
func getRouteMap(cluster *core.Cluster, keys []string) RouteMap {
m := make(RouteMap)
for _, key := range keys {
slot := cluster.GetSlot(key)
node := cluster.PickNode(slot)
m[node] = append(m[node], key)
}
return m
}
// execMSetInLocal executes msets in local node
func execMSetInLocal(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
if len(cmdLine) < 3 {
return protocol.MakeArgNumErrReply("mset")
}
@@ -22,15 +37,34 @@ func execMSet_(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine) redis
return cluster.LocalExec(c, cmdLine)
}
func requestRollback(cluster *core.Cluster, c redis.Connection, txId string, routeMap map[string][]string) {
func execMSet(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply {
if len(cmdLine) < 3 || len(cmdLine)%2 != 1 {
return protocol.MakeArgNumErrReply("mset")
}
var keys []string
for i := 1; i < len(cmdLine); i += 2 {
keys = append(keys, string(cmdLine[i]))
}
routeMap := getRouteMap(cluster, keys)
if len(routeMap) == 1 {
// only one node, do it fast
for node := range routeMap {
cmdLine[0] = []byte("mset_")
return cluster.Relay(node, c, cmdLine)
}
}
return execMSetSlow(cluster, c, cmdLine, routeMap)
}
func requestRollback(cluster *core.Cluster, c redis.Connection, txId string, routeMap RouteMap) {
rollbackCmd := utils.ToCmdLine("rollback", txId)
for node := range routeMap {
for node := range routeMap {
cluster.Relay(node, c, rollbackCmd)
}
}
// execMSetSlow execute mset through tcc
func execMSetSlow(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine, routeMap map[string][]string) redis.Reply {
func execMSetSlow(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine, routeMap RouteMap) redis.Reply {
txId := utils.RandString(6)
keyValues := make(map[string][]byte)
@@ -62,7 +96,7 @@ func execMSetSlow(cluster *core.Cluster, c redis.Connection, cmdLine CmdLine, ro
// send commit request
commiteCmd := utils.ToCmdLine("commit", txId)
for node := range nodePrepareCmdMap {
for node := range nodePrepareCmdMap {
reply := cluster.Relay(node, c, commiteCmd)
if protocol.IsErrorReply(reply) {
requestRollback(cluster, c, txId, routeMap)

View File

@@ -0,0 +1,20 @@
package commands
import (
"testing"
"github.com/hdt3213/godis/cluster/core"
"github.com/hdt3213/godis/lib/utils"
"github.com/hdt3213/godis/redis/connection"
)
func TestMset(t *testing.T) {
id1 := "1"
id2 := "2"
nodes := core.MakeTestCluster([]string{id1, id2})
node1 := nodes[id1]
c := connection.NewFakeConn()
// 1, 2 will be routed to node1 and node2, see MakeTestCluster
res := execMSet(node1, c, utils.ToCmdLine("mset", "1", "1", "2", "2"))
println(res)
}