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[hack.String(pair.Member)]; !ok {
destMap[hack.String(pair.Member)] = pair.Score * weights[i]
} else {
destMap[hack.String(pair.Member)] = aggregateFunc(score, pair.Score*weights[i])
}
}
}
t := db.zsetBatch
t.Lock()
defer t.Unlock()
db.zDelete(t, destKey)
for member, score := range destMap {
if err := checkZSetKMSize(destKey, []byte(member)); err != nil {
return 0, err
}
if _, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
return 0, err
}
}
var n = int64(len(destMap))
sk := db.zEncodeSizeKey(destKey)
t.Put(sk, PutInt64(n))
if err := t.Commit(); err != nil {
return 0, err
}
return n, nil
}
func (w *httpWriter) writeFVPairArray(lst []ledis.FVPair) {
m := make(map[string]string)
for _, elem := range lst {
m[hack.String(elem.Field)] = hack.String(elem.Value)
}
w.genericWrite(m)
}
//inner command, only for replication
//REPLCONF <option> <value> <option> <value> ...
func replconfCommand(c *client) error {
args := c.args
if len(args)%2 != 0 {
return ErrCmdParams
}
if !c.app.ldb.ReplicationUsed() {
return ledis.ErrRplNotSupport
}
//now only support "listening-port"
for i := 0; i < len(args); i += 2 {
switch strings.ToLower(hack.String(args[i])) {
case "listening-port":
var host string
var err error
if _, err = num.ParseUint16(hack.String(args[i+1])); err != nil {
return err
}
if host, _, err = net.SplitHostPort(c.remoteAddr); err != nil {
return err
} else {
c.slaveListeningAddr = net.JoinHostPort(host, hack.String(args[i+1]))
}
c.app.addSlave(c)
default:
return ErrSyntax
}
}
c.resp.writeStatus(OK)
return nil
}
// XSELECT db THEN command
func (c *respClient) handleXSelectCmd() error {
if len(c.args) <= 2 {
// invalid command format
return fmt.Errorf("invalid format for XSELECT, must XSELECT db THEN your command")
}
if hack.String(upperSlice(c.args[1])) != "THEN" {
// invalid command format, just resturn here
return fmt.Errorf("invalid format for XSELECT, must XSELECT db THEN your command")
}
index, err := strconv.Atoi(hack.String(c.args[0]))
if err != nil {
return fmt.Errorf("invalid db for XSELECT, err %v", err)
}
db, err := c.app.ldb.Select(index)
if err != nil {
return fmt.Errorf("invalid db for XSELECT, err %v", err)
}
c.db = db
c.cmd = hack.String(lowerSlice(c.args[2]))
c.args = c.args[3:]
return nil
}
func zparseRange(c *client, a1 []byte, a2 []byte) (start int, stop int, err error) {
if start, err = strconv.Atoi(hack.String(a1)); err != nil {
return
}
if stop, err = strconv.Atoi(hack.String(a2)); err != nil {
return
}
return
}
func (db *DB) sInterGeneric(keys ...[]byte) ([][]byte, error) {
destMap := make(map[string]bool)
members, err := db.SMembers(keys[0])
if err != nil {
return nil, err
}
for _, m := range members {
destMap[hack.String(m)] = true
}
for _, key := range keys[1:] {
if err := checkKeySize(key); err != nil {
return nil, err
}
members, err := db.SMembers(key)
if err != nil {
return nil, err
} else if len(members) == 0 {
return nil, err
}
tempMap := make(map[string]bool)
for _, member := range members {
if err := checkKeySize(member); err != nil {
return nil, err
}
if _, ok := destMap[hack.String(member)]; ok {
tempMap[hack.String(member)] = true //mark this item as selected
}
}
destMap = tempMap //reduce the size of the result set
if len(destMap) == 0 {
return nil, nil
}
}
slice := make([][]byte, len(destMap))
idx := 0
for k, v := range destMap {
if !v {
continue
}
slice[idx] = []byte(k)
idx++
}
return slice, nil
}
func (w *luaWriter) writeFVPairArray(lst []ledis.FVPair) {
if lst == nil {
w.l.PushBoolean(false)
return
}
w.l.CreateTable(len(lst)*2, 0)
for i, v := range lst {
w.l.PushString(hack.String(v.Field))
w.l.RawSeti(-2, 2*i+1)
w.l.PushString(hack.String(v.Value))
w.l.RawSeti(-2, 2*i+2)
}
}
func zrangeGeneric(c *client, reverse bool) error {
args := c.args
if len(args) < 3 {
return ErrCmdParams
}
key := args[0]
start, stop, err := zparseRange(c, args[1], args[2])
if err != nil {
return ErrValue
}
args = args[3:]
var withScores bool = false
if len(args) > 0 {
if len(args) != 1 {
return ErrCmdParams
}
if strings.ToLower(hack.String(args[0])) == "withscores" {
withScores = true
} else {
return ErrSyntax
}
}
if datas, err := c.db.ZRangeGeneric(key, start, stop, reverse); err != nil {
return err
} else {
c.resp.writeScorePairArray(datas, withScores)
}
return nil
}
func listAdaptor(db *DB) *adaptor {
adp := new(adaptor)
adp.showIdent = func() string {
return "list-adptor"
}
adp.set = func(k []byte, v []byte) (int64, error) {
eles := make([][]byte, 0)
for i := 0; i < 3; i++ {
e := []byte(hack.String(v) + fmt.Sprintf("_%d", i))
eles = append(eles, e)
}
if n, err := db.LPush(k, eles...); err != nil {
return 0, err
} else {
return n, nil
}
}
adp.exists = func(k []byte) (int64, error) {
if llen, err := db.LLen(k); err != nil || llen <= 0 {
return 0, err
} else {
return 1, nil
}
}
adp.del = db.LClear
adp.expire = db.LExpire
adp.expireAt = db.LExpireAt
adp.ttl = db.LTTL
return adp
}
func (w *httpWriter) writeScorePairArray(lst []ledis.ScorePair, withScores bool) {
var arr []string
if withScores {
arr = make([]string, 2*len(lst))
for i, data := range lst {
arr[2*i] = hack.String(data.Member)
arr[2*i+1] = strconv.FormatInt(data.Score, 10)
}
} else {
arr = make([]string, len(lst))
for i, data := range lst {
arr[i] = hack.String(data.Member)
}
}
w.genericWrite(arr)
}
func scriptLoadCommand(c *client) error {
s := c.app.script
l := s.l
if len(c.args) != 2 {
return ErrCmdParams
}
h := sha1.Sum(c.args[1])
key := hex.EncodeToString(h[0:20])
if r := l.LoadString(hack.String(c.args[1])); r != 0 {
err := fmt.Errorf("%s", l.ToString(-1))
l.Pop(1)
return err
} else {
l.PushValue(-1)
l.SetGlobal(key)
s.chunks[key] = struct{}{}
}
c.resp.writeBulk(hack.Slice(key))
return nil
}
func (w *httpWriter) writeBulk(b []byte) {
if b == nil {
w.genericWrite(nil)
} else {
w.genericWrite(hack.String(b))
}
}
func (w *luaWriter) writeBulk(b []byte) {
if b == nil {
w.l.PushBoolean(false)
} else {
w.l.PushString(hack.String(b))
}
}
func selectCommand(c *client) error {
if len(c.args) != 1 {
return ErrCmdParams
}
if index, err := strconv.Atoi(hack.String(c.args[0])); err != nil {
return err
} else {
// if c.db.IsInMulti() {
// if err := c.script.Select(index); err != nil {
// return err
// } else {
// c.db = c.script.DB
// }
// } else {
// if db, err := c.ldb.Select(index); err != nil {
// return err
// } else {
// c.db = db
// }
// }
if db, err := c.ldb.Select(index); err != nil {
return err
} else {
c.db = db
}
c.resp.writeStatus(OK)
}
return nil
}
func scriptCommand(c *client) error {
s := c.app.script
l := s.l
s.Lock()
base := l.GetTop()
defer func() {
l.SetTop(base)
s.Unlock()
}()
args := c.args
if len(args) < 1 {
return ErrCmdParams
}
switch strings.ToLower(hack.String(args[0])) {
case "load":
return scriptLoadCommand(c)
case "exists":
return scriptExistsCommand(c)
case "flush":
return scriptFlushCommand(c)
default:
return fmt.Errorf("invalid script %s", args[0])
}
return nil
}
func (db *DB) sUnionGeneric(keys ...[]byte) ([][]byte, error) {
dstMap := make(map[string]bool)
for _, key := range keys {
if err := checkKeySize(key); err != nil {
return nil, err
}
members, err := db.SMembers(key)
if err != nil {
return nil, err
}
for _, member := range members {
dstMap[hack.String(member)] = true
}
}
slice := make([][]byte, len(dstMap))
idx := 0
for k, v := range dstMap {
if !v {
continue
}
slice[idx] = []byte(k)
idx++
}
return slice, nil
}
func parseScanArgs(args [][]byte) (cursor []byte, match string, count int, desc bool, err error) {
cursor = args[0]
args = args[1:]
count = 10
desc = false
for i := 0; i < len(args); {
switch strings.ToUpper(hack.String(args[i])) {
case "MATCH":
if i+1 >= len(args) {
err = ErrCmdParams
return
}
match = hack.String(args[i+1])
i++
case "COUNT":
if i+1 >= len(args) {
err = ErrCmdParams
return
}
count, err = strconv.Atoi(hack.String(args[i+1]))
if err != nil {
return
}
i++
case "ASC":
desc = false
case "DESC":
desc = true
default:
err = fmt.Errorf("invalid argument %s", args[i])
return
}
i++
}
return
}
func StrUint64(v []byte, err error) (uint64, error) {
if err != nil {
return 0, err
} else if v == nil {
return 0, nil
} else {
return strconv.ParseUint(hack.String(v), 10, 64)
}
}
func luaSetGlobalArray(l *lua.State, name string, ay [][]byte) {
l.NewTable()
for i := 0; i < len(ay); i++ {
l.PushString(hack.String(ay[i]))
l.RawSeti(-2, i+1)
}
l.SetGlobal(name)
}
func StrInt8(v []byte, err error) (int8, error) {
if err != nil {
return 0, err
} else if v == nil {
return 0, nil
} else {
res, err := strconv.ParseInt(hack.String(v), 10, 8)
return int8(res), err
}
}
func (db *DB) sDiffGeneric(keys ...[]byte) ([][]byte, error) {
destMap := make(map[string]bool)
members, err := db.SMembers(keys[0])
if err != nil {
return nil, err
}
for _, m := range members {
destMap[hack.String(m)] = true
}
for _, k := range keys[1:] {
members, err := db.SMembers(k)
if err != nil {
return nil, err
}
for _, m := range members {
if _, ok := destMap[hack.String(m)]; !ok {
continue
} else if ok {
delete(destMap, hack.String(m))
}
}
// O - A = O, O is zero set.
if len(destMap) == 0 {
return nil, nil
}
}
slice := make([][]byte, len(destMap))
idx := 0
for k, v := range destMap {
if !v {
continue
}
slice[idx] = []byte(k)
idx++
}
return slice, nil
}
func (w *httpWriter) writeSliceArray(lst [][]byte) {
arr := make([]interface{}, len(lst))
for i, elem := range lst {
if elem == nil {
arr[i] = nil
} else {
arr[i] = hack.String(elem)
}
}
w.genericWrite(arr)
}
func zparseMemberRange(minBuf []byte, maxBuf []byte) (min []byte, max []byte, rangeType uint8, err error) {
rangeType = store.RangeClose
if strings.ToLower(hack.String(minBuf)) == "-" {
min = nil
} else {
if len(minBuf) == 0 {
err = ErrCmdParams
return
}
if minBuf[0] == '(' {
rangeType |= store.RangeLOpen
min = minBuf[1:]
} else if minBuf[0] == '[' {
min = minBuf[1:]
} else {
err = ErrCmdParams
return
}
}
if strings.ToLower(hack.String(maxBuf)) == "+" {
max = nil
} else {
if len(maxBuf) == 0 {
err = ErrCmdParams
return
}
if maxBuf[0] == '(' {
rangeType |= store.RangeROpen
max = maxBuf[1:]
} else if maxBuf[0] == '[' {
max = maxBuf[1:]
} else {
err = ErrCmdParams
return
}
}
return
}
func (w *luaWriter) writeSliceArray(lst [][]byte) {
if lst == nil {
w.l.PushBoolean(false)
return
}
w.l.CreateTable(len(lst), 0)
for i, v := range lst {
w.l.PushString(hack.String(v))
w.l.RawSeti(-2, i+1)
}
}
func (m *migrateKeyLocker) Lock(key []byte) bool {
m.m.Lock()
defer m.m.Unlock()
k := hack.String(key)
_, ok := m.locks[k]
if ok {
return false
}
m.locks[k] = struct{}{}
return true
}
func hashAdaptor(db *DB) *adaptor {
adp := new(adaptor)
adp.showIdent = func() string {
return "hash-adptor"
}
adp.set = func(k []byte, v []byte) (int64, error) {
datas := make([]FVPair, 0)
for i := 0; i < 3; i++ {
suffix := fmt.Sprintf("_%d", i)
pair := FVPair{
Field: []byte(hack.String(k) + suffix),
Value: []byte(hack.String(v) + suffix)}
datas = append(datas, pair)
}
if err := db.HMset(k, datas...); err != nil {
return 0, err
} else {
return int64(len(datas)), nil
}
}
adp.exists = func(k []byte) (int64, error) {
if hlen, err := db.HLen(k); err != nil || hlen <= 0 {
return 0, err
} else {
return 1, nil
}
}
adp.del = db.HClear
adp.expire = db.HExpire
adp.expireAt = db.HExpireAt
adp.ttl = db.HTTL
return adp
}
// XSCAN type cursor [MATCH match] [COUNT count] [ASC|DESC]
func xscanCommand(c *client) error {
args := c.args
if len(args) < 2 {
return ErrCmdParams
}
var dataType ledis.DataType
switch strings.ToUpper(hack.String(args[0])) {
case "KV":
dataType = ledis.KV
case "HASH":
dataType = ledis.HASH
case "LIST":
dataType = ledis.LIST
case "SET":
dataType = ledis.SET
case "ZSET":
dataType = ledis.ZSET
default:
return fmt.Errorf("invalid key type %s", args[0])
}
cursor, match, count, desc, err := parseScanArgs(args[1:])
if err != nil {
return err
}
var ay [][]byte
if !desc {
ay, err = c.db.Scan(dataType, cursor, count, false, match)
} else {
ay, err = c.db.RevScan(dataType, cursor, count, false, match)
}
if err != nil {
return err
}
data := make([]interface{}, 2)
if len(ay) < count {
data[0] = []byte("")
} else {
data[0] = ay[len(ay)-1]
}
data[1] = ay
c.resp.writeArray(data)
return nil
}
func infoCommand(c *client) error {
if len(c.args) > 1 {
return ErrCmdParams
}
var section string
if len(c.args) == 1 {
section = strings.ToLower(hack.String(c.args[0]))
}
buf := c.app.info.Dump(section)
c.resp.writeBulk(buf)
return nil
}
func (w *luaWriter) writeScorePairArray(lst []ledis.ScorePair, withScores bool) {
if lst == nil {
w.l.PushBoolean(false)
return
}
if withScores {
w.l.CreateTable(len(lst)*2, 0)
for i, v := range lst {
w.l.PushString(hack.String(v.Member))
w.l.RawSeti(-2, 2*i+1)
w.l.PushString(hack.String(num.FormatInt64ToSlice(v.Score)))
w.l.RawSeti(-2, 2*i+2)
}
} else {
w.l.CreateTable(len(lst), 0)
for i, v := range lst {
w.l.PushString(hack.String(v.Member))
w.l.RawSeti(-2, i+1)
}
}
}
func slaveofCommand(c *client) error {
args := c.args
if len(args) != 2 && len(args) != 3 {
return ErrCmdParams
}
masterAddr := ""
restart := false
readonly := false
if strings.ToLower(hack.String(args[0])) == "no" &&
strings.ToLower(hack.String(args[1])) == "one" {
//stop replication, use master = ""
if len(args) == 3 && strings.ToLower(hack.String(args[2])) == "readonly" {
readonly = true
}
} else {
if _, err := strconv.ParseInt(hack.String(args[1]), 10, 16); err != nil {
return err
}
masterAddr = fmt.Sprintf("%s:%s", args[0], args[1])
if len(args) == 3 && strings.ToLower(hack.String(args[2])) == "restart" {
restart = true
}
}
if err := c.app.slaveof(masterAddr, restart, readonly); err != nil {
return err
}
c.resp.writeStatus(OK)
return nil
}