github.com/hdt3213/godis@v1.2.9/cluster/mset.go (about) 1 package cluster 2 3 import ( 4 "fmt" 5 "github.com/hdt3213/godis/interface/redis" 6 "github.com/hdt3213/godis/lib/utils" 7 "github.com/hdt3213/godis/redis/protocol" 8 "strconv" 9 ) 10 11 const keyExistsErr = "key exists" 12 13 // MGet atomically get multi key-value from cluster, writeKeys can be distributed on any node 14 func MGet(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply { 15 if len(cmdLine) < 2 { 16 return protocol.MakeErrReply("ERR wrong number of arguments for 'mget' command") 17 } 18 keys := make([]string, len(cmdLine)-1) 19 for i := 1; i < len(cmdLine); i++ { 20 keys[i-1] = string(cmdLine[i]) 21 } 22 23 resultMap := make(map[string][]byte) 24 groupMap := cluster.groupBy(keys) 25 for peer, group := range groupMap { 26 resp := cluster.relay(peer, c, makeArgs("MGET", group...)) 27 if protocol.IsErrorReply(resp) { 28 errReply := resp.(protocol.ErrorReply) 29 return protocol.MakeErrReply(fmt.Sprintf("ERR during get %s occurs: %v", group[0], errReply.Error())) 30 } 31 arrReply, _ := resp.(*protocol.MultiBulkReply) 32 for i, v := range arrReply.Args { 33 key := group[i] 34 resultMap[key] = v 35 } 36 } 37 result := make([][]byte, len(keys)) 38 for i, k := range keys { 39 result[i] = resultMap[k] 40 } 41 return protocol.MakeMultiBulkReply(result) 42 } 43 44 // MSet atomically sets multi key-value in cluster, writeKeys can be distributed on any node 45 func MSet(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply { 46 argCount := len(cmdLine) - 1 47 if argCount%2 != 0 || argCount < 1 { 48 return protocol.MakeErrReply("ERR wrong number of arguments for 'mset' command") 49 } 50 51 size := argCount / 2 52 keys := make([]string, size) 53 valueMap := make(map[string]string) 54 for i := 0; i < size; i++ { 55 keys[i] = string(cmdLine[2*i+1]) 56 valueMap[keys[i]] = string(cmdLine[2*i+2]) 57 } 58 59 groupMap := cluster.groupBy(keys) 60 if len(groupMap) == 1 && allowFastTransaction { // do fast 61 for peer := range groupMap { 62 return cluster.relay(peer, c, cmdLine) 63 } 64 } 65 66 //prepare 67 var errReply redis.Reply 68 txID := cluster.idGenerator.NextID() 69 txIDStr := strconv.FormatInt(txID, 10) 70 rollback := false 71 for peer, group := range groupMap { 72 peerArgs := []string{txIDStr, "MSET"} 73 for _, k := range group { 74 peerArgs = append(peerArgs, k, valueMap[k]) 75 } 76 var resp redis.Reply 77 if peer == cluster.self { 78 resp = execPrepare(cluster, c, makeArgs("Prepare", peerArgs...)) 79 } else { 80 resp = cluster.relay(peer, c, makeArgs("Prepare", peerArgs...)) 81 } 82 if protocol.IsErrorReply(resp) { 83 errReply = resp 84 rollback = true 85 break 86 } 87 } 88 if rollback { 89 // rollback 90 requestRollback(cluster, c, txID, groupMap) 91 } else { 92 _, errReply = requestCommit(cluster, c, txID, groupMap) 93 rollback = errReply != nil 94 } 95 if !rollback { 96 return &protocol.OkReply{} 97 } 98 return errReply 99 100 } 101 102 // MSetNX sets multi key-value in database, only if none of the given writeKeys exist and all given writeKeys are on the same node 103 func MSetNX(cluster *Cluster, c redis.Connection, cmdLine CmdLine) redis.Reply { 104 argCount := len(cmdLine) - 1 105 if argCount%2 != 0 || argCount < 1 { 106 return protocol.MakeErrReply("ERR wrong number of arguments for 'msetnx' command") 107 } 108 109 size := argCount / 2 110 keys := make([]string, size) 111 valueMap := make(map[string]string) 112 for i := 0; i < size; i++ { 113 keys[i] = string(cmdLine[2*i+1]) 114 valueMap[keys[i]] = string(cmdLine[2*i+2]) 115 } 116 117 groupMap := cluster.groupBy(keys) 118 if len(groupMap) == 1 && allowFastTransaction { // do fast 119 for peer := range groupMap { 120 return cluster.relay(peer, c, cmdLine) 121 } 122 } 123 124 // prepare procedure: 125 // 1. Normal tcc preparation (undo log and lock related keys) 126 // 2. Peer checks whether any key already exists, If so it will return keyExistsErr. Then coordinator will request rollback over all participated nodes 127 var errReply redis.Reply 128 txID := cluster.idGenerator.NextID() 129 txIDStr := strconv.FormatInt(txID, 10) 130 rollback := false 131 for node, group := range groupMap { 132 nodeArgs := []string{txIDStr, "MSETNX"} 133 for _, k := range group { 134 nodeArgs = append(nodeArgs, k, valueMap[k]) 135 } 136 resp := cluster.relayPrepare(node, c, makeArgs("Prepare", nodeArgs...)) 137 if protocol.IsErrorReply(resp) { 138 re := resp.(protocol.ErrorReply) 139 if re.Error() == keyExistsErr { 140 errReply = protocol.MakeIntReply(0) 141 } else { 142 errReply = resp 143 } 144 rollback = true 145 break 146 } 147 } 148 if rollback { 149 // rollback 150 requestRollback(cluster, c, txID, groupMap) 151 return errReply 152 } 153 _, errReply = requestCommit(cluster, c, txID, groupMap) 154 rollback = errReply != nil 155 if !rollback { 156 return protocol.MakeIntReply(1) 157 } 158 return errReply 159 } 160 161 func prepareMSetNx(cluster *Cluster, conn redis.Connection, cmdLine CmdLine) redis.Reply { 162 args := cmdLine[1:] 163 if len(args)%2 != 0 { 164 return protocol.MakeSyntaxErrReply() 165 } 166 size := len(args) / 2 167 values := make([][]byte, size) 168 keys := make([]string, size) 169 for i := 0; i < size; i++ { 170 keys[i] = string(args[2*i]) 171 values[i] = args[2*i+1] 172 } 173 re := cluster.db.ExecWithLock(conn, utils.ToCmdLine2("ExistIn", keys...)) 174 if protocol.IsErrorReply(re) { 175 return re 176 } 177 _, ok := re.(*protocol.EmptyMultiBulkReply) 178 if !ok { 179 return protocol.MakeErrReply(keyExistsErr) 180 } 181 return protocol.MakeOkReply() 182 } 183 184 func init() { 185 registerPrepareFunc("MSetNx", prepareMSetNx) 186 }