Example #1
0
File: set.go Project: scozss/setdb
func Sadd(args [][]byte, wb *levigo.WriteBatch) interface{} {
	var newMembers uint32
	key := NewKeyBuffer(SetKey, args[0], len(args[1]))
	mk := metaKey(args[0])
	card, err := scard(mk, nil)
	if err != nil {
		return err
	}

	for _, member := range args[1:] {
		key.SetSuffix(member)
		if card > 0 {
			res, err := DB.Get(DefaultReadOptions, key.Key())
			if err != nil {
				return err
			}
			if res != nil {
				continue
			}
		}
		wb.Put(key.Key(), []byte{})
		newMembers++
	}
	if newMembers > 0 {
		setCard(mk, card+newMembers, wb)
	}
	return newMembers
}
Example #2
0
File: hash.go Project: scozss/setdb
func Hincrby(args [][]byte, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	length, err := hlen(mk, nil)
	if err != nil {
		return err
	}
	key := NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key()
	res, err := DB.Get(DefaultReadOptions, key)
	if err != nil {
		return err
	}

	var current int64
	if res != nil {
		current, err = bconv.ParseInt(res, 10, 64)
		if err != nil {
			return InvalidIntError
		}
	}
	increment, err := bconv.ParseInt(args[2], 10, 64)
	if err != nil {
		return InvalidIntError
	}
	result := strconv.AppendInt(nil, current+increment, 10)
	wb.Put(key, result)

	// if is a new key, increment the hash length
	if res == nil {
		setHlen(mk, length+1, wb)
	}
	return result
}
Example #3
0
File: hash.go Project: scozss/setdb
func Hdel(args [][]byte, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	length, err := hlen(mk, nil)
	if err != nil {
		return err
	}
	if length == 0 {
		return 0
	}

	var deleted uint32
	key := NewKeyBuffer(HashKey, args[0], len(args[1]))
	for _, field := range args[1:] {
		key.SetSuffix(field)
		res, err := DB.Get(ReadWithoutCacheFill, key.Key())
		if err != nil {
			return err
		}
		if res == nil {
			continue
		}
		wb.Delete(key.Key())
		deleted++
	}
	if deleted == length {
		wb.Delete(mk)
	} else if deleted > 0 {
		setHlen(mk, length-deleted, wb)
	}
	return deleted
}
Example #4
0
File: hash.go Project: scozss/setdb
func Hincrbyfloat(args [][]byte, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	length, err := hlen(mk, nil)
	if err != nil {
		return err
	}
	key := NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key()
	res, err := DB.Get(DefaultReadOptions, key)
	if err != nil {
		return err
	}

	var current float64
	if res != nil {
		current, err = bconv.ParseFloat(res, 64)
		if err != nil {
			return fmt.Errorf("hash value is not a valid float")
		}
	}
	increment, err := bconv.ParseFloat(args[2], 64)
	if err != nil {
		return fmt.Errorf("value is not a valid float")
	}
	result := strconv.AppendFloat(nil, current+increment, 'f', -1, 64)
	wb.Put(key, result)

	// if is a new key, increment the hash length
	if res == nil {
		setHlen(mk, length+1, wb)
	}
	return result
}
Example #5
0
File: hash.go Project: scozss/setdb
func Hmset(args [][]byte, wb *levigo.WriteBatch) interface{} {
	if (len(args)-1)%2 != 0 {
		return fmt.Errorf("wrong number of arguments for 'hmset' command")
	}

	mk := metaKey(args[0])
	length, err := hlen(mk, nil)
	if err != nil {
		return err
	}

	var added uint32
	key := NewKeyBuffer(HashKey, args[0], len(args[1]))
	for i := 1; i < len(args); i += 2 {
		key.SetSuffix(args[i])
		var res []byte
		if length > 0 {
			res, err = DB.Get(DefaultReadOptions, key.Key())
			if err != nil {
				return err
			}
		}
		if res == nil {
			added++
		}
		wb.Put(key.Key(), args[i+1])
	}
	if added > 0 {
		setHlen(mk, length+added, wb)
	}
	return ReplyOK
}
Example #6
0
File: set.go Project: scozss/setdb
func Srem(args [][]byte, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	card, err := scard(mk, nil)
	if err != nil {
		return err
	}
	if card == 0 {
		return card
	}
	var deleted uint32
	key := NewKeyBuffer(SetKey, args[0], len(args[1]))
	for _, member := range args[1:] {
		key.SetSuffix(member)
		res, err := DB.Get(ReadWithoutCacheFill, key.Key())
		if err != nil {
			return err
		}
		if res == nil {
			continue
		}
		wb.Delete(key.Key())
		deleted++
	}
	if deleted == card {
		wb.Delete(mk)
	} else if deleted > 0 { // decrement the cardinality
		setCard(mk, card-deleted, wb)
	}
	return deleted
}
Example #7
0
File: list.go Project: scozss/setdb
func lpush(args [][]byte, left bool, create bool, wb *levigo.WriteBatch) (interface{}, error) {
	mk := metaKey(args[0])
	l, err := llen(mk, nil)
	if err != nil {
		return nil, err
	}
	if create || l.length > 0 {
		key := NewKeyBuffer(ListKey, args[0], 8)
		for _, value := range args[1:] {
			l.length++
			var seq int64
			if left {
				seq = l.left
				l.left--
			} else {
				seq = l.right
				l.right++
			}
			// To sort negative ints in order before positive, we subtract math.MinInt64
			// which wraps the numbers around and sorts correctly
			binary.BigEndian.PutUint64(key.SuffixForRead(8), uint64(seq-math.MinInt64))
			wb.Put(key.Key(), value)
		}
		setLlen(mk, l, wb)
	}
	return l.length, nil
}
Example #8
0
File: list.go Project: scozss/setdb
func setLlen(key []byte, l *listDetails, wb *levigo.WriteBatch) {
	data := make([]byte, 22)
	data[0] = ListLengthValue
	binary.BigEndian.PutUint32(data[1:], l.length)
	data[5] = l.flags
	binary.BigEndian.PutUint64(data[6:], uint64(l.left))
	binary.BigEndian.PutUint64(data[14:], uint64(l.right))
	wb.Put(key, data)
}
Example #9
0
func AddIndex(index [2]string, key []byte, indexDb *levigo.DB, wb *levigo.WriteBatch) error {
	searchKey := []byte(index[0] + "~" + index[1])
	keys, err := indexDb.Get(LReadOptions, searchKey)
	if err != nil {
		return err
	}
	keys = appendDataKey(keys, key)
	wb.Put(searchKey, keys)
	return nil
}
Example #10
0
File: set.go Project: scozss/setdb
func combineSet(keys [][]byte, op int, wb *levigo.WriteBatch) interface{} {
	var count uint32
	res := []interface{}{}
	members := make(chan *iterSetMember)
	var storeKey *KeyBuffer
	var mk []byte

	if wb != nil {
		mk = metaKey(keys[0])
		_, err := delKey(mk, wb)
		if err != nil {
			return err
		}
		storeKey = NewKeyBuffer(SetKey, keys[0], 0)
		keys = keys[1:]
	}

	go multiSetIter(keys, members, op != setUnion)

combine:
	for m := range members {
		switch op {
		case setInter:
			for _, k := range m.exists {
				if !k {
					continue combine
				}
			}
		case setDiff:
			for i, k := range m.exists {
				if i == 0 && !k || i > 0 && k {
					continue combine
				}
			}
		}
		if wb != nil {
			storeKey.SetSuffix(m.member)
			wb.Put(storeKey.Key(), []byte{})
			count++
		} else {
			res = append(res, m.member)
		}
	}

	if wb != nil {
		if count > 0 {
			setCard(mk, count, wb)
		}
		return count
	}

	return res
}
Example #11
0
File: hash.go Project: scozss/setdb
func DelHash(key []byte, wb *levigo.WriteBatch) {
	it := DB.NewIterator(ReadWithoutCacheFill)
	defer it.Close()
	iterKey := NewKeyBuffer(HashKey, key, 0)
	for it.Seek(iterKey.Key()); it.Valid(); it.Next() {
		k := it.Key()
		if !iterKey.IsPrefixOf(k) {
			break
		}
		wb.Delete(k)
	}
}
Example #12
0
// TODO: refactor with above.
func RemoveIndex(index [2]string, key []byte, indexDb *levigo.DB, wb *levigo.WriteBatch) error {
	searchKey := []byte(index[0] + "~" + index[1])
	keys, err := indexDb.Get(LReadOptions, searchKey)
	if err != nil {
		return err
	}
	keys = removeDataKey(keys, key)
	if len(keys) > 0 {
		wb.Put(searchKey, keys)
	} else {
		wb.Delete(searchKey)
	}
	return nil
}
Example #13
0
File: set.go Project: scozss/setdb
func DelSet(key []byte, wb *levigo.WriteBatch) {
	it := DB.NewIterator(ReadWithoutCacheFill)
	defer it.Close()
	iterKey := NewKeyBuffer(SetKey, key, 0)
	for it.Seek(iterKey.Key()); it.Valid(); it.Next() {
		k := it.Key()
		// If the prefix of the current key doesn't match the iteration key,
		// we have reached the end of the set
		if !iterKey.IsPrefixOf(k) {
			break
		}
		wb.Delete(k)
	}
}
Example #14
0
func delKey(key []byte, wb *levigo.WriteBatch) (deleted bool, err error) {
	res, err := DB.Get(ReadWithoutCacheFill, key)
	if err != nil {
		return
	}
	if res == nil {
		return
	}
	if len(res) < 1 {
		return false, InvalidDataError
	}
	del(key[1:], res[0], wb)
	wb.Delete(key)
	return true, nil
}
Example #15
0
File: list.go Project: scozss/setdb
func lpop(key []byte, left bool, wb *levigo.WriteBatch) (interface{}, error) {
	mk := metaKey(key)
	l, err := llen(mk, nil)
	if err != nil {
		return nil, err
	}
	if l.length == 0 {
		return nil, nil
	}

	iterKey := NewKeyBuffer(ListKey, key, 0)
	it := DB.NewIterator(ReadWithoutCacheFill)
	defer it.Close()
	if !left {
		iterKey.ReverseIterKey()
	}
	it.Seek(iterKey.Key())
	if !left {
		it.Prev()
	}
	if !it.Valid() {
		return nil, nil
	}
	k := it.Key()
	if !iterKey.IsPrefixOf(k) {
		return nil, nil
	}
	res := it.Value()

	wb.Delete(k)
	l.length--
	if l.length == 0 {
		wb.Delete(mk)
	} else {
		// decode the sequence number from the list item key
		seq := int64(binary.BigEndian.Uint64(k[len(key)+5:])) + math.MinInt64
		if left {
			l.left = seq
		} else {
			l.right = seq
		}
		setLlen(mk, l, wb)
	}

	return res, nil
}
Example #16
0
func set(k []byte, v []byte, wb *levigo.WriteBatch) error {
	mk := metaKey(k)
	res, err := DB.Get(DefaultReadOptions, mk)
	if err != nil {
		return err
	}

	// If there is a non-string key here, let's delete it first
	if len(res) > 0 && res[0] != StringLengthValue {
		del(k, res[0], wb)
	}

	setStringLen(mk, len(v), wb)
	wb.Put(stringKey(k), v)

	return nil
}
Example #17
0
File: zset.go Project: scozss/setdb
func DelZset(key []byte, wb *levigo.WriteBatch) {
	// TODO: count keys to verify everything works as expected?
	it := DB.NewIterator(ReadWithoutCacheFill)
	defer it.Close()
	iterKey := NewKeyBuffer(ZSetKey, key, 0)
	scoreKey := NewKeyBuffer(ZScoreKey, key, 0)

	for it.Seek(iterKey.Key()); it.Valid(); it.Next() {
		k := it.Key()
		// If the prefix of the current key doesn't match the iteration key,
		// we have reached the end of the zset
		if !iterKey.IsPrefixOf(k) {
			break
		}
		wb.Delete(k)
		setZScoreKeyMember(scoreKey, k[len(iterKey.Key()):])
		setZScoreKeyScore(scoreKey, btof(it.Value()))
		wb.Delete(scoreKey.Key())
	}
}
Example #18
0
func (s CommandSuite) TestCommands(c *C) {
	for _, t := range tests {
		cmd := commands[t.command]
		var wb *levigo.WriteBatch
		if cmd.writes {
			wb = levigo.NewWriteBatch()
		}
		var args [][]byte
		if t.args != "" {
			if cmd.arity > 0 {
				args = bytes.SplitN([]byte(t.args), []byte(" "), cmd.arity)
			} else {
				args = bytes.Split([]byte(t.args), []byte(" "))
			}
		}
		cmd.lockKeys(args)
		res := cmd.function(args, wb)
		if cmd.writes {
			err := DB.Write(DefaultWriteOptions, wb)
			c.Assert(err, IsNil)
			wb.Close()
		}
		cmd.unlockKeys(args)
		if stream, ok := res.(*cmdReplyStream); ok {
			items := make([]interface{}, 0, int(stream.size))
			for item := range stream.items {
				items = append(items, item)
			}
			res = items
		}
		if reply, ok := res.(rawReply); ok {
			if reply[0] == '+' {
				res = string(reply[1 : len(reply)-2])
			}
		}
		c.Assert(res, DeepEquals, t.response, Commentf("%s %s, obtained=%s expected=%s", t.command, t.args, res, t.response))
	}
}
Example #19
0
File: hash.go Project: scozss/setdb
func hset(args [][]byte, overwrite bool, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	length, err := hlen(mk, nil)
	if err != nil {
		return err
	}
	var res []byte
	key := NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key()
	if length > 0 {
		res, err = DB.Get(DefaultReadOptions, key)
		if err != nil {
			return err
		}
	}
	if overwrite || res == nil {
		wb.Put(key, args[2])
	}
	if res == nil {
		setHlen(mk, length+1, wb)
		return 1
	}
	return 0
}
Example #20
0
File: set.go Project: scozss/setdb
func Spop(args [][]byte, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	card, err := scard(mk, nil)
	if err != nil {
		return err
	}
	if card == 0 {
		return nil
	}
	key := NewKeyBuffer(SetKey, args[0], 1)
	member := srand(key)
	if member == nil {
		return nil
	}
	key.SetSuffix(member)
	wb.Delete(key.Key())
	if card == 1 { // we're removing the last remaining member
		wb.Delete(mk)
	} else {
		setCard(mk, card-1, wb)
	}
	return member
}
Example #21
0
File: zset.go Project: scozss/setdb
func Zrem(args [][]byte, wb *levigo.WriteBatch) interface{} {
	mk := metaKey(args[0])
	card, err := zcard(mk, nil)
	if err != nil {
		return err
	}
	if card == 0 {
		return 0
	}

	var deleted uint32
	setKey := NewKeyBuffer(ZSetKey, args[0], len(args[1]))
	scoreKey := NewKeyBuffer(ZScoreKey, args[0], 8+len(args[1]))
	// Delete each of the members
	for _, member := range args[1:] {
		setKey.SetSuffix(member)
		res, err := DB.Get(ReadWithoutCacheFill, setKey.Key())
		if err != nil {
			return nil
		}
		if res == nil {
			continue
		}
		if len(res) != 8 {
			return InvalidDataError
		}

		score := btof(res)
		setZScoreKeyMember(scoreKey, member)
		setZScoreKeyScore(scoreKey, score)
		wb.Delete(setKey.Key())
		wb.Delete(scoreKey.Key())
		deleted++
	}
	if deleted == card { // We deleted all of the members, so delete the meta key
		wb.Delete(mk)
	} else if deleted > 0 { // Decrement the cardinality
		setZcard(mk, card-deleted, wb)
	}

	return deleted
}
Example #22
0
File: zset.go Project: scozss/setdb
func zrangebyscore(args [][]byte, flag zrangeFlag, wb *levigo.WriteBatch) interface{} {
	// use a snapshot for this read so that the zcard is consistent
	snapshot := DB.NewSnapshot()
	opts := levigo.NewReadOptions()
	opts.SetSnapshot(snapshot)
	defer opts.Close()
	defer DB.ReleaseSnapshot(snapshot)

	mk := metaKey(args[0])
	card, err := zcard(mk, opts)
	if err != nil {
		return err
	}

	var minExclusive, maxExclusive bool
	if args[1][0] == '(' {
		minExclusive = true
		args[1] = args[1][1:]
	}
	if args[2][0] == '(' {
		maxExclusive = true
		args[2] = args[2][1:]
	}
	min, err := bconv.ParseFloat(args[1], 64)
	max, err2 := bconv.ParseFloat(args[2], 64)
	if err != nil || err2 != nil {
		return fmt.Errorf("min or max is not a float")
	}
	if flag == zrangeReverse {
		min, max = max, min
		minExclusive, maxExclusive = maxExclusive, minExclusive
	}

	if card == 0 || min > max {
		if flag <= zrangeReverse {
			return []interface{}{}
		} else {
			return 0
		}
	}

	// shortcut for zcount of -Inf to +Inf
	if flag == zrangeCount && math.IsInf(min, -1) && math.IsInf(max, 1) {
		return card
	}

	var withscores bool
	var offset, total int64 = 0, -1
	if flag <= zrangeReverse && len(args) > 3 {
		if len(args) == 4 || len(args) == 7 {
			if EqualIgnoreCase(args[3], []byte("withscores")) {
				withscores = true
			}
		} else if len(args) != 6 {
			return SyntaxError
		}
		if len(args) >= 6 {
			if EqualIgnoreCase(args[len(args)-3], []byte("limit")) {
				offset, err = bconv.ParseInt(args[len(args)-2], 10, 64)
				total, err2 = bconv.ParseInt(args[len(args)-1], 10, 64)
				if err != nil || err2 != nil {
					return InvalidIntError
				}
				if offset < 0 || total < 1 {
					return []interface{}{}
				}
			} else {
				return SyntaxError
			}
		}
	}

	it := DB.NewIterator(opts)
	defer it.Close()

	var deleted, count uint32
	var deleteKey *KeyBuffer
	if flag == zrangeDelete {
		deleteKey = NewKeyBuffer(ZSetKey, args[0], 0)
	}

	res := []interface{}{}
	iterKey := NewKeyBuffer(ZScoreKey, args[0], 8)
	if flag != zrangeReverse {
		setZScoreKeyScore(iterKey, min)
	} else {
		setZScoreKeyScore(iterKey, max)
		iterKey.ReverseIterKey()
	}
	it.Seek(iterKey.Key())
	for i := int64(0); it.Valid(); i++ {
		if flag == zrangeReverse {
			it.Prev()
		}
		if i >= offset {
			if total > -1 && i-offset >= total {
				break
			}
			k := it.Key()
			if !iterKey.IsPrefixOf(k) {
				break
			}
			score, member := parseZScoreKey(it.Key(), len(args[0]))
			if (!minExclusive && score < min) || (minExclusive && score <= min) {
				if flag != zrangeReverse {
					it.Next()
				}
				continue
			}
			if (!maxExclusive && score > max) || (maxExclusive && score >= max) {
				break
			}
			if flag <= zrangeReverse {
				res = append(res, member)
				if withscores {
					res = append(res, ftoa(score))
				}
			}
			if flag == zrangeDelete {
				deleteKey.SetSuffix(member)
				wb.Delete(k)
				wb.Delete(deleteKey.Key())
				deleted++
			}
			if flag == zrangeCount {
				count++
			}
		}
		if flag != zrangeReverse {
			it.Next()
		}
	}

	if flag == zrangeDelete && deleted == card {
		wb.Delete(mk)
	} else if deleted > 0 {
		setZcard(mk, card-deleted, wb)
	}
	if flag == zrangeDelete {
		return deleted
	}
	if flag == zrangeCount {
		return count
	}

	return res
}
Example #23
0
func setStringLen(key []byte, length int, wb *levigo.WriteBatch) {
	meta := make([]byte, 5)
	meta[0] = StringLengthValue
	binary.BigEndian.PutUint32(meta[1:], uint32(length))
	wb.Put(key, meta)
}
Example #24
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]
	}
}
Example #25
0
File: set.go Project: scozss/setdb
func setCard(key []byte, card uint32, wb *levigo.WriteBatch) {
	data := make([]byte, 5)
	data[0] = SetCardValue
	binary.BigEndian.PutUint32(data[1:], card)
	wb.Put(key, data)
}
Example #26
0
File: hash.go Project: scozss/setdb
func setHlen(key []byte, length uint32, wb *levigo.WriteBatch) {
	data := make([]byte, 5)
	data[0] = HashLengthValue
	binary.BigEndian.PutUint32(data[1:], length)
	wb.Put(key, data)
}
Example #27
0
File: zset.go Project: 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
}
Example #28
0
File: zset.go Project: scozss/setdb
func zadd(args [][]byte, wb *levigo.WriteBatch, incr bool) interface{} {
	var newMembers uint32
	var score float64
	scoreBytes := make([]byte, 8)
	setKey := NewKeyBuffer(ZSetKey, args[0], len(args[2]))
	scoreKey := NewKeyBuffer(ZScoreKey, args[0], 8+len(args[2]))

	mk := metaKey(args[0])
	card, err := zcard(mk, nil)
	if err != nil {
		return err
	}

	// Iterate through each of the score/member pairs
	for i := 1; i < len(args); i += 2 {
		var err error
		score, err = bconv.ParseFloat(args[i], 64)
		if err != nil {
			return fmt.Errorf("'%s' is not a valid float", string(args[1]))
		}

		// Check if the member exists
		setKey.SetSuffix(args[i+1])
		var res []byte
		if card > 0 {
			res, err = DB.Get(DefaultReadOptions, setKey.Key())
			if err != nil {
				return err
			}
		}

		// set the score key with 8 empty bytes before the member for the score
		setZScoreKeyMember(scoreKey, args[i+1])
		if res != nil { // We got a score from the db, so the member already exists
			if len(res) != 8 {
				return InvalidDataError
			}
			actualScore := math.Float64frombits(binary.BigEndian.Uint64(res))
			if incr { // this is a ZINCRBY, so increment the score
				score += actualScore
			}
			if score == actualScore { // Member already exists with the same score, do nothing
				continue
			}

			// Delete score key for member
			setZScoreKeyScore(scoreKey, actualScore)
			wb.Delete(scoreKey.Key())
		} else { // No score found, we're adding a new member
			newMembers++
		}

		// Store the set and score keys
		binary.BigEndian.PutUint64(scoreBytes, math.Float64bits(score))
		setZScoreKeyScore(scoreKey, score)
		wb.Put(setKey.Key(), scoreBytes)
		wb.Put(scoreKey.Key(), []byte{}) // The score key is only used for sorting, the value is empty
	}

	// Update the set metadata with the new cardinality
	if newMembers > 0 {
		setZcard(mk, card+newMembers, wb)
	}

	if incr { // This is a ZINCRBY, return the new score
		return ftoa(score)
	}
	return newMembers
}
Example #29
0
func DelString(key []byte, wb *levigo.WriteBatch) {
	wb.Delete(stringKey(key))
}