egg_models/utils/cache/redis_pool_cluster.go

package cache

import (
	"strconv"
	"time"

	"github.com/go-redis/redis"
)

type RedisClusterPool struct {
	client *redis.ClusterClient
}

func NewRedisClusterPool(addrs []string) (*RedisClusterPool, error) {
	opt := &redis.ClusterOptions{
		Addrs:        addrs,
		PoolSize:     512,
		PoolTimeout:  10 * time.Second,
		IdleTimeout:  10 * time.Second,
		DialTimeout:  10 * time.Second,
		ReadTimeout:  3 * time.Second,
		WriteTimeout: 3 * time.Second,
	}
	c := redis.NewClusterClient(opt)
	if err := c.Ping().Err(); err != nil {
		return nil, err
	}
	return &RedisClusterPool{client: c}, nil
}

func (p *RedisClusterPool) Get(key string) (interface{}, error) {
	res, err := p.client.Get(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return []byte(res), nil
}
func (p *RedisClusterPool) Set(key string, value interface{}) error {
	err := p.client.Set(key, value, 0).Err()
	return convertError(err)
}
func (p *RedisClusterPool) GetSet(key string, value interface{}) (interface{}, error) {
	res, err := p.client.GetSet(key, value).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return []byte(res), nil
}
func (p *RedisClusterPool) SetNx(key string, value interface{}) (int64, error) {
	res, err := p.client.SetNX(key, value, 0).Result()
	if err != nil {
		return 0, convertError(err)
	}
	if res {
		return 1, nil
	}
	return 0, nil
}
func (p *RedisClusterPool) SetEx(key string, value interface{}, timeout int64) error {
	_, err := p.client.Set(key, value, time.Duration(timeout)*time.Second).Result()
	if err != nil {
		return convertError(err)
	}
	return nil
}

// nil表示成功，ErrNil表示数据库内已经存在这个key，其他表示数据库发生错误
func (p *RedisClusterPool) SetNxEx(key string, value interface{}, timeout int64) error {
	res, err := p.client.SetNX(key, value, time.Duration(timeout)*time.Second).Result()
	if err != nil {
		return convertError(err)
	}
	if res {
		return nil
	}
	return ErrNil
}
func (p *RedisClusterPool) MGet(keys ...string) ([]interface{}, error) {
	res, err := p.client.MGet(keys...).Result()
	return res, convertError(err)
}

// 为确保多个key映射到同一个slot，每个key最好加上hash tag，如:{test}
func (p *RedisClusterPool) MSet(kvs map[string]interface{}) error {
	pairs := make([]string, 0, len(kvs)*2)
	for k, v := range kvs {
		val, err := String(v, nil)
		if err != nil {
			return err
		}
		pairs = append(pairs, k, val)
	}
	return convertError(p.client.MSet(pairs).Err())
}

// 为确保多个key映射到同一个slot，每个key最好加上hash tag，如:{test}
func (p *RedisClusterPool) MSetNX(kvs map[string]interface{}) (bool, error) {
	pairs := make([]string, 0, len(kvs)*2)
	for k, v := range kvs {
		val, err := String(v, nil)
		if err != nil {
			return false, err
		}
		pairs = append(pairs, k, val)
	}
	res, err := p.client.MSetNX(pairs).Result()
	return res, convertError(err)
}
func (p *RedisClusterPool) ExpireAt(key string, timestamp int64) (int64, error) {
	res, err := p.client.ExpireAt(key, time.Unix(timestamp, 0)).Result()
	if err != nil {
		return 0, convertError(err)
	}
	if res {
		return 1, nil
	}
	return 0, nil
}
func (p *RedisClusterPool) Del(keys ...string) (int64, error) {
	args := make([]interface{}, 0, len(keys))
	for _, key := range keys {
		args = append(args, key)
	}
	res, err := p.client.Del(keys...).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) Incr(key string) (int64, error) {
	res, err := p.client.Incr(key).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) IncrBy(key string, delta int64) (int64, error) {
	res, err := p.client.IncrBy(key, delta).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) Expire(key string, duration int64) (int64, error) {
	res, err := p.client.Expire(key, time.Duration(duration)*time.Second).Result()
	if err != nil {
		return 0, convertError(err)
	}
	if res {
		return 1, nil
	}
	return 0, nil
}
func (p *RedisClusterPool) Exists(key string) (bool, error) { // todo (bool, error)
	res, err := p.client.Exists(key).Result()
	if err != nil {
		return false, convertError(err)
	}
	if res > 0 {
		return true, nil
	}
	return false, nil
}
func (p *RedisClusterPool) HGet(key string, field string) (interface{}, error) {
	res, err := p.client.HGet(key, field).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return []byte(res), nil
}
func (p *RedisClusterPool) HLen(key string) (int64, error) {
	res, err := p.client.HLen(key).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) HSet(key string, field string, val interface{}) error {
	value, err := String(val, nil)
	if err != nil && err != ErrNil {
		return err
	}
	_, err = p.client.HSet(key, field, value).Result()
	if err != nil {
		return convertError(err)
	}
	return nil
}
func (p *RedisClusterPool) HDel(key string, fields ...string) (int64, error) {
	args := make([]interface{}, 0, len(fields)+1)
	args = append(args, key)
	for _, field := range fields {
		args = append(args, field)
	}
	res, err := p.client.HDel(key, fields...).Result()
	if err != nil {
		return 0, convertError(err)
	}
	return res, nil
}

func (p *RedisClusterPool) HMGet(key string, fields ...string) (interface{}, error) {
	args := make([]interface{}, 0, len(fields)+1)
	args = append(args, key)
	for _, field := range fields {
		args = append(args, field)
	}
	if len(fields) == 0 {
		return nil, ErrNil
	}
	res, err := p.client.HMGet(key, fields...).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) HMSet(key string, kvs ...interface{}) error {
	if len(kvs) == 0 {
		return nil
	}
	if len(kvs)%2 != 0 {
		return ErrWrongArgsNum
	}
	var err error
	v := map[string]interface{}{} // todo change
	v["field"], err = String(kvs[0], nil)
	if err != nil && err != ErrNil {
		return err
	}
	v["value"], err = String(kvs[1], nil)
	if err != nil && err != ErrNil {
		return err
	}
	pairs := make([]string, 0, len(kvs)-2)
	if len(kvs) > 2 {
		for _, kv := range kvs[2:] {
			kvString, err := String(kv, nil)
			if err != nil && err != ErrNil {
				return err
			}
			pairs = append(pairs, kvString)
		}
	}
	v["paris"] = pairs
	_, err = p.client.HMSet(key, v).Result()
	if err != nil {
		return convertError(err)
	}
	return nil
}

func (p *RedisClusterPool) HKeys(key string) ([]string, error) {
	res, err := p.client.HKeys(key).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) HVals(key string) ([]interface{}, error) {
	res, err := p.client.HVals(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	rs := make([]interface{}, 0, len(res))
	for _, res := range res {
		rs = append(rs, res)
	}
	return rs, nil
}
func (p *RedisClusterPool) HGetAll(key string) (map[string]string, error) {
	vals, err := p.client.HGetAll(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return vals, nil
}
func (p *RedisClusterPool) HIncrBy(key, field string, delta int64) (int64, error) {
	res, err := p.client.HIncrBy(key, field, delta).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) ZAdd(key string, kvs ...interface{}) (int64, error) {
	args := make([]interface{}, 0, len(kvs)+1)
	args = append(args, key)
	args = append(args, kvs...)
	if len(kvs) == 0 {
		return 0, nil
	}
	if len(kvs)%2 != 0 {
		return 0, ErrWrongArgsNum
	}
	zs := make([]redis.Z, len(kvs)/2)
	for i := 0; i < len(kvs); i += 2 {
		idx := i / 2
		score, err := Float64(kvs[i], nil)
		if err != nil && err != ErrNil {
			return 0, err
		}
		zs[idx].Score = score
		zs[idx].Member = kvs[i+1]
	}
	res, err := p.client.ZAdd(key, zs...).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) ZRem(key string, members ...string) (int64, error) {
	args := make([]interface{}, 0, len(members))
	args = append(args, key)
	for _, member := range members {
		args = append(args, member)
	}
	res, err := p.client.ZRem(key, members).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, err
}

func (p *RedisClusterPool) ZRange(key string, min, max int64, withScores bool) (interface{}, error) {
	res := make([]interface{}, 0)
	if withScores {
		zs, err := p.client.ZRangeWithScores(key, min, max).Result()
		if err != nil {
			return nil, convertError(err)
		}
		for _, z := range zs {
			res = append(res, z.Member, strconv.FormatFloat(z.Score, 'f', -1, 64))
		}
	} else {
		ms, err := p.client.ZRange(key, min, max).Result()
		if err != nil {
			return nil, convertError(err)
		}
		for _, m := range ms {
			res = append(res, m)
		}
	}
	return res, nil
}
func (p *RedisClusterPool) ZRangeByScoreWithScore(key string, min, max int64) (map[string]int64, error) {
	opt := new(redis.ZRangeBy)
	opt.Min = strconv.FormatInt(int64(min), 10)
	opt.Max = strconv.FormatInt(int64(max), 10)
	opt.Count = -1
	opt.Offset = 0
	vals, err := p.client.ZRangeByScoreWithScores(key, *opt).Result()
	if err != nil {
		return nil, convertError(err)
	}
	res := make(map[string]int64, len(vals))
	for _, val := range vals {
		key, err := String(val.Member, nil)
		if err != nil && err != ErrNil {
			return nil, err
		}
		res[key] = int64(val.Score)
	}
	return res, nil
}
func (p *RedisClusterPool) LRange(key string, start, stop int64) (interface{}, error) {
	res, err := p.client.LRange(key, start, stop).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) LSet(key string, index int, value interface{}) error {
	err := p.client.LSet(key, int64(index), value).Err()
	return convertError(err)
}
func (p *RedisClusterPool) LLen(key string) (int64, error) {
	res, err := p.client.LLen(key).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) LRem(key string, count int, value interface{}) (int, error) {
	val, _ := value.(string)
	res, err := p.client.LRem(key, int64(count), val).Result()
	if err != nil {
		return int(res), convertError(err)
	}
	return int(res), nil
}
func (p *RedisClusterPool) TTl(key string) (int64, error) {
	duration, err := p.client.TTL(key).Result()
	if err != nil {
		return int64(duration.Seconds()), convertError(err)
	}
	return int64(duration.Seconds()), nil
}
func (p *RedisClusterPool) LPop(key string) (interface{}, error) {
	res, err := p.client.LPop(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) RPop(key string) (interface{}, error) {
	res, err := p.client.RPop(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) BLPop(key string, timeout int) (interface{}, error) {
	res, err := p.client.BLPop(time.Duration(timeout)*time.Second, key).Result()
	if err != nil {
		// 兼容redis 2.x
		if err == redis.Nil {
			return nil, ErrNil
		}
		return nil, err
	}
	return res[1], nil
}
func (p *RedisClusterPool) BRPop(key string, timeout int) (interface{}, error) {
	res, err := p.client.BRPop(time.Duration(timeout)*time.Second, key).Result()
	if err != nil {
		// 兼容redis 2.x
		if err == redis.Nil {
			return nil, ErrNil
		}
		return nil, convertError(err)
	}
	return res[1], nil
}
func (p *RedisClusterPool) LPush(key string, value ...interface{}) error {
	args := make([]interface{}, 0, len(value)+1)
	args = append(args, key)
	args = append(args, value...)
	vals := make([]string, 0, len(value))
	for _, v := range value {
		val, err := String(v, nil)
		if err != nil && err != ErrNil {
			return err
		}
		vals = append(vals, val)
	}
	_, err := p.client.LPush(key, vals).Result() // todo ...
	if err != nil {
		return convertError(err)
	}
	return nil
}
func (p *RedisClusterPool) RPush(key string, value ...interface{}) error {
	args := make([]interface{}, 0, len(value)+1)
	args = append(args, key)
	args = append(args, value...)
	vals := make([]string, 0, len(value))
	for _, v := range value {
		val, err := String(v, nil)
		if err != nil && err != ErrNil {
			if err == ErrNil {
				continue
			}
			return err
		}
		if val == "" {
			continue
		}
		vals = append(vals, val)
	}
	_, err := p.client.RPush(key, vals).Result() // todo ...
	if err != nil {
		return convertError(err)
	}
	return nil
}

// 为确保srcKey跟destKey映射到同一个slot，srcKey和destKey需要加上hash tag，如:{test}
func (p *RedisClusterPool) BRPopLPush(srcKey string, destKey string, timeout int) (interface{}, error) {
	res, err := p.client.BRPopLPush(srcKey, destKey, time.Duration(timeout)*time.Second).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}

// 为确保srcKey跟destKey映射到同一个slot，srcKey和destKey需要加上hash tag，如:{test}
func (p *RedisClusterPool) RPopLPush(srcKey string, destKey string) (interface{}, error) {
	res, err := p.client.RPopLPush(srcKey, destKey).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) SAdd(key string, members ...interface{}) (int64, error) {
	args := make([]interface{}, 0, len(members)+1)
	args = append(args, key)
	args = append(args, members...)
	ms := make([]string, 0, len(members))
	for _, member := range members {
		m, err := String(member, nil)
		if err != nil && err != ErrNil {
			return 0, err
		}
		ms = append(ms, m)
	}
	res, err := p.client.SAdd(key, ms).Result() // todo ...
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) SPop(key string) ([]byte, error) {
	res, err := p.client.SPop(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return []byte(res), nil
}
func (p *RedisClusterPool) SIsMember(key string, member interface{}) (bool, error) {
	m, _ := member.(string)
	res, err := p.client.SIsMember(key, m).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) SRem(key string, members ...interface{}) (int64, error) {
	args := make([]interface{}, 0, len(members)+1)
	args = append(args, key)
	args = append(args, members...)
	ms := make([]string, 0, len(members))
	for _, member := range members {
		m, err := String(member, nil)
		if err != nil && err != ErrNil {
			return 0, err
		}
		ms = append(ms, m)
	}
	res, err := p.client.SRem(key, ms).Result() // todo ...
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) SMembers(key string) ([]string, error) {
	res, err := p.client.SMembers(key).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) ScriptLoad(luaScript string) (interface{}, error) {
	res, err := p.client.ScriptLoad(luaScript).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) EvalSha(sha1 string, numberKeys int, keysArgs ...interface{}) (interface{}, error) {
	vals := make([]interface{}, 0, len(keysArgs)+2)
	vals = append(vals, sha1, numberKeys)
	vals = append(vals, keysArgs...)
	keys := make([]string, 0, numberKeys)
	args := make([]string, 0, len(keysArgs)-numberKeys)
	for i, value := range keysArgs {
		val, err := String(value, nil)
		if err != nil && err != ErrNil {
			return nil, err
		}
		if i < numberKeys {
			keys = append(keys, val)
		} else {
			args = append(args, val)
		}
	}
	res, err := p.client.EvalSha(sha1, keys, args).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) Eval(luaScript string, numberKeys int, keysArgs ...interface{}) (interface{}, error) {
	vals := make([]interface{}, 0, len(keysArgs)+2)
	vals = append(vals, luaScript, numberKeys)
	vals = append(vals, keysArgs...)
	keys := make([]string, 0, numberKeys)
	args := make([]string, 0, len(keysArgs)-numberKeys)
	for i, value := range keysArgs {
		val, err := String(value, nil)
		if err != nil && err != ErrNil {
			return nil, err
		}
		if i < numberKeys {
			keys = append(keys, val)
		} else {
			args = append(args, val)
		}
	}
	res, err := p.client.Eval(luaScript, keys, args).Result()
	if err != nil {
		return nil, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) GetBit(key string, offset int64) (int64, error) {
	res, err := p.client.GetBit(key, offset).Result()
	if err != nil {
		return res, convertError(err)
	}
	return res, nil
}
func (p *RedisClusterPool) SetBit(key string, offset uint32, value int) (int, error) {
	res, err := p.client.SetBit(key, int64(offset), value).Result()
	return int(res), convertError(err)
}
func (p *RedisClusterPool) GetClient() *redis.ClusterClient {
	return pools
}