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  }