Esempio n. 1
0
File: zset.go Progetto: scozss/setdb
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
}
Esempio n. 2
0
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]
	}
}
Esempio n. 3
0
File: zset.go Progetto: scozss/setdb
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
}