ledisdb/ledis/t_zset.go

1007 lines
20 KiB
Go
Raw Normal View History

2014-05-09 10:49:22 +04:00
package ledis
2014-05-08 06:54:33 +04:00
import (
"bytes"
"encoding/binary"
"errors"
2014-07-25 13:58:00 +04:00
"github.com/siddontang/ledisdb/store"
"time"
2014-05-08 06:54:33 +04:00
)
const (
2014-06-04 14:32:23 +04:00
MinScore int64 = -1<<63 + 1
MaxScore int64 = 1<<63 - 1
InvalidScore int64 = -1 << 63
AggregateSum byte = 0
AggregateMin byte = 1
AggregateMax byte = 2
)
2014-05-16 04:56:32 +04:00
type ScorePair struct {
Score int64
Member []byte
}
2014-05-08 06:54:33 +04:00
var errZSizeKey = errors.New("invalid zsize key")
var errZSetKey = errors.New("invalid zset key")
var errZScoreKey = errors.New("invalid zscore key")
var errScoreOverflow = errors.New("zset score overflow")
var errInvalidAggregate = errors.New("invalid aggregate")
var errInvalidWeightNum = errors.New("invalid weight number")
var errInvalidSrcKeyNum = errors.New("invalid src key number")
2014-05-08 06:54:33 +04:00
const (
zsetNScoreSep byte = '<'
zsetPScoreSep byte = zsetNScoreSep + 1
zsetStopScoreSep byte = zsetPScoreSep + 1
zsetStartMemSep byte = ':'
zsetStopMemSep byte = zsetStartMemSep + 1
)
2014-05-20 04:41:24 +04:00
func checkZSetKMSize(key []byte, member []byte) error {
2014-05-23 08:23:20 +04:00
if len(key) > MaxKeySize || len(key) == 0 {
2014-06-05 11:46:38 +04:00
return errKeySize
2014-05-23 08:23:20 +04:00
} else if len(member) > MaxZSetMemberSize || len(member) == 0 {
2014-06-05 11:46:38 +04:00
return errZSetMemberSize
2014-05-20 04:41:24 +04:00
}
return nil
}
func (db *DB) zEncodeSizeKey(key []byte) []byte {
buf := make([]byte, len(key)+2)
buf[0] = db.index
buf[1] = ZSizeType
2014-05-08 06:54:33 +04:00
2014-05-20 04:41:24 +04:00
copy(buf[2:], key)
2014-05-08 06:54:33 +04:00
return buf
}
2014-05-20 04:41:24 +04:00
func (db *DB) zDecodeSizeKey(ek []byte) ([]byte, error) {
if len(ek) < 2 || ek[0] != db.index || ek[1] != ZSizeType {
2014-05-08 06:54:33 +04:00
return nil, errZSizeKey
}
2014-05-20 04:41:24 +04:00
return ek[2:], nil
2014-05-08 06:54:33 +04:00
}
2014-05-20 04:41:24 +04:00
func (db *DB) zEncodeSetKey(key []byte, member []byte) []byte {
2014-05-25 10:01:35 +04:00
buf := make([]byte, len(key)+len(member)+5)
2014-05-08 06:54:33 +04:00
pos := 0
2014-05-20 04:41:24 +04:00
buf[pos] = db.index
pos++
buf[pos] = ZSetType
2014-05-08 06:54:33 +04:00
pos++
2014-05-20 04:41:24 +04:00
binary.BigEndian.PutUint16(buf[pos:], uint16(len(key)))
pos += 2
2014-05-08 06:54:33 +04:00
copy(buf[pos:], key)
pos += len(key)
2014-05-25 10:01:35 +04:00
buf[pos] = zsetStartMemSep
pos++
2014-05-08 06:54:33 +04:00
copy(buf[pos:], member)
return buf
}
2014-05-20 04:41:24 +04:00
func (db *DB) zDecodeSetKey(ek []byte) ([]byte, []byte, error) {
if len(ek) < 5 || ek[0] != db.index || ek[1] != ZSetType {
2014-05-08 06:54:33 +04:00
return nil, nil, errZSetKey
}
2014-05-20 04:41:24 +04:00
keyLen := int(binary.BigEndian.Uint16(ek[2:]))
2014-05-25 10:01:35 +04:00
if keyLen+5 > len(ek) {
2014-05-08 06:54:33 +04:00
return nil, nil, errZSetKey
}
2014-05-20 04:41:24 +04:00
key := ek[4 : 4+keyLen]
2014-05-25 10:01:35 +04:00
if ek[4+keyLen] != zsetStartMemSep {
return nil, nil, errZSetKey
}
member := ek[5+keyLen:]
2014-05-08 06:54:33 +04:00
return key, member, nil
}
2014-05-25 10:01:35 +04:00
func (db *DB) zEncodeStartSetKey(key []byte) []byte {
k := db.zEncodeSetKey(key, nil)
return k
}
func (db *DB) zEncodeStopSetKey(key []byte) []byte {
k := db.zEncodeSetKey(key, nil)
k[len(k)-1] = zsetStartMemSep + 1
return k
}
2014-05-20 04:41:24 +04:00
func (db *DB) zEncodeScoreKey(key []byte, member []byte, score int64) []byte {
buf := make([]byte, len(key)+len(member)+14)
2014-05-08 06:54:33 +04:00
pos := 0
2014-05-20 04:41:24 +04:00
buf[pos] = db.index
pos++
buf[pos] = ZScoreType
2014-05-08 06:54:33 +04:00
pos++
2014-05-20 04:41:24 +04:00
binary.BigEndian.PutUint16(buf[pos:], uint16(len(key)))
pos += 2
2014-05-08 06:54:33 +04:00
copy(buf[pos:], key)
pos += len(key)
if score < 0 {
buf[pos] = zsetNScoreSep
} else {
buf[pos] = zsetPScoreSep
}
pos++
binary.BigEndian.PutUint64(buf[pos:], uint64(score))
pos += 8
buf[pos] = zsetStartMemSep
pos++
copy(buf[pos:], member)
return buf
}
2014-05-20 04:41:24 +04:00
func (db *DB) zEncodeStartScoreKey(key []byte, score int64) []byte {
return db.zEncodeScoreKey(key, nil, score)
2014-05-08 06:54:33 +04:00
}
2014-05-20 04:41:24 +04:00
func (db *DB) zEncodeStopScoreKey(key []byte, score int64) []byte {
k := db.zEncodeScoreKey(key, nil, score)
2014-05-08 06:54:33 +04:00
k[len(k)-1] = zsetStopMemSep
return k
}
2014-05-20 04:41:24 +04:00
func (db *DB) zDecodeScoreKey(ek []byte) (key []byte, member []byte, score int64, err error) {
if len(ek) < 14 || ek[0] != db.index || ek[1] != ZScoreType {
2014-05-08 06:54:33 +04:00
err = errZScoreKey
return
}
2014-05-20 04:41:24 +04:00
keyLen := int(binary.BigEndian.Uint16(ek[2:]))
2014-05-08 06:54:33 +04:00
if keyLen+14 > len(ek) {
err = errZScoreKey
return
}
2014-05-20 04:41:24 +04:00
key = ek[4 : 4+keyLen]
pos := 4 + keyLen
2014-05-08 06:54:33 +04:00
if (ek[pos] != zsetNScoreSep) && (ek[pos] != zsetPScoreSep) {
err = errZScoreKey
return
}
pos++
score = int64(binary.BigEndian.Uint64(ek[pos:]))
pos += 8
if ek[pos] != zsetStartMemSep {
err = errZScoreKey
return
}
pos++
member = ek[pos:]
return
}
func (db *DB) zSetItem(t *tx, key []byte, score int64, member []byte) (int64, error) {
2014-05-08 06:54:33 +04:00
if score <= MinScore || score >= MaxScore {
return 0, errScoreOverflow
}
var exists int64 = 0
2014-05-20 04:41:24 +04:00
ek := db.zEncodeSetKey(key, member)
if v, err := db.db.Get(ek); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else if v != nil {
exists = 1
2014-05-09 05:17:28 +04:00
if s, err := Int64(v, err); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else {
2014-05-20 04:41:24 +04:00
sk := db.zEncodeScoreKey(key, member, s)
2014-05-08 06:54:33 +04:00
t.Delete(sk)
}
}
2014-05-09 05:17:28 +04:00
t.Put(ek, PutInt64(score))
2014-05-08 06:54:33 +04:00
2014-05-20 04:41:24 +04:00
sk := db.zEncodeScoreKey(key, member, score)
2014-05-08 06:54:33 +04:00
t.Put(sk, []byte{})
return exists, nil
}
func (db *DB) zDelItem(t *tx, key []byte, member []byte, skipDelScore bool) (int64, error) {
2014-05-20 04:41:24 +04:00
ek := db.zEncodeSetKey(key, member)
if v, err := db.db.Get(ek); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else if v == nil {
//not exists
return 0, nil
} else {
//exists
if !skipDelScore {
//we must del score
2014-05-09 05:17:28 +04:00
if s, err := Int64(v, err); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else {
2014-05-20 04:41:24 +04:00
sk := db.zEncodeScoreKey(key, member, s)
2014-05-08 06:54:33 +04:00
t.Delete(sk)
}
}
}
t.Delete(ek)
return 1, nil
}
func (db *DB) zDelete(t *tx, key []byte) int64 {
delMembCnt, _ := db.zRemRange(t, key, MinScore, MaxScore, 0, -1)
// todo : log err
return delMembCnt
}
func (db *DB) zExpireAt(key []byte, when int64) (int64, error) {
t := db.zsetTx
t.Lock()
defer t.Unlock()
if zcnt, err := db.ZCard(key); err != nil || zcnt == 0 {
return 0, err
} else {
db.expireAt(t, ZSetType, key, when)
if err := t.Commit(); err != nil {
return 0, err
}
}
2014-05-08 06:54:33 +04:00
return 1, nil
}
2014-05-16 04:56:32 +04:00
func (db *DB) ZAdd(key []byte, args ...ScorePair) (int64, error) {
if len(args) == 0 {
return 0, nil
}
t := db.zsetTx
2014-05-08 06:54:33 +04:00
t.Lock()
defer t.Unlock()
var num int64 = 0
2014-05-16 04:56:32 +04:00
for i := 0; i < len(args); i++ {
score := args[i].Score
member := args[i].Member
2014-05-08 06:54:33 +04:00
2014-05-20 04:41:24 +04:00
if err := checkZSetKMSize(key, member); err != nil {
return 0, err
}
if n, err := db.zSetItem(t, key, score, member); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else if n == 0 {
//add new
num++
}
}
if _, err := db.zIncrSize(t, key, num); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
}
//todo add binlog
err := t.Commit()
return num, err
}
func (db *DB) zIncrSize(t *tx, key []byte, delta int64) (int64, error) {
2014-05-20 04:41:24 +04:00
sk := db.zEncodeSizeKey(key)
size, err := Int64(db.db.Get(sk))
2014-05-08 06:54:33 +04:00
if err != nil {
return 0, err
} else {
size += delta
if size <= 0 {
size = 0
t.Delete(sk)
db.rmExpire(t, ZSetType, key)
2014-05-08 06:54:33 +04:00
} else {
2014-05-09 05:17:28 +04:00
t.Put(sk, PutInt64(size))
2014-05-08 06:54:33 +04:00
}
}
return size, nil
}
func (db *DB) ZCard(key []byte) (int64, error) {
2014-05-20 04:41:24 +04:00
if err := checkKeySize(key); err != nil {
return 0, err
}
sk := db.zEncodeSizeKey(key)
return Int64(db.db.Get(sk))
2014-05-08 06:54:33 +04:00
}
2014-06-04 14:32:23 +04:00
func (db *DB) ZScore(key []byte, member []byte) (int64, error) {
2014-05-20 04:41:24 +04:00
if err := checkZSetKMSize(key, member); err != nil {
2014-06-04 14:32:23 +04:00
return InvalidScore, err
2014-05-20 04:41:24 +04:00
}
2014-06-04 14:32:23 +04:00
var score int64 = InvalidScore
2014-05-20 04:41:24 +04:00
2014-06-04 14:32:23 +04:00
k := db.zEncodeSetKey(key, member)
if v, err := db.db.Get(k); err != nil {
return InvalidScore, err
} else if v == nil {
return InvalidScore, ErrScoreMiss
2014-06-04 14:32:23 +04:00
} else {
if score, err = Int64(v, nil); err != nil {
return InvalidScore, err
}
2014-05-09 05:17:28 +04:00
}
2014-06-04 14:32:23 +04:00
return score, nil
2014-05-08 06:54:33 +04:00
}
2014-05-16 04:56:32 +04:00
func (db *DB) ZRem(key []byte, members ...[]byte) (int64, error) {
if len(members) == 0 {
return 0, nil
}
t := db.zsetTx
2014-05-08 06:54:33 +04:00
t.Lock()
defer t.Unlock()
var num int64 = 0
2014-05-16 04:56:32 +04:00
for i := 0; i < len(members); i++ {
2014-05-20 04:41:24 +04:00
if err := checkZSetKMSize(key, members[i]); err != nil {
return 0, err
}
if n, err := db.zDelItem(t, key, members[i], false); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else if n == 1 {
num++
}
}
if _, err := db.zIncrSize(t, key, -num); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
}
err := t.Commit()
return num, err
}
2014-06-04 14:32:23 +04:00
func (db *DB) ZIncrBy(key []byte, delta int64, member []byte) (int64, error) {
2014-05-20 04:41:24 +04:00
if err := checkZSetKMSize(key, member); err != nil {
2014-06-04 14:32:23 +04:00
return InvalidScore, err
2014-05-20 04:41:24 +04:00
}
t := db.zsetTx
2014-05-08 06:54:33 +04:00
t.Lock()
defer t.Unlock()
2014-05-20 04:41:24 +04:00
ek := db.zEncodeSetKey(key, member)
var oldScore int64 = 0
v, err := db.db.Get(ek)
2014-05-08 06:54:33 +04:00
if err != nil {
2014-06-04 14:32:23 +04:00
return InvalidScore, err
} else if v == nil {
db.zIncrSize(t, key, 1)
} else {
if oldScore, err = Int64(v, err); err != nil {
2014-06-04 14:32:23 +04:00
return InvalidScore, err
2014-05-08 06:54:33 +04:00
}
}
newScore := oldScore + delta
if newScore >= MaxScore || newScore <= MinScore {
return InvalidScore, errScoreOverflow
}
2014-05-08 06:54:33 +04:00
sk := db.zEncodeScoreKey(key, member, newScore)
2014-05-20 04:41:24 +04:00
t.Put(sk, []byte{})
t.Put(ek, PutInt64(newScore))
if v != nil {
// so as to update score, we must delete the old one
oldSk := db.zEncodeScoreKey(key, member, oldScore)
t.Delete(oldSk)
}
2014-05-08 06:54:33 +04:00
err = t.Commit()
return newScore, err
2014-05-08 06:54:33 +04:00
}
func (db *DB) ZCount(key []byte, min int64, max int64) (int64, error) {
2014-05-20 04:41:24 +04:00
if err := checkKeySize(key); err != nil {
return 0, err
}
minKey := db.zEncodeStartScoreKey(key, min)
maxKey := db.zEncodeStopScoreKey(key, max)
2014-05-08 06:54:33 +04:00
2014-07-25 13:58:00 +04:00
rangeType := store.RangeROpen
2014-05-08 06:54:33 +04:00
2014-06-19 13:19:40 +04:00
it := db.db.RangeLimitIterator(minKey, maxKey, rangeType, 0, -1)
2014-05-08 06:54:33 +04:00
var n int64 = 0
for ; it.Valid(); it.Next() {
n++
}
it.Close()
return n, nil
}
func (db *DB) zrank(key []byte, member []byte, reverse bool) (int64, error) {
2014-05-20 04:41:24 +04:00
if err := checkZSetKMSize(key, member); err != nil {
return 0, err
}
k := db.zEncodeSetKey(key, member)
2014-05-08 06:54:33 +04:00
it := db.db.NewIterator()
defer it.Close()
if v := it.Find(k); v == nil {
2014-05-08 06:54:33 +04:00
return -1, nil
} else {
if s, err := Int64(v, nil); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else {
2014-07-25 13:58:00 +04:00
var rit *store.RangeLimitIterator
2014-05-08 06:54:33 +04:00
2014-05-20 04:41:24 +04:00
sk := db.zEncodeScoreKey(key, member, s)
2014-05-08 06:54:33 +04:00
if !reverse {
2014-05-20 04:41:24 +04:00
minKey := db.zEncodeStartScoreKey(key, MinScore)
2014-07-25 13:58:00 +04:00
rit = store.NewRangeIterator(it, &store.Range{minKey, sk, store.RangeClose})
2014-05-08 06:54:33 +04:00
} else {
2014-05-20 04:41:24 +04:00
maxKey := db.zEncodeStopScoreKey(key, MaxScore)
2014-07-25 13:58:00 +04:00
rit = store.NewRevRangeIterator(it, &store.Range{sk, maxKey, store.RangeClose})
2014-05-08 06:54:33 +04:00
}
var lastKey []byte = nil
var n int64 = 0
for ; rit.Valid(); rit.Next() {
2014-05-08 06:54:33 +04:00
n++
lastKey = rit.BufKey(lastKey)
2014-05-08 06:54:33 +04:00
}
2014-05-20 04:41:24 +04:00
if _, m, _, err := db.zDecodeScoreKey(lastKey); err == nil && bytes.Equal(m, member) {
2014-05-08 06:54:33 +04:00
n--
return n, nil
}
}
}
return -1, nil
}
2014-07-25 13:58:00 +04:00
func (db *DB) zIterator(key []byte, min int64, max int64, offset int, count int, reverse bool) *store.RangeLimitIterator {
2014-05-20 04:41:24 +04:00
minKey := db.zEncodeStartScoreKey(key, min)
maxKey := db.zEncodeStopScoreKey(key, max)
2014-05-08 06:54:33 +04:00
if !reverse {
2014-07-25 13:58:00 +04:00
return db.db.RangeLimitIterator(minKey, maxKey, store.RangeClose, offset, count)
2014-05-08 06:54:33 +04:00
} else {
2014-07-25 13:58:00 +04:00
return db.db.RevRangeLimitIterator(minKey, maxKey, store.RangeClose, offset, count)
2014-05-08 06:54:33 +04:00
}
}
func (db *DB) zRemRange(t *tx, key []byte, min int64, max int64, offset int, count int) (int64, error) {
2014-05-20 04:41:24 +04:00
if len(key) > MaxKeySize {
2014-06-05 11:46:38 +04:00
return 0, errKeySize
2014-05-20 04:41:24 +04:00
}
it := db.zIterator(key, min, max, offset, count, false)
2014-05-08 06:54:33 +04:00
var num int64 = 0
for ; it.Valid(); it.Next() {
2014-07-07 11:36:45 +04:00
sk := it.RawKey()
_, m, _, err := db.zDecodeScoreKey(sk)
2014-05-08 06:54:33 +04:00
if err != nil {
continue
}
if n, err := db.zDelItem(t, key, m, true); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
} else if n == 1 {
num++
}
t.Delete(sk)
2014-05-08 06:54:33 +04:00
}
it.Close()
2014-05-08 06:54:33 +04:00
if _, err := db.zIncrSize(t, key, -num); err != nil {
2014-05-08 06:54:33 +04:00
return 0, err
}
return num, nil
2014-05-08 06:54:33 +04:00
}
func (db *DB) zRange(key []byte, min int64, max int64, offset int, count int, reverse bool) ([]ScorePair, error) {
2014-05-20 04:41:24 +04:00
if len(key) > MaxKeySize {
2014-06-05 11:46:38 +04:00
return nil, errKeySize
2014-05-20 04:41:24 +04:00
}
if offset < 0 {
return []ScorePair{}, nil
}
nv := 64
if count > 0 {
nv = count
}
v := make([]ScorePair, 0, nv)
2014-07-25 13:58:00 +04:00
var it *store.RangeLimitIterator
//if reverse and offset is 0, count < 0, we may use forward iterator then reverse
2014-07-25 13:58:00 +04:00
//because store iterator prev is slower than next
if !reverse || (offset == 0 && count < 0) {
it = db.zIterator(key, min, max, offset, count, false)
} else {
it = db.zIterator(key, min, max, offset, count, true)
}
2014-05-08 06:54:33 +04:00
for ; it.Valid(); it.Next() {
2014-05-20 04:41:24 +04:00
_, m, s, err := db.zDecodeScoreKey(it.Key())
2014-05-08 06:54:33 +04:00
//may be we will check key equal?
if err != nil {
continue
}
v = append(v, ScorePair{Member: m, Score: s})
2014-05-08 06:54:33 +04:00
}
it.Close()
2014-05-08 06:54:33 +04:00
if reverse && (offset == 0 && count < 0) {
for i, j := 0, len(v)-1; i < j; i, j = i+1, j-1 {
v[i], v[j] = v[j], v[i]
}
}
2014-05-08 06:54:33 +04:00
return v, nil
}
2014-05-12 11:08:59 +04:00
func (db *DB) zParseLimit(key []byte, start int, stop int) (offset int, count int, err error) {
if start < 0 || stop < 0 {
//refer redis implementation
var size int64
size, err = db.ZCard(key)
if err != nil {
return
}
llen := int(size)
if start < 0 {
start = llen + start
}
if stop < 0 {
stop = llen + stop
}
if start < 0 {
start = 0
}
if start >= llen {
offset = -1
return
}
}
if start > stop {
offset = -1
return
}
offset = start
count = (stop - start) + 1
return
}
func (db *DB) ZClear(key []byte) (int64, error) {
t := db.zsetTx
t.Lock()
defer t.Unlock()
rmCnt, err := db.zRemRange(t, key, MinScore, MaxScore, 0, -1)
if err == nil {
err = t.Commit()
}
return rmCnt, err
}
func (db *DB) ZMclear(keys ...[]byte) (int64, error) {
t := db.zsetTx
t.Lock()
defer t.Unlock()
for _, key := range keys {
if _, err := db.zRemRange(t, key, MinScore, MaxScore, 0, -1); err != nil {
return 0, err
}
}
err := t.Commit()
return int64(len(keys)), err
}
func (db *DB) ZRange(key []byte, start int, stop int) ([]ScorePair, error) {
return db.ZRangeGeneric(key, start, stop, false)
}
//min and max must be inclusive
//if no limit, set offset = 0 and count = -1
func (db *DB) ZRangeByScore(key []byte, min int64, max int64,
offset int, count int) ([]ScorePair, error) {
return db.ZRangeByScoreGeneric(key, min, max, offset, count, false)
}
func (db *DB) ZRank(key []byte, member []byte) (int64, error) {
return db.zrank(key, member, false)
}
func (db *DB) ZRemRangeByRank(key []byte, start int, stop int) (int64, error) {
offset, count, err := db.zParseLimit(key, start, stop)
if err != nil {
return 0, err
}
var rmCnt int64
t := db.zsetTx
t.Lock()
defer t.Unlock()
rmCnt, err = db.zRemRange(t, key, MinScore, MaxScore, offset, count)
if err == nil {
err = t.Commit()
}
return rmCnt, err
}
//min and max must be inclusive
func (db *DB) ZRemRangeByScore(key []byte, min int64, max int64) (int64, error) {
t := db.zsetTx
t.Lock()
defer t.Unlock()
rmCnt, err := db.zRemRange(t, key, min, max, 0, -1)
if err == nil {
err = t.Commit()
}
return rmCnt, err
}
func (db *DB) ZRevRange(key []byte, start int, stop int) ([]ScorePair, error) {
return db.ZRangeGeneric(key, start, stop, true)
}
func (db *DB) ZRevRank(key []byte, member []byte) (int64, error) {
return db.zrank(key, member, true)
}
//min and max must be inclusive
//if no limit, set offset = 0 and count = -1
func (db *DB) ZRevRangeByScore(key []byte, min int64, max int64, offset int, count int) ([]ScorePair, error) {
return db.ZRangeByScoreGeneric(key, min, max, offset, count, true)
}
func (db *DB) ZRangeGeneric(key []byte, start int, stop int, reverse bool) ([]ScorePair, error) {
offset, count, err := db.zParseLimit(key, start, stop)
if err != nil {
return nil, err
}
return db.zRange(key, MinScore, MaxScore, offset, count, reverse)
}
//min and max must be inclusive
//if no limit, set offset = 0 and count = -1
func (db *DB) ZRangeByScoreGeneric(key []byte, min int64, max int64,
offset int, count int, reverse bool) ([]ScorePair, error) {
return db.zRange(key, min, max, offset, count, reverse)
2014-05-12 11:08:59 +04:00
}
func (db *DB) zFlush() (drop int64, err error) {
t := db.zsetTx
t.Lock()
defer t.Unlock()
minKey := make([]byte, 2)
minKey[0] = db.index
minKey[1] = ZSetType
maxKey := make([]byte, 2)
maxKey[0] = db.index
maxKey[1] = ZScoreType + 1
2014-07-25 13:58:00 +04:00
it := db.db.RangeLimitIterator(minKey, maxKey, store.RangeROpen, 0, -1)
defer it.Close()
for ; it.Valid(); it.Next() {
t.Delete(it.RawKey())
drop++
if drop&1023 == 0 {
if err = t.Commit(); err != nil {
return
}
}
}
db.expFlush(t, ZSetType)
err = t.Commit()
return
}
2014-05-25 10:01:35 +04:00
func (db *DB) ZScan(key []byte, member []byte, count int, inclusive bool) ([]ScorePair, error) {
var minKey []byte
if member != nil {
if err := checkZSetKMSize(key, member); err != nil {
return nil, err
}
minKey = db.zEncodeSetKey(key, member)
} else {
minKey = db.zEncodeStartSetKey(key)
}
maxKey := db.zEncodeStopSetKey(key)
if count <= 0 {
count = defaultScanCount
}
v := make([]ScorePair, 0, 2*count)
2014-07-25 13:58:00 +04:00
rangeType := store.RangeROpen
2014-05-25 10:01:35 +04:00
if !inclusive {
2014-07-25 13:58:00 +04:00
rangeType = store.RangeOpen
2014-05-25 10:01:35 +04:00
}
2014-06-19 13:19:40 +04:00
it := db.db.RangeLimitIterator(minKey, maxKey, rangeType, 0, count)
2014-05-25 10:01:35 +04:00
for ; it.Valid(); it.Next() {
if _, m, err := db.zDecodeSetKey(it.Key()); err != nil {
continue
} else {
score, _ := Int64(it.Value(), nil)
v = append(v, ScorePair{Member: m, Score: score})
}
}
it.Close()
2014-05-25 10:01:35 +04:00
return v, nil
}
func (db *DB) ZExpire(key []byte, duration int64) (int64, error) {
if duration <= 0 {
return 0, errExpireValue
}
return db.zExpireAt(key, time.Now().Unix()+duration)
}
func (db *DB) ZExpireAt(key []byte, when int64) (int64, error) {
if when <= time.Now().Unix() {
return 0, errExpireValue
}
return db.zExpireAt(key, when)
}
func (db *DB) ZTTL(key []byte) (int64, error) {
if err := checkKeySize(key); err != nil {
return -1, err
}
return db.ttl(ZSetType, key)
}
2014-06-23 07:12:20 +04:00
func (db *DB) ZPersist(key []byte) (int64, error) {
if err := checkKeySize(key); err != nil {
return 0, err
}
t := db.zsetTx
t.Lock()
defer t.Unlock()
n, err := db.rmExpire(t, ZSetType, key)
2014-06-23 07:12:20 +04:00
if err != nil {
return 0, err
}
err = t.Commit()
return n, err
}
func getAggregateFunc(aggregate byte) func(int64, int64) int64 {
switch aggregate {
case AggregateSum:
return func(a int64, b int64) int64 {
return a + b
}
case AggregateMax:
return func(a int64, b int64) int64 {
if a > b {
return a
}
return b
}
case AggregateMin:
return func(a int64, b int64) int64 {
if a > b {
return b
}
return a
}
}
return nil
}
func (db *DB) ZUnionStore(destKey []byte, srcKeys [][]byte, weights []int64, aggregate byte) (int64, error) {
var destMap = map[string]int64{}
aggregateFunc := getAggregateFunc(aggregate)
if aggregateFunc == nil {
return 0, errInvalidAggregate
}
if len(srcKeys) < 1 {
return 0, errInvalidSrcKeyNum
}
if weights != nil {
if len(srcKeys) != len(weights) {
return 0, errInvalidWeightNum
}
} else {
weights = make([]int64, len(srcKeys))
for i := 0; i < len(weights); i++ {
weights[i] = 1
}
}
for i, key := range srcKeys {
scorePairs, err := db.ZRange(key, 0, -1)
if err != nil {
return 0, err
}
for _, pair := range scorePairs {
if score, ok := destMap[String(pair.Member)]; !ok {
destMap[String(pair.Member)] = pair.Score
} else {
destMap[String(pair.Member)] = aggregateFunc(score, pair.Score*weights[i])
}
}
}
t := db.zsetTx
t.Lock()
defer t.Unlock()
db.zDelete(t, destKey)
var num int64 = 0
for member, score := range destMap {
if err := checkZSetKMSize(destKey, []byte(member)); err != nil {
return 0, err
}
if n, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
return 0, err
} else if n == 0 {
//add new
num++
}
}
if _, err := db.zIncrSize(t, destKey, num); err != nil {
return 0, err
}
//todo add binlog
2014-08-13 05:57:13 +04:00
if err := t.Commit(); err != nil {
return 0, err
}
return int64(len(destMap)), nil
}
func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, aggregate byte) (int64, error) {
var destMap = map[string]int64{}
aggregateFunc := getAggregateFunc(aggregate)
if aggregateFunc == nil {
return 0, errInvalidAggregate
}
if len(srcKeys) < 1 {
return 0, errInvalidSrcKeyNum
}
if weights != nil {
if len(srcKeys) != len(weights) {
return 0, errInvalidWeightNum
}
} else {
weights = make([]int64, len(srcKeys))
for i := 0; i < len(weights); i++ {
weights[i] = 1
}
}
var keptMembers [][]byte
for i, key := range srcKeys {
scorePairs, err := db.ZRange(key, 0, -1)
if err != nil {
return 0, err
}
for _, pair := range scorePairs {
if score, ok := destMap[String(pair.Member)]; !ok {
destMap[String(pair.Member)] = pair.Score
} else {
keptMembers = append(keptMembers, pair.Member)
destMap[String(pair.Member)] = aggregateFunc(score, pair.Score*weights[i])
}
}
}
t := db.zsetTx
t.Lock()
defer t.Unlock()
db.zDelete(t, destKey)
var num int64 = 0
for _, member := range keptMembers {
score := destMap[String(member)]
if err := checkZSetKMSize(destKey, member); err != nil {
return 0, err
}
if n, err := db.zSetItem(t, destKey, score, member); err != nil {
return 0, err
} else if n == 0 {
//add new
num++
}
}
2014-08-13 05:57:13 +04:00
2014-08-13 06:26:35 +04:00
if _, err := db.zIncrSize(t, destKey, num); err != nil {
return 0, err
}
2014-08-13 05:57:13 +04:00
//todo add binlog
if err := t.Commit(); err != nil {
return 0, err
}
return int64(len(keptMembers)), nil
}