兼容go redis cluster的pipeline批量

// xiaorui.cc

func (cluster *Cluster) Do(cmd string, args ...interface{}) (interface{}, error) {
    if len(args) < 1 {
    return nil, fmt.Errorf("Do: no key found in args")
    }

    if cmd == "MSET" || cmd == "MSETNX" {
    return cluster.multiSet(cmd, args...)
    }

    if cmd == "MGET" {
    return cluster.multiGet(cmd, args...)
    }
    ...
}


func (cluster *Cluster) multiSet(cmd string, args ...interface{}) (interface{}, error) {
    ...
    // 命令分组
    tasks := make([]*multiTask, 0)

    cluster.rwLock.RLock()
    for i := 0; i < len(args); i += 2 {
        key, err := key(args[i])
        if err != nil {
            cluster.rwLock.RUnlock()
            return nil, fmt.Errorf("multiSet: invalid key %v", args[i])
        }

        slot := hash(key)

        var j int
        for j = 0; j < len(tasks); j++ {
            // 相同的slot的key放在一个multiTask的slice里。
            if tasks[j].slot == slot {
                tasks[j].args = append(tasks[j].args, args[i])   // key
                tasks[j].args = append(tasks[j].args, args[i+1]) // value

                break
            }
        }

        if j == len(tasks) {
            node := cluster.slots[slot]
            if node == nil {
                cluster.rwLock.RUnlock()
                return nil, fmt.Errorf("multiSet: %s[%d] no node found", key, slot)
            }

            task := &multiTask{
                node: node,
                slot: slot,
                cmd:  cmd,
                args: []interface{}{args[i], args[i+1]},
                done: make(chan int),
            }
            tasks = append(tasks, task)
        }
    }
    cluster.rwLock.RUnlock()

    // 每个slot队列为为一个并发
    for i := range tasks {
        go handleSetTask(tasks[i])
    }

    // 确定都执行完
    for i := range tasks {
        <-tasks[i].done
    }

    for i := range tasks {
        _, err := String(tasks[i].reply, tasks[i].err)
        if err != nil {
            return nil, err
        }
    }

    return "OK", nil
}

type multiTask struct {
    node *redisNode
    slot uint16

    cmd  string
    args []interface{}

    reply   interface{}
    replies []interface{}
    err     error

    done chan int
}

// xiaorui.cc

// NewBatch create a new batch to pack mutiple commands.
func (cluster *Cluster) NewBatch() *Batch {
    return &Batch{
        cluster: cluster,
        batches: make([]nodeBatch, 0),
        index:   make([]int, 0),
    }
}

// Put add a redis command to batch, DO NOT put MGET/MSET/MSETNX.
func (batch *Batch) Put(cmd string, args ...interface{}) error {
    if len(args) < 1 {
        return fmt.Errorf("Put: no key found in args")
    }

    if cmd == "MGET" || cmd == "MSET" || cmd == "MSETNX" {
        return fmt.Errorf("Put: %s not supported", cmd)
    }

    node, err := batch.cluster.getNodeByKey(args[0])
    if err != nil {
        return fmt.Errorf("Put: %v", err)
    }

    var i int
    for i = 0; i < len(batch.batches); i++ {
        if batch.batches[i].node == node {
            batch.batches[i].cmds = append(batch.batches[i].cmds,
                nodeCommand{cmd: cmd, args: args})

            batch.index = append(batch.index, i)
            break
        }
    }
    ...
}

func (cluster *Cluster) RunBatch(bat *Batch) ([]interface{}, error) {
    for i := range bat.batches {
        go doBatch(&bat.batches[i])
    }

    for i := range bat.batches {
        <-bat.batches[i].done
    }

    var replies []interface{}
    for _, i := range bat.index {
        if bat.batches[i].err != nil {
            return nil, bat.batches[i].err
        }

        replies = append(replies, bat.batches[i].cmds[0].reply)
        bat.batches[i].cmds = bat.batches[i].cmds[1:]
    }

    return replies, nil
}

func doBatch(batch *nodeBatch) {
    conn, err := batch.node.getConn()
    ...

    for i := range batch.cmds {
        conn.send(batch.cmds[i].cmd, batch.cmds[i].args...)
    }

    err = conn.flush()
    ...

    for i := range batch.cmds {
        reply, err := conn.receive()
        if err != nil {
            batch.err = err
            conn.shutdown()
            batch.done <- 1
            return
        }

        batch.cmds[i].reply, batch.cmds[i].err = reply, err
    }

    ...
}

// xiaorui.cc

func (c *ClusterClient) defaultProcessTxPipeline(cmds []Cmder) error {
        ...
	cmdsMap := c.mapCmdsBySlot(cmds)
	for slot, cmds := range cmdsMap {
		node, err := state.slotMasterNode(slot)
		if err != nil {
			setCmdsErr(cmds, err)
			continue
		}
		cmdsMap := map[*clusterNode][]Cmder{node: cmds}

		for attempt := 0; attempt <= c.opt.MaxRedirects; attempt++ {
			if attempt > 0 {
				time.Sleep(c.retryBackoff(attempt))
			}

			failedCmds := newCmdsMap()
			var wg sync.WaitGroup

			for node, cmds := range cmdsMap {
				wg.Add(1)
				go func(node *clusterNode, cmds []Cmder) {
					defer wg.Done()

					cn, err := node.Client.getConn()
					if err != nil {
						if err == pool.ErrClosed {
							c.mapCmdsByNode(cmds, failedCmds)
						} else {
							setCmdsErr(cmds, err)
						}
						return
					}

					err = c.txPipelineProcessCmds(node, cn, cmds, failedCmds)
					node.Client.releaseConnStrict(cn, err)
				}(node, cmds)
			}

			wg.Wait()
			if len(failedCmds.m) == 0 {
				break
			}
			cmdsMap = failedCmds.m
		}
	}

	return cmdsFirstErr(cmds)
}