agent_ad/app/utils/cache/redis.go

package cache

import (
	"encoding/json"
	"errors"
	"log"
	"strings"
	"time"

	redigo "github.com/gomodule/redigo/redis"
)

// configuration
type Config struct {
	Server   string
	Password string
	MaxIdle  int // Maximum number of idle connections in the pool.

	// Maximum number of connections allocated by the pool at a given time.
	// When zero, there is no limit on the number of connections in the pool.
	MaxActive int

	// Close connections after remaining idle for this duration. If the value
	// is zero, then idle connections are not closed. Applications should set
	// the timeout to a value less than the server's timeout.
	IdleTimeout time.Duration

	// If Wait is true and the pool is at the MaxActive limit, then Get() waits
	// for a connection to be returned to the pool before returning.
	Wait           bool
	KeyPrefix      string // prefix to all keys; example is "dev environment name"
	KeyDelimiter   string // delimiter to be used while appending keys; example is ":"
	KeyPlaceholder string // placeholder to be parsed using given arguments to obtain a final key; example is "?"
}

var pool *redigo.Pool
var conf *Config

func NewRedis(addr string) {
	if addr == "" {
		panic("\nredis connect string cannot be empty\n")
	}
	pool = &redigo.Pool{
		MaxIdle:     redisMaxIdleConn,
		IdleTimeout: redisIdleTTL,
		MaxActive:   redisMaxActive,
		// MaxConnLifetime: redisDialTTL,
		Wait: true,
		Dial: func() (redigo.Conn, error) {
			c, err := redigo.Dial("tcp", addr,
				redigo.DialConnectTimeout(redisDialTTL),
				redigo.DialReadTimeout(redisReadTTL),
				redigo.DialWriteTimeout(redisWriteTTL),
			)
			if err != nil {
				log.Println("Redis Dial failed: ", err)
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redigo.Conn, t time.Time) error {
			_, err := c.Do("PING")
			if err != nil {
				log.Println("Unable to ping to redis server:", err)
			}
			return err
		},
	}
	conn := pool.Get()
	defer conn.Close()
	if conn.Err() != nil {
		println("\nredis connect " + addr + " error: " + conn.Err().Error())
	} else {
		println("\nredis connect " + addr + " success!\n")
	}
}

func Do(cmd string, args ...interface{}) (reply interface{}, err error) {
	conn := pool.Get()
	defer conn.Close()
	return conn.Do(cmd, args...)
}

func GetPool() *redigo.Pool {
	return pool
}

func ParseKey(key string, vars []string) (string, error) {
	arr := strings.Split(key, conf.KeyPlaceholder)
	actualKey := ""
	if len(arr) != len(vars)+1 {
		return "", errors.New("redis/connection.go: Insufficient arguments to parse key")
	} else {
		for index, val := range arr {
			if index == 0 {
				actualKey = arr[index]
			} else {
				actualKey += vars[index-1] + val
			}
		}
	}
	return getPrefixedKey(actualKey), nil
}

func getPrefixedKey(key string) string {
	return conf.KeyPrefix + conf.KeyDelimiter + key
}
func StripEnvKey(key string) string {
	return strings.TrimLeft(key, conf.KeyPrefix+conf.KeyDelimiter)
}
func SplitKey(key string) []string {
	return strings.Split(key, conf.KeyDelimiter)
}
func Expire(key string, ttl int) (interface{}, error) {
	return Do("EXPIRE", key, ttl)
}
func Persist(key string) (interface{}, error) {
	return Do("PERSIST", key)
}

func Del(key string) (interface{}, error) {
	return Do("DEL", key)
}
func Set(key string, data interface{}) (interface{}, error) {
	// set
	return Do("SET", key, data)
}
func SetNX(key string, data interface{}) (interface{}, error) {
	return Do("SETNX", key, data)
}
func SetEx(key string, data interface{}, ttl int) (interface{}, error) {
	return Do("SETEX", key, ttl, data)
}

func SetJson(key string, data interface{}, ttl int) bool {
	c, err := json.Marshal(data)
	if err != nil {
		return false
	}
	if ttl < 1 {
		_, err = Set(key, c)
	} else {
		_, err = SetEx(key, c, ttl)
	}
	if err != nil {
		return false
	}
	return true
}

func GetJson(key string, dst interface{}) error {
	b, err := GetBytes(key)
	if err != nil {
		return err
	}
	if err = json.Unmarshal(b, dst); err != nil {
		return err
	}
	return nil
}

func Get(key string) (interface{}, error) {
	// get
	return Do("GET", key)
}
func GetTTL(key string) (time.Duration, error) {
	ttl, err := redigo.Int64(Do("TTL", key))
	return time.Duration(ttl) * time.Second, err
}
func GetBytes(key string) ([]byte, error) {
	return redigo.Bytes(Do("GET", key))
}
func GetString(key string) (string, error) {
	return redigo.String(Do("GET", key))
}
func GetStringMap(key string) (map[string]string, error) {
	return redigo.StringMap(Do("GET", key))
}
func GetInt(key string) (int, error) {
	return redigo.Int(Do("GET", key))
}
func GetInt64(key string) (int64, error) {
	return redigo.Int64(Do("GET", key))
}
func GetStringLength(key string) (int, error) {
	return redigo.Int(Do("STRLEN", key))
}
func ZAdd(key string, score float64, data interface{}) (interface{}, error) {
	return Do("ZADD", key, score, data)
}
func ZAddNX(key string, score float64, data interface{}) (interface{}, error) {
	return Do("ZADD", key, "NX", score, data)
}
func ZRem(key string, data interface{}) (interface{}, error) {
	return Do("ZREM", key, data)
}
func ZRange(key string, start int, end int, withScores bool) ([]interface{}, error) {
	if withScores {
		return redigo.Values(Do("ZRANGE", key, start, end, "WITHSCORES"))
	}
	return redigo.Values(Do("ZRANGE", key, start, end))
}
func ZRemRangeByScore(key string, start int64, end int64) ([]interface{}, error) {
	return redigo.Values(Do("ZREMRANGEBYSCORE", key, start, end))
}
func ZCard(setName string) (int64, error) {
	return redigo.Int64(Do("ZCARD", setName))
}
func ZScan(setName string) (int64, error) {
	return redigo.Int64(Do("ZCARD", setName))
}
func SAdd(setName string, data interface{}) (interface{}, error) {
	return Do("SADD", setName, data)
}
func SCard(setName string) (int64, error) {
	return redigo.Int64(Do("SCARD", setName))
}
func SIsMember(setName string, data interface{}) (bool, error) {
	return redigo.Bool(Do("SISMEMBER", setName, data))
}
func SMembers(setName string) ([]string, error) {
	return redigo.Strings(Do("SMEMBERS", setName))
}
func SRem(setName string, data interface{}) (interface{}, error) {
	return Do("SREM", setName, data)
}
func HSet(key string, HKey string, data interface{}) (interface{}, error) {
	return Do("HSET", key, HKey, data)
}

func HGet(key string, HKey string) (interface{}, error) {
	return Do("HGET", key, HKey)
}

func HMGet(key string, hashKeys ...string) ([]interface{}, error) {
	ret, err := Do("HMGET", key, hashKeys)
	if err != nil {
		return nil, err
	}
	reta, ok := ret.([]interface{})
	if !ok {
		return nil, errors.New("result not an array")
	}
	return reta, nil
}

func HMSet(key string, hashKeys []string, vals []interface{}) (interface{}, error) {
	if len(hashKeys) == 0 || len(hashKeys) != len(vals) {
		var ret interface{}
		return ret, errors.New("bad length")
	}
	input := []interface{}{key}
	for i, v := range hashKeys {
		input = append(input, v, vals[i])
	}
	return Do("HMSET", input...)
}

func HGetString(key string, HKey string) (string, error) {
	return redigo.String(Do("HGET", key, HKey))
}
func HGetFloat(key string, HKey string) (float64, error) {
	f, err := redigo.Float64(Do("HGET", key, HKey))
	return f, err
}
func HGetInt(key string, HKey string) (int, error) {
	return redigo.Int(Do("HGET", key, HKey))
}
func HGetInt64(key string, HKey string) (int64, error) {
	return redigo.Int64(Do("HGET", key, HKey))
}
func HGetBool(key string, HKey string) (bool, error) {
	return redigo.Bool(Do("HGET", key, HKey))
}
func HDel(key string, HKey string) (interface{}, error) {
	return Do("HDEL", key, HKey)
}

func HGetAll(key string) (map[string]interface{}, error) {
	vals, err := redigo.Values(Do("HGETALL", key))
	if err != nil {
		return nil, err
	}
	num := len(vals) / 2
	result := make(map[string]interface{}, num)
	for i := 0; i < num; i++ {
		key, _ := redigo.String(vals[2*i], nil)
		result[key] = vals[2*i+1]
	}
	return result, nil
}

func FlushAll() bool {
	res, _ := redigo.String(Do("FLUSHALL"))
	if res == "" {
		return false
	}
	return true
}

// NOTE: Use this in production environment with extreme care.
// Read more here:https://redigo.io/commands/keys
func Keys(pattern string) ([]string, error) {
	return redigo.Strings(Do("KEYS", pattern))
}

func HKeys(key string) ([]string, error) {
	return redigo.Strings(Do("HKEYS", key))
}

func Exists(key string) bool {
	count, err := redigo.Int(Do("EXISTS", key))
	if count == 0 || err != nil {
		return false
	}
	return true
}

func Incr(key string) (int64, error) {
	return redigo.Int64(Do("INCR", key))
}

func Decr(key string) (int64, error) {
	return redigo.Int64(Do("DECR", key))
}

func IncrBy(key string, incBy int64) (int64, error) {
	return redigo.Int64(Do("INCRBY", key, incBy))
}

func DecrBy(key string, decrBy int64) (int64, error) {
	return redigo.Int64(Do("DECRBY", key))
}

func IncrByFloat(key string, incBy float64) (float64, error) {
	return redigo.Float64(Do("INCRBYFLOAT", key, incBy))
}

func DecrByFloat(key string, decrBy float64) (float64, error) {
	return redigo.Float64(Do("DECRBYFLOAT", key, decrBy))
}

// use for message queue
func LPush(key string, data interface{}) (interface{}, error) {
	// set
	return Do("LPUSH", key, data)
}

func LPop(key string) (interface{}, error) {
	return Do("LPOP", key)
}

func LPopString(key string) (string, error) {
	return redigo.String(Do("LPOP", key))
}
func LPopFloat(key string) (float64, error) {
	f, err := redigo.Float64(Do("LPOP", key))
	return f, err
}
func LPopInt(key string) (int, error) {
	return redigo.Int(Do("LPOP", key))
}
func LPopInt64(key string) (int64, error) {
	return redigo.Int64(Do("LPOP", key))
}

func RPush(key string, data interface{}) (interface{}, error) {
	// set
	return Do("RPUSH", key, data)
}

func RPop(key string) (interface{}, error) {
	return Do("RPOP", key)
}

func RPopString(key string) (string, error) {
	return redigo.String(Do("RPOP", key))
}
func RPopFloat(key string) (float64, error) {
	f, err := redigo.Float64(Do("RPOP", key))
	return f, err
}
func RPopInt(key string) (int, error) {
	return redigo.Int(Do("RPOP", key))
}
func RPopInt64(key string) (int64, error) {
	return redigo.Int64(Do("RPOP", key))
}

func Scan(cursor int64, pattern string, count int64) (int64, []string, error) {
	var items []string
	var newCursor int64

	values, err := redigo.Values(Do("SCAN", cursor, "MATCH", pattern, "COUNT", count))
	if err != nil {
		return 0, nil, err
	}
	values, err = redigo.Scan(values, &newCursor, &items)
	if err != nil {
		return 0, nil, err
	}
	return newCursor, items, nil
}

func LPushMax(key string, data ...interface{}) (interface{}, error) {
	// set
	return Do("LPUSH", key, data)
}

func SelectDb(db int) (interface{}, error) {
	return Do("SELECT", db)
}