github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/redis/cache.go (about)

     1  package redis
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/gomodule/redigo/redis"
     8  	"github.com/sirupsen/logrus"
     9  )
    10  
    11  // Cache is Redis cache adapter.
    12  type RedisCache struct {
    13  	Pool *redis.Pool // redis connection pool
    14  }
    15  
    16  // NewRedisCache create new redis cache with default collection name.
    17  func NewRedisCache() *RedisCache {
    18  	return &RedisCache{}
    19  }
    20  
    21  func RedisDo(conn redis.Conn) func(cmd string, args ...interface{}) (interface{}, error) {
    22  	return func(cmd string, args ...interface{}) (interface{}, error) {
    23  		res, err := conn.Do(cmd, args...)
    24  		if err != nil {
    25  			logrus.WithField("redis_cmd", cmd).Warningf("%v", err.Error())
    26  		}
    27  		return res, err
    28  	}
    29  }
    30  
    31  func formatMs(dur time.Duration) int64 {
    32  	if dur > 0 && dur < time.Millisecond {
    33  		return 0
    34  	}
    35  	return int64(dur / time.Millisecond)
    36  }
    37  
    38  func formatSec(dur time.Duration) int64 {
    39  	if dur > 0 && dur < time.Second {
    40  		return 0
    41  	}
    42  	return int64(dur / time.Second)
    43  }
    44  
    45  func usePrecise(dur time.Duration) bool {
    46  	return dur < time.Second || dur%time.Second != 0
    47  }
    48  
    49  // Start start redis cache adapter.
    50  // config is like {"key":"collection key","conn":"connection info","dbNum":"0"}
    51  // the cache item in redis are stored forever,
    52  // so no gc operation.
    53  func (cache *RedisCache) Start(config *Redis) error {
    54  	if err := cache.connectInit(config); err != nil {
    55  		return err
    56  	}
    57  
    58  	c := cache.Pool.Get()
    59  	defer c.Close()
    60  
    61  	return c.Err()
    62  }
    63  
    64  // connect to redis.
    65  func (cache *RedisCache) connectInit(config *Redis) error {
    66  	dialFunc := func() (c redis.Conn, err error) {
    67  		c, err = redis.DialTimeout(
    68  			config.Protocol,
    69  			fmt.Sprintf("%s:%d", config.Host, config.Port),
    70  			config.ConnectTimeout,
    71  			config.ReadTimeout,
    72  			config.WriteTimeout,
    73  		)
    74  		if err != nil {
    75  			return
    76  		}
    77  
    78  		if config.Password != "" {
    79  			if _, err := RedisDo(c)("AUTH", config.Password.String()); err != nil {
    80  				c.Close()
    81  				return c, err
    82  			}
    83  		}
    84  
    85  		_, selectErr := RedisDo(c)("SELECT", config.DB)
    86  		if selectErr != nil {
    87  			c.Close()
    88  			return nil, selectErr
    89  		}
    90  
    91  		return
    92  	}
    93  
    94  	// initialize a new pool
    95  	cache.Pool = &redis.Pool{
    96  		Dial:        dialFunc,
    97  		MaxIdle:     config.MaxIdle,
    98  		MaxActive:   config.MaxActive,
    99  		IdleTimeout: config.IdleTimeout,
   100  		Wait:        true,
   101  	}
   102  
   103  	return nil
   104  }
   105  
   106  // Redis `SET key value [expiration] NX` command.
   107  // Zero expiration means the key has no expiration time.
   108  func (cache *RedisCache) SetNX(key string, value interface{}, expiration time.Duration) (bool, error) {
   109  	conn := cache.Pool.Get()
   110  	defer conn.Close()
   111  	var err error
   112  	if expiration == 0 {
   113  		// Use old `SETNX` to support old Redis versions.
   114  		_, err = redis.String(RedisDo(conn)("setnx", key, value))
   115  	} else {
   116  		if usePrecise(expiration) {
   117  			_, err = redis.String(RedisDo(conn)("set", key, value, "nx", "px", formatMs(expiration)))
   118  		} else {
   119  			_, err = redis.String(RedisDo(conn)("set", key, value, "nx", "ex", formatSec(expiration)))
   120  		}
   121  	}
   122  
   123  	if err != nil {
   124  		if err != redis.ErrNil {
   125  			return false, err
   126  		} else {
   127  			return false, nil
   128  		}
   129  	} else {
   130  		return true, err
   131  	}
   132  }
   133  
   134  func (cache *RedisCache) Get(key string) ([]byte, error) {
   135  	conn := cache.Pool.Get()
   136  	defer conn.Close()
   137  
   138  	return redis.Bytes(RedisDo(conn)("GET", key))
   139  }
   140  
   141  func (cache *RedisCache) Incr(key string) (int64, error) {
   142  	conn := cache.Pool.Get()
   143  	defer conn.Close()
   144  	return redis.Int64(RedisDo(conn)("INCR", key))
   145  }
   146  
   147  func (cache *RedisCache) MGet(key []interface{}) (interface{}, error) {
   148  	conn := cache.Pool.Get()
   149  	defer conn.Close()
   150  	return RedisDo(conn)("MGET", key...)
   151  }
   152  
   153  func (cache *RedisCache) MGetValue(keys []interface{}) ([]interface{}, error) {
   154  	conn := cache.Pool.Get()
   155  	defer conn.Close()
   156  	return redis.Values(RedisDo(conn)("MGET", keys...))
   157  }
   158  
   159  func (cache *RedisCache) HSet(key, field string, value interface{}) (interface{}, error) {
   160  	conn := cache.Pool.Get()
   161  	defer conn.Close()
   162  	return RedisDo(conn)("HSET", key, field, value)
   163  }
   164  
   165  func (cache *RedisCache) HSetWithExpire(key string, timeout time.Duration) (interface{}, error) {
   166  	conn := cache.Pool.Get()
   167  	defer conn.Close()
   168  
   169  	reply, err := RedisDo(conn)("HSET", key, nil, nil)
   170  	if err != nil {
   171  		return reply, err
   172  	}
   173  
   174  	return RedisDo(conn)("EXPIRE", key, formatSec(timeout))
   175  }
   176  
   177  func (cache *RedisCache) HMset(value []interface{}) (interface{}, error) {
   178  	conn := cache.Pool.Get()
   179  	defer conn.Close()
   180  	return RedisDo(conn)("HMSET", value...)
   181  }
   182  
   183  func (cache *RedisCache) HGet(key, field string) (interface{}, error) {
   184  	conn := cache.Pool.Get()
   185  	defer conn.Close()
   186  	return RedisDo(conn)("HGET", key, field)
   187  }
   188  
   189  func (cache *RedisCache) HIncrBy(key, field string, value interface{}) (int64, error) {
   190  	conn := cache.Pool.Get()
   191  	defer conn.Close()
   192  	return redis.Int64(RedisDo(conn)("HINCRBY", key, field, value))
   193  }
   194  
   195  func (cache *RedisCache) HMGet(key string, fields []string) (interface{}, error) {
   196  	conn := cache.Pool.Get()
   197  	defer conn.Close()
   198  
   199  	var args []interface{}
   200  	args = append(args, key)
   201  	for _, field := range fields {
   202  		args = append(args, field)
   203  	}
   204  
   205  	return RedisDo(conn)("HMGET", args...)
   206  }
   207  
   208  func (cache *RedisCache) GetString(key string) (string, error) {
   209  	conn := cache.Pool.Get()
   210  	defer conn.Close()
   211  	return redis.String(RedisDo(conn)("GET", key))
   212  }
   213  
   214  func (cache *RedisCache) GetStringMap(key string) (map[string]string, error) {
   215  	conn := cache.Pool.Get()
   216  	defer conn.Close()
   217  	return redis.StringMap(RedisDo(conn)("HGETALL", key))
   218  }
   219  
   220  func (cache *RedisCache) HGetAll(key string) (interface{}, error) {
   221  	conn := cache.Pool.Get()
   222  	defer conn.Close()
   223  	return RedisDo(conn)("HGETALL", key)
   224  }
   225  
   226  func (cache *RedisCache) GetInts(key string) ([]int, error) {
   227  	conn := cache.Pool.Get()
   228  	defer conn.Close()
   229  	return redis.Ints(RedisDo(conn)("GET", key))
   230  }
   231  
   232  func (cache *RedisCache) Expire(key string, timeout time.Duration) error {
   233  	conn := cache.Pool.Get()
   234  	defer conn.Close()
   235  	_, err := RedisDo(conn)("EXPIRE", key, formatSec(timeout))
   236  	return err
   237  }
   238  
   239  func (cache *RedisCache) Set(key string, bytes interface{}, timeout time.Duration) error {
   240  	conn := cache.Pool.Get()
   241  	defer conn.Close()
   242  	if timeout == -1 {
   243  		_, err := RedisDo(conn)("SET", key, bytes)
   244  		return err
   245  	} else {
   246  		_, err := RedisDo(conn)("SET", key, bytes, "EX", formatSec(timeout))
   247  		return err
   248  	}
   249  }
   250  
   251  func (cache *RedisCache) Del(key string) error {
   252  	conn := cache.Pool.Get()
   253  	defer conn.Close()
   254  	_, err := RedisDo(conn)("DEL", key)
   255  
   256  	return err
   257  }
   258  
   259  func (cache *RedisCache) Exists(key string) (bool, error) {
   260  	conn := cache.Pool.Get()
   261  	defer conn.Close()
   262  	exists, err := redis.Bool(RedisDo(conn)("EXISTS", key))
   263  	if err != nil {
   264  		return false, err
   265  	}
   266  
   267  	if exists {
   268  		return true, nil
   269  	} else {
   270  		return false, nil
   271  	}
   272  }
   273  
   274  func (cache *RedisCache) ZRange(key string, start, end int, withscores bool) ([]string, error) {
   275  	conn := cache.Pool.Get()
   276  	defer conn.Close()
   277  	var res []string
   278  	var err error
   279  	if withscores {
   280  		res, err = redis.Strings(RedisDo(conn)("ZRANGE", key, start, end, "withscores"))
   281  	} else {
   282  		res, err = redis.Strings(RedisDo(conn)("ZRANGE", key, start, end))
   283  	}
   284  	return res, err
   285  }
   286  
   287  func (cache *RedisCache) ZRangeInts(key string, start, end int, withscores bool) ([]int, error) {
   288  	conn := cache.Pool.Get()
   289  	defer conn.Close()
   290  	var res []int
   291  	var err error
   292  	if withscores {
   293  		res, err = redis.Ints(RedisDo(conn)("ZRANGE", key, start, end, "withscores"))
   294  	} else {
   295  		res, err = redis.Ints(RedisDo(conn)("ZRANGE", key, start, end))
   296  	}
   297  	return res, err
   298  }
   299  
   300  func (cache *RedisCache) ZRevrange(key string, start, end int, withscores bool) ([]int, error) {
   301  	conn := cache.Pool.Get()
   302  	defer conn.Close()
   303  	var res []int
   304  	var err error
   305  	if withscores {
   306  		res, err = redis.Ints(RedisDo(conn)("ZREVRANGE", key, start, end, "withscores"))
   307  	} else {
   308  		res, err = redis.Ints(RedisDo(conn)("ZREVRANGE", key, start, end))
   309  	}
   310  
   311  	return res, err
   312  }
   313  
   314  func (cache *RedisCache) ZRevrangeByScore(key string, max_num, min_num int, withscores bool, offset, count int) ([]int, error) {
   315  	conn := cache.Pool.Get()
   316  	defer conn.Close()
   317  	var res []int
   318  	var err error
   319  	if !withscores {
   320  		res, err = redis.Ints(RedisDo(conn)("ZREVRANGEBYSCORE", key, max_num, min_num, "limit", offset, count))
   321  	} else {
   322  		res, err = redis.Ints(RedisDo(conn)("ZREVRANGEBYSCORE", key, max_num, min_num, "withscores", "limit", offset, count))
   323  	}
   324  	return res, err
   325  }
   326  func (cache *RedisCache) ZRangeByScore(key string, min_num, max_num int64, withscores bool, offset, count int) ([]int, error) {
   327  	conn := cache.Pool.Get()
   328  	defer conn.Close()
   329  	var res []int
   330  	var err error
   331  	if withscores {
   332  		res, err = redis.Ints(RedisDo(conn)("ZREVRANGEBYSCORE", key, max_num, min_num, "limit", offset, count))
   333  	} else {
   334  		res, err = redis.Ints(RedisDo(conn)("ZREVRANGEBYSCORE", key, max_num, min_num, "withscores", "limit", offset, count))
   335  	}
   336  	return res, err
   337  }
   338  
   339  func (cache *RedisCache) ZScore(key, member string) (interface{}, error) {
   340  	conn := cache.Pool.Get()
   341  	defer conn.Close()
   342  
   343  	res, err := RedisDo(conn)("ZSCORE", key, member)
   344  	return res, err
   345  }
   346  
   347  func (cache *RedisCache) ZAdd(key string, value, member interface{}) (interface{}, error) {
   348  	conn := cache.Pool.Get()
   349  	defer conn.Close()
   350  
   351  	res, err := RedisDo(conn)("ZADD", key, value, member)
   352  	return res, err
   353  }
   354  
   355  func (cache *RedisCache) SAdd(key string, items string) (int, error) {
   356  	//var err error
   357  	conn := cache.Pool.Get()
   358  	defer conn.Close()
   359  	res, err := redis.Int(RedisDo(conn)("SADD", key, items))
   360  	return res, err
   361  }
   362  
   363  func (cache *RedisCache) SIsMember(key string, items string) (int, error) {
   364  	//var err error
   365  	conn := cache.Pool.Get()
   366  	defer conn.Close()
   367  	res, err := redis.Int(RedisDo(conn)("SISMEMBER", key, items))
   368  	return res, err
   369  }
   370  
   371  func (cache *RedisCache) RPush(key string, value interface{}) error {
   372  	conn := cache.Pool.Get()
   373  	defer conn.Close()
   374  	_, err := RedisDo(conn)("RPUSH", key, value)
   375  	return err
   376  }
   377  
   378  func (cache *RedisCache) RPushBatch(keys []interface{}) error {
   379  	conn := cache.Pool.Get()
   380  	defer conn.Close()
   381  	_, err := RedisDo(conn)("RPUSH", keys...)
   382  	return err
   383  }
   384  
   385  func (cache *RedisCache) LRange(key string, start, end int) ([]interface{}, error) {
   386  	conn := cache.Pool.Get()
   387  	defer conn.Close()
   388  	result, err := redis.Values(RedisDo(conn)("LRANGE", key, start, end))
   389  	return result, err
   390  }
   391  
   392  func (cache *RedisCache) LRem(key string, value interface{}) error {
   393  	conn := cache.Pool.Get()
   394  	defer conn.Close()
   395  	_, err := RedisDo(conn)("LREM", key, 0, value)
   396  	return err
   397  }
   398  
   399  func (cache *RedisCache) TTL(key string) (int, error) {
   400  	conn := cache.Pool.Get()
   401  	defer conn.Close()
   402  	return redis.Int(RedisDo(conn)("ttl", key))
   403  }