func ZunionInterKeys(args [][]byte) [][]byte {
numKeys, err := bconv.Atoi(args[1])
// don't return any keys if the response will be a syntax error
if err != nil || len(args) < 2+numKeys {
return nil
}
keys := make([][]byte, 1, 1+numKeys)
keys[0] = args[0]
keyloop:
for _, k := range args[2 : 2+numKeys] {
for _, key := range keys {
// skip keys that are already in the array
if bytes.Equal(k, key) {
continue keyloop
}
}
keys = append(keys, k)
}
return keys
}
func protocolHandler(c *client) {
// Read a length (looks like "$3\r\n")
readLength := func(prefix byte) (length int, err error) {
b, err := c.r.ReadByte()
if err != nil {
return
}
if b != prefix {
writeProtocolError(c.w, "invalid length")
return
}
l, overflowed, err := c.r.ReadLine() // Read bytes will look like "123"
if err != nil {
return
}
if overflowed {
writeProtocolError(c.w, "length line too long")
return
}
if len(l) == 0 {
writeProtocolError(c.w, "missing length")
return
}
length, err = bconv.Atoi(l)
if err != nil {
writeProtocolError(c.w, "length is not a valid integer")
return
}
return
}
runCommand := func(args [][]byte) (err error) {
if len(args) == 0 {
writeProtocolError(c.w, "missing command")
return
}
// lookup the command
command, ok := commands[UnsafeBytesToString(bytes.ToLower(args[0]))]
if !ok {
writeError(c.w, "unknown command '"+string(args[0])+"'")
return
}
// check command arity, negative arity means >= n
if (command.arity < 0 && len(args)-1 < -command.arity) || (command.arity >= 0 && len(args)-1 > command.arity) {
writeError(c.w, "wrong number of arguments for '"+string(args[0])+"' command")
return
}
// call the command and respond
var wb *levigo.WriteBatch
if command.writes {
wb = levigo.NewWriteBatch()
defer wb.Close()
}
command.lockKeys(args[1:])
res := command.function(args[1:], wb)
if command.writes {
if _, ok := res.(error); !ok { // only write the batch if the return value is not an error
err = DB.Write(DefaultWriteOptions, wb)
}
if err != nil {
writeError(c.w, "data write error: "+err.Error())
return
}
}
command.unlockKeys(args[1:])
writeReply(c.w, res)
return
}
processInline := func() error {
line, err := c.r.ReadBytes('\n')
if err != nil {
return err
}
return runCommand(bytes.Split(line[:len(line)-2], []byte(" ")))
}
scratch := make([]byte, 2)
args := [][]byte{}
// Client event loop, each iteration handles a command
for {
// check if we're using the old inline protocol
b, err := c.r.Peek(1)
if err != nil {
return
}
if b[0] != '*' {
err = processInline()
if err != nil {
return
}
continue
}
// Step 1: get the number of arguments
argCount, err := readLength('*')
if err != nil {
return
}
// read the arguments
for i := 0; i < argCount; i++ {
length, err := readLength('$')
if err != nil {
return
}
// Read the argument bytes
args = append(args, make([]byte, length))
_, err = io.ReadFull(c.r, args[i])
if err != nil {
return
}
// The argument has a trailing \r\n that we need to discard
c.r.Read(scratch) // TODO: make sure these bytes are read
}
err = runCommand(args)
if err != nil {
return
}
// Truncate arguments for the next run
args = args[:0]
}
}
func combineZset(args [][]byte, op int, wb *levigo.WriteBatch) interface{} {
var count uint32
res := []interface{}{}
members := make(chan *iterZsetMember)
var setKey, scoreKey *KeyBuffer
scoreBytes := make([]byte, 8)
mk := metaKey(args[0])
if wb != nil {
_, err := delKey(mk, wb)
if err != nil {
return err
}
setKey = NewKeyBuffer(ZSetKey, args[0], 0)
scoreKey = NewKeyBuffer(ZScoreKey, args[0], 0)
}
numKeys, err := bconv.Atoi(args[1])
if err != nil {
return InvalidIntError
}
aggregate := zsetAggSum
weights := make([]float64, numKeys)
scores := make([]float64, 0, numKeys)
for i := 0; i < numKeys; i++ {
weights[i] = 1
}
argOffset := 2 + numKeys
if len(args) > argOffset {
if len(args) > argOffset+numKeys {
if EqualIgnoreCase(args[argOffset], []byte("weights")) {
argOffset += numKeys + 1
for i, w := range args[numKeys+3 : argOffset] {
weights[i], err = bconv.ParseFloat(w, 64)
if err != nil {
return fmt.Errorf("weight value is not a float")
}
}
} else {
return SyntaxError
}
}
if len(args) > argOffset {
if len(args) == argOffset+2 && EqualIgnoreCase(args[argOffset], []byte("aggregate")) {
agg := bytes.ToLower(args[argOffset+1])
switch {
case bytes.Equal(agg, []byte("sum")):
aggregate = zsetAggSum
case bytes.Equal(agg, []byte("min")):
aggregate = zsetAggMin
case bytes.Equal(agg, []byte("max")):
aggregate = zsetAggMax
default:
return SyntaxError
}
} else {
return SyntaxError
}
}
}
go multiZsetIter(args[2:numKeys+2], members, op != zsetUnion)
combine:
for m := range members {
if op == zsetInter {
for _, k := range m.exists {
if !k {
continue combine
}
}
}
scores = scores[:0]
for i, k := range m.exists {
if k {
scores = append(scores, m.scores[i])
}
}
var score float64
for i, s := range scores {
scores[i] = s * weights[i]
}
switch aggregate {
case zsetAggSum:
for _, s := range scores {
score += s
}
case zsetAggMin:
sort.Float64s(scores)
score = scores[0]
case zsetAggMax:
sort.Float64s(scores)
score = scores[len(scores)-1]
}
if wb != nil {
setKey.SetSuffix(m.member)
binary.BigEndian.PutUint64(scoreBytes, math.Float64bits(score))
setZScoreKeyMember(scoreKey, m.member)
setZScoreKeyScore(scoreKey, score)
wb.Put(setKey.Key(), scoreBytes)
wb.Put(scoreKey.Key(), []byte{})
count++
} else {
res = append(res, m.member, ftoa(score))
}
}
if wb != nil {
if count > 0 {
setZcard(mk, count, wb)
}
return count
}
return res
}
